Compute Diffusion

diffusion

FUNCTION	DESCRIPTION
local_diffusion	Estimate position-dependent diffusion from trajectory displacements.
diffusion_tensor	Compute diffusion tensor \(D = J G^{-1} J^T\) from model Jacobian.
log_det_diffusion_tensor	Compute log determinant of diffusion tensor \(D = J G^{-1} J^T\).
inverse_diffusion_tensor	Compute inverse metric tensor from diffusion tensor \(D = J G^{-1} J^T\).
decompose_diffusion_tensor	Compute inverse and log-determinant of a diffusion tensor via single

Functions

local_diffusion(X: torch.Tensor, dt: float = 1.0, bandwidth: float | None = None) -> torch.Tensor

Estimate position-dependent diffusion from trajectory displacements.

Uses the short-time displacement covariance to estimate the diagonal diffusion tensor at each frame:

\(D_j(x_i) = \frac{\sum_k K(x_i - x_k)\, (\Delta x_{k,j})^2} {2\, \Delta t \sum_k K(x_i - x_k)}\)

where \(K\) is a Gaussian kernel and \(\Delta x_k = x_{k+1} - x_k\). This exploits the fact that the displacement covariance at short times is \(\operatorname{Cov}(\Delta x \mid x) = 2 k_BT D_0(x)\, \Delta t + O(\Delta t^2)\), so the diffusion is determined purely by the fluctuations, independent of the drift (forces).

The result can be passed as metric_inv to compute_diffusion_tensor for the Fixman-corrected invariant free energy.

PARAMETER	DESCRIPTION
X	Trajectory positions with shape (N, d). Frames must be consecutive (constant timestep dt between frames). TYPE: Tensor
dt	Timestep between consecutive frames. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
bandwidth	Gaussian kernel bandwidth for smoothing. If None, uses Silverman's rule: \(h = \left(\frac{4}{d+2}\right)^{1/(d+4)} N^{-1/(d+4)} \bar{\sigma}\) where \(\bar{\sigma}\) is the mean standard deviation across dimensions. TYPE: float | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Per-sample diagonal diffusion tensor with shape (N, d). The last frame is assigned the same value as the second-to-last frame since no forward displacement is available.

Source code in spectre/compute/diffusion.py

def local_diffusion(
    X: torch.Tensor,
    dt: float = 1.0,
    bandwidth: float | None = None,
) -> torch.Tensor:
    """
    Estimate position-dependent diffusion from trajectory displacements.

    Uses the short-time displacement covariance to estimate the diagonal
    diffusion tensor at each frame:

    $D_j(x_i) = \\frac{\\sum_k K(x_i - x_k)\\, (\\Delta x_{k,j})^2}
    {2\\, \\Delta t \\sum_k K(x_i - x_k)}$

    where $K$ is a Gaussian kernel and $\\Delta x_k = x_{k+1} - x_k$. This
    exploits the fact that the displacement covariance at short times is
    $\\operatorname{Cov}(\\Delta x \\mid x) = 2 k_BT D_0(x)\\, \\Delta t +
    O(\\Delta t^2)$, so the diffusion is determined purely by the fluctuations,
    independent of the drift (forces).

    The result can be passed as ``metric_inv`` to `compute_diffusion_tensor`
    for the Fixman-corrected invariant free energy.

    Parameters
    ----------
    X : torch.Tensor
        Trajectory positions with shape (N, d). Frames must be consecutive
        (constant timestep `dt` between frames).
    dt : float, optional, by default 1.0
        Timestep between consecutive frames.
    bandwidth : float | None, optional, by default None
        Gaussian kernel bandwidth for smoothing. If ``None``, uses Silverman's
        rule: $h = \\left(\\frac{4}{d+2}\\right)^{1/(d+4)} N^{-1/(d+4)}
        \\bar{\\sigma}$ where $\\bar{\\sigma}$ is the mean standard deviation
        across dimensions.

    Returns
    -------
    torch.Tensor
        Per-sample diagonal diffusion tensor with shape (N, d). The last frame
        is assigned the same value as the second-to-last frame since no forward
        displacement is available.
    """
    if X.ndim != 2:
        raise ValueError(f"Expected `X` to have 2 dimensions, got {X.ndim}.")

    N, d = X.shape

    if N < 2:
        raise ValueError("Need at least 2 frames to estimate diffusion.")

    # Displacements: shape (N-1, d).
    dx = X[1:] - X[:-1]
    dx_sq = dx**2  # (N-1, d)

    # Positions corresponding to each displacement.
    X_dx = X[:-1]  # (N-1, d)

    if bandwidth is None:
        sigma = X_dx.std(dim=0).mean()
        bandwidth = float(
            (4.0 / (d + 2)) ** (1.0 / (d + 4)) * (N - 1) ** (-1.0 / (d + 4)) * sigma
        )

    # Kernel-smoothed estimate at each frame.
    # For each frame i in X, compute weighted average of dx_sq over
    # nearby displacement frames.
    D_local = torch.zeros(N, d, device=X.device, dtype=X.dtype)

    inv_bw_sq = 1.0 / (bandwidth**2)

    # Pairwise squared distances: X (N, d) vs X_dx (N-1, d).
    dists_sq = torch.cdist(X, X_dx).pow(2)  # (N, N-1)
    weights = torch.exp(-0.5 * inv_bw_sq * dists_sq)  # (N, N-1)

    # Weighted average: D_j(x_i) = sum_k w_ik * dx_sq_kj / (2*dt * sum_k w_ik)
    w_sum = weights.sum(dim=1, keepdim=True).clamp(min=1e-30)  # (N, 1)
    D_local = torch.matmul(weights, dx_sq) / (2.0 * dt * w_sum)  # (N, d)

    return D_local

diffusion_tensor(X: torch.Tensor, model: torch.nn.Module, mode: Literal['forward', 'reverse', 'auto'] = 'auto', metric_inv: torch.Tensor | None = None) -> torch.Tensor

Compute diffusion tensor \(D = J G^{-1} J^T\) from model Jacobian.

When metric_inv is None, assumes a Euclidean metric (\(G = I\)) and computes \(D = J J^T\). When provided, computes the Fixman-corrected diffusion tensor \(D = J G^{-1} J^T\), where \(G^{-1}\) is the inverse metric of the input coordinate space (e.g., inverse mass matrix).

PARAMETER	DESCRIPTION
X	Input data tensor with shape (N, d). TYPE: Tensor
model	Neural network model mapping R^d -> R^k. TYPE: Module
mode	Autodiff mode for Jacobian computation. Forward mode scales with input dimension d, reverse mode scales with output dimension k. Use "reverse" when d >> k (e.g., Cartesian coordinates as input). "auto" selects reverse when d > k, forward otherwise. TYPE: Literal["forward", "reverse", "auto"], optional, by default "auto" DEFAULT: 'auto'
metric_inv	Inverse metric tensor of the input coordinate space. Supports: None: Euclidean metric, \(D = J J^T\). Shape (d,): diagonal inverse metric, shared across samples (e.g., 1 / masses per DOF). Computes \(D = J \operatorname{diag}(g) J^T\) efficiently without materializing the full matrix. Shape (N, d): per-sample diagonal inverse metric for position-dependent diffusion (e.g., estimated from local displacement covariance). Shape (d, d): full inverse metric, shared across all samples. Computes \(D = J G^{-1} J^T\). Shape (N, d, d): per-sample full inverse metric. Computes \(D_i = J_i G_i^{-1} J_i^T\). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Diffusion tensor with shape (N,) for k=1 or (N, k, k) for k>1.

Source code in spectre/compute/diffusion.py

def diffusion_tensor(
    X: torch.Tensor,
    model: torch.nn.Module,
    mode: Literal["forward", "reverse", "auto"] = "auto",
    metric_inv: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute diffusion tensor $D = J G^{-1} J^T$ from model Jacobian.

    When `metric_inv` is ``None``, assumes a Euclidean metric ($G = I$) and
    computes $D = J J^T$. When provided, computes the Fixman-corrected
    diffusion tensor $D = J G^{-1} J^T$, where $G^{-1}$ is the inverse metric
    of the input coordinate space (e.g., inverse mass matrix).

    Parameters
    ----------
    X : torch.Tensor
        Input data tensor with shape (N, d).

    model : torch.nn.Module
        Neural network model mapping R^d -> R^k.

    mode : Literal["forward", "reverse", "auto"], optional, by default "auto"
        Autodiff mode for Jacobian computation. Forward mode scales with input
        dimension `d`, reverse mode scales with output dimension `k`. Use
        ``"reverse"`` when `d` >> `k` (e.g., Cartesian coordinates as input).
        ``"auto"`` selects reverse when `d` > `k`, forward otherwise.

    metric_inv : torch.Tensor | None, optional, by default None
        Inverse metric tensor of the input coordinate space. Supports:

        - ``None``: Euclidean metric, $D = J J^T$.
        - Shape ``(d,)``: diagonal inverse metric, shared across samples (e.g.,
          ``1 / masses`` per DOF). Computes $D = J \\operatorname{diag}(g) J^T$
          efficiently without materializing the full matrix.
        - Shape ``(N, d)``: per-sample diagonal inverse metric for
          position-dependent diffusion (e.g., estimated from local displacement
          covariance).
        - Shape ``(d, d)``: full inverse metric, shared across all samples.
          Computes $D = J G^{-1} J^T$.
        - Shape ``(N, d, d)``: per-sample full inverse metric. Computes $D_i =
          J_i G_i^{-1} J_i^T$.

    Returns
    -------
    torch.Tensor
        Diffusion tensor with shape (N,) for k=1 or (N, k, k) for k>1.
    """
    jac_fn = _select_jac_fn(X, model, mode)
    batched_jac = vmap(jac_fn(model))(X)  # Shape: (N, k, d)
    if batched_jac.ndim == 3:
        batched_jac = batched_jac.squeeze(1)

    # Apply inverse metric: scale J columns for diagonal metrics,
    # or compute J @ G_inv @ J^T directly for full metrics.
    _full_metric = False
    if metric_inv is not None:
        if metric_inv.ndim == 1:
            # Shared diagonal: scale columns by sqrt(g). Shape (d,).
            scale = torch.sqrt(metric_inv.clamp(min=0.0))
            batched_jac = batched_jac * scale
        elif metric_inv.ndim == 2 and metric_inv.shape[0] == metric_inv.shape[1]:
            # Square matrix: full shared metric (d, d).
            _full_metric = True
        elif metric_inv.ndim == 2:
            # Non-square 2D: per-sample diagonal (N, d).
            scale = torch.sqrt(metric_inv.clamp(min=0.0))
            if batched_jac.ndim == 2:
                # k=1: J shape (N, d), scale shape (N, d)
                batched_jac = batched_jac * scale
            else:
                # k>1: J shape (N, k, d), scale shape (N, 1, d)
                batched_jac = batched_jac * scale.unsqueeze(1)
        elif metric_inv.ndim == 3:
            # Per-sample full metric: shape (N, d, d).
            _full_metric = True
        else:
            raise ValueError(
                f"Expected `metric_inv` to have 1-3 dimensions, got {metric_inv.ndim}."
            )

    if _full_metric and metric_inv is not None:
        # Full metric path: D = J @ G_inv @ J^T
        G_inv = metric_inv
        if batched_jac.ndim == 2:
            # k=1: J shape (N, d), treat as (N, 1, d)
            J = batched_jac.unsqueeze(1)
            JG = torch.matmul(J, G_inv)
            D = torch.matmul(JG, J.transpose(-1, -2)).squeeze(-1).squeeze(-1)
        else:
            # k>1: J shape (N, k, d)
            JG = torch.matmul(batched_jac, G_inv)  # (N, k, d)
            D = torch.matmul(JG, batched_jac.transpose(-1, -2))  # (N, k, k)
    else:
        # Euclidean or diagonal (already applied via column scaling).
        if batched_jac.ndim == 2:  # Case k=1: J shape is (N, d)
            D = torch.sum(batched_jac**2, dim=-1)  # (N,)
        else:  # Case k>1: J shape is (N, k, d)
            D = torch.matmul(batched_jac, batched_jac.transpose(-1, -2))  # (N, k, k)

    return D

log_det_diffusion_tensor(X: torch.Tensor, model: torch.nn.Module, mode: Literal['forward', 'reverse', 'auto'] = 'auto', metric_inv: torch.Tensor | None = None) -> torch.Tensor

Compute log determinant of diffusion tensor \(D = J G^{-1} J^T\).

Computes the Jacobian of a neural network model and calculates the log determinant of the diffusion tensor for invariant free energy.

PARAMETER	DESCRIPTION
X	Input data tensor with shape (N, d). TYPE: Tensor
model	Neural network model mapping R^d -> R^k. TYPE: Module
mode	Autodiff mode for Jacobian computation. Forward mode scales with input dimension d, reverse mode scales with output dimension k. Use "reverse" when d >> k (e.g., Cartesian coordinates as input). "auto" selects reverse when d > k, forward otherwise. TYPE: Literal["forward", "reverse", "auto"], optional, by default "auto" DEFAULT: 'auto'
metric_inv	Inverse metric tensor of the input coordinate space. See compute_diffusion_tensor for supported shapes. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Log determinant of diffusion tensor with shape (N,).

Examples:

>>> import torch
>>> model = torch.nn.Linear(3, 2)
>>> X = torch.randn(10, 3)
>>> log_det = compute_log_det_diffusion_tensor(X, model)
>>> log_det.shape
torch.Size([10])

Source code in spectre/compute/diffusion.py

def log_det_diffusion_tensor(
    X: torch.Tensor,
    model: torch.nn.Module,
    mode: Literal["forward", "reverse", "auto"] = "auto",
    metric_inv: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute log determinant of diffusion tensor $D = J G^{-1} J^T$.

    Computes the Jacobian of a neural network model and calculates the log
    determinant of the diffusion tensor for invariant free energy.

    Parameters
    ----------
    X : torch.Tensor
        Input data tensor with shape (N, d).

    model : torch.nn.Module
        Neural network model mapping R^d -> R^k.

    mode : Literal["forward", "reverse", "auto"], optional, by default "auto"
        Autodiff mode for Jacobian computation. Forward mode scales with input
        dimension `d`, reverse mode scales with output dimension `k`. Use
        ``"reverse"`` when `d` >> `k` (e.g., Cartesian coordinates as input).
        ``"auto"`` selects reverse when `d` > `k`, forward otherwise.

    metric_inv : torch.Tensor | None, optional, by default None
        Inverse metric tensor of the input coordinate space. See
        `compute_diffusion_tensor` for supported shapes.

    Returns
    -------
    torch.Tensor
        Log determinant of diffusion tensor with shape (N,).

    Examples
    --------
    >>> import torch
    >>> model = torch.nn.Linear(3, 2)
    >>> X = torch.randn(10, 3)
    >>> log_det = compute_log_det_diffusion_tensor(X, model)
    >>> log_det.shape
    torch.Size([10])
    """
    D = diffusion_tensor(X, model, mode=mode, metric_inv=metric_inv)

    if D.ndim == 1:  # Case k=1: scalar diffusion
        log_det_D = torch.log(D.clamp(min=torch.finfo(D.dtype).tiny))
    else:  # Case k>1: matrix diffusion
        D = D + 1e-6 * torch.eye(D.shape[-1], device=X.device)
        L = torch.linalg.cholesky(D)
        log_det_D = 2 * torch.sum(
            torch.log(
                torch.diagonal(L, dim1=-2, dim2=-1).clamp(min=torch.finfo(L.dtype).tiny)
            ),
            dim=-1,
        )

    return log_det_D

inverse_diffusion_tensor(X: torch.Tensor, model: torch.nn.Module, mode: Literal['forward', 'reverse', 'auto'] = 'auto', metric_inv: torch.Tensor | None = None) -> torch.Tensor

Compute inverse metric tensor from diffusion tensor \(D = J G^{-1} J^T\).

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples, in_features). TYPE: Tensor
model	Neural network model mapping R^d -> R^k. TYPE: Module
mode	Autodiff mode for Jacobian computation. Forward mode scales with input dimension d, reverse mode scales with output dimension k. Use "reverse" when d >> k (e.g., Cartesian coordinates as input). "auto" selects reverse when d > k, forward otherwise. TYPE: Literal["forward", "reverse", "auto"], optional, by default "auto" DEFAULT: 'auto'
metric_inv	Inverse metric tensor of the input coordinate space. See compute_diffusion_tensor for supported shapes. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Inverse metric tensors with shape (n_samples, out_features, out_features).

Source code in spectre/compute/diffusion.py

def inverse_diffusion_tensor(
    X: torch.Tensor,
    model: torch.nn.Module,
    mode: Literal["forward", "reverse", "auto"] = "auto",
    metric_inv: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute inverse metric tensor from diffusion tensor $D = J G^{-1} J^T$.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor with shape (n_samples, in_features).

    model : torch.nn.Module
        Neural network model mapping R^d -> R^k.

    mode : Literal["forward", "reverse", "auto"], optional, by default "auto"
        Autodiff mode for Jacobian computation. Forward mode scales with input
        dimension `d`, reverse mode scales with output dimension `k`. Use
        ``"reverse"`` when `d` >> `k` (e.g., Cartesian coordinates as input).
        ``"auto"`` selects reverse when `d` > `k`, forward otherwise.

    metric_inv : torch.Tensor | None, optional, by default None
        Inverse metric tensor of the input coordinate space. See
        `compute_diffusion_tensor` for supported shapes.

    Returns
    -------
    torch.Tensor
        Inverse metric tensors with shape (n_samples, out_features,
        out_features).
    """
    D = diffusion_tensor(X, model, mode=mode, metric_inv=metric_inv)

    if D.ndim == 1:  # Case out_features=1
        D_inv = (1.0 / (D + 1e-8)).unsqueeze(-1).unsqueeze(-1)
    else:  # Case out_features>1
        D = D + 1e-6 * torch.eye(D.shape[-1], device=X.device, dtype=X.dtype)
        try:
            L = torch.linalg.cholesky(D)
            D_inv = torch.cholesky_inverse(L)
        except RuntimeError:
            D_inv = torch.linalg.inv(D)

    return D_inv

decompose_diffusion_tensor(D: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]

Compute inverse and log-determinant of a diffusion tensor via single Cholesky factorization.

Avoids redundant computation when both \(D^{-1}\) and \(\log\det D\) are needed (e.g., Mahalanobis distance and invariant kernel weighting). For the matrix case (\(k > 1\)), a single Cholesky factor \(L\) yields \(D^{-1}\) via back-substitution and \(\log\det D = 2\sum\log \operatorname{diag}(L)\).

PARAMETER	DESCRIPTION
D	Diffusion tensor with shape (N,) for scalar diffusion (\(k = 1\)) or (N, k, k) for matrix diffusion (\(k > 1\)). Must be positive (semi-)definite; a small ridge is added internally for numerical stability. TYPE: Tensor

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	D_inv: Inverse diffusion tensor with shape (N, 1, 1) for \(k = 1\) or (N, k, k) for \(k > 1\). log_det_D: Log-determinant with shape (N,).

Source code in spectre/compute/diffusion.py

def decompose_diffusion_tensor(
    D: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute inverse and log-determinant of a diffusion tensor via single
    Cholesky factorization.

    Avoids redundant computation when both $D^{-1}$ and $\\log\\det D$ are
    needed (e.g., Mahalanobis distance and invariant kernel weighting). For the
    matrix case ($k > 1$), a single Cholesky factor $L$ yields $D^{-1}$ via
    back-substitution and $\\log\\det D = 2\\sum\\log \\operatorname{diag}(L)$.

    Parameters
    ----------
    D : torch.Tensor
        Diffusion tensor with shape ``(N,)`` for scalar diffusion ($k = 1$) or
        ``(N, k, k)`` for matrix diffusion ($k > 1$). Must be positive
        (semi-)definite; a small ridge is added internally for numerical
        stability.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        - ``D_inv``: Inverse diffusion tensor with shape ``(N, 1, 1)`` for
          $k = 1$ or ``(N, k, k)`` for $k > 1$.
        - ``log_det_D``: Log-determinant with shape ``(N,)``.
    """
    if D.ndim == 1:
        D_inv = (1.0 / (D + 1e-8)).unsqueeze(-1).unsqueeze(-1)
        log_det_D = torch.log(D.clamp(min=torch.finfo(D.dtype).tiny))
    else:
        D_reg = D + 1e-6 * torch.eye(D.shape[-1], device=D.device, dtype=D.dtype)
        try:
            L = torch.linalg.cholesky(D_reg)
            D_inv = torch.cholesky_inverse(L)
            log_det_D = 2.0 * torch.sum(
                torch.log(
                    torch.diagonal(L, dim1=-2, dim2=-1).clamp(
                        min=torch.finfo(L.dtype).tiny
                    )
                ),
                dim=-1,
            )
        except RuntimeError:
            D_inv = torch.linalg.inv(D_reg)
            _, log_det_D = torch.linalg.slogdet(D_reg)

    return D_inv, log_det_D