Compute Free Energy

free_energy

CLASS	DESCRIPTION
FreeEnergy	Computes free-energy landscape using kernel density estimation.

FUNCTION	DESCRIPTION
compute_kt	Compute thermal energy \(k_B T\) for a given temperature \(T\) and units.
find_local_extrema	Find local extrema (minima or maxima).

Classes

FreeEnergy(bw_method: Literal['scott', 'silverman'] = 'scott', kt: float = compute_kt(temperature=300.0, units='kj/mol'), eps: float = torch.finfo(torch.float32).eps, store_data: bool = False, grid_extent_margin: float = 0.2, device: torch.device | str | None = None)

Computes free-energy landscape using kernel density estimation.

PARAMETER	DESCRIPTION
bw_method	Bandwidth estimate procedure. TYPE: Literal["scott", "silverman"], optional, by default "scott" DEFAULT: 'scott'
kt	Energy unit for free energy. compute_kt(temperature=300.0, units="kj/mol") (~2.494) corresponds to \(k_B T\) at 300 K in kJ/mol units. TYPE: float, optional, by default `compute_kt(300.0, units="kj/mol")` DEFAULT: compute_kt(temperature=300.0, units='kj/mol')
eps	Small value to avoid numerical issues. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps
store_data	If True, stores data samples and weights. TYPE: bool, optional, by default False DEFAULT: False
grid_extent_margin	Grid margin estimated as the fraction of difference between the maximal and minimal values. When data has zero range (all identical values), uses margin of 1.0. TYPE: float, optional, by default 0.2 DEFAULT: 0.2
device	Device to run computations on. TYPE: device DEFAULT: None

ATTRIBUTE	DESCRIPTION
x_bins	Number of bins along x axis. TYPE: int
y_bins	Number of bins along y axis if free energy is computed from a 2D array. For a 1D array is None. TYPE: int | None
extent	Limit values along each dimension. If free energy is computed from a 2D array, extent is [x_min, x_max, y_min, y_max]. For a 1D array contains only [x_min, x_max]. TYPE: tuple[float, float]
grid	Grid on which free energy is computed. TYPE: Tensor

METHOD	DESCRIPTION
check_fitted	Check if FreeEnergy is fitted.
fit	Calculates the standard free energy or invariant free energy (with
fit_transform	Fit FreeEnergy and return free energy and errors.
grid_extent	Calculates grid margins for FreeEnergy for a given variable.
gradient	Compute gradient of free energy landscape.
minimum_free_energy_path	Find minimum free energy path using string methods.
compute_states	Perform state analysis using segmentation.
local_extrema	Find local extrema (minima or maxima) in log-density calculated by
barrier_heights	Calculate energy barriers between adjacent minima.
committor	Compute committor function q(x) = P(reach B before A | start at x).
free_energy_difference	Calculate free energy difference between two states.

Source code in spectre/compute/free_energy.py

def __init__(
    self,
    bw_method: Literal["scott", "silverman"] = "scott",
    kt: float = compute_kt(temperature=300.0, units="kj/mol"),
    eps: float = torch.finfo(torch.float32).eps,
    store_data: bool = False,
    grid_extent_margin: float = 0.2,
    device: torch.device | str | None = None,
) -> None:
    if isinstance(bw_method, str):
        allowed_bw_methods = ["scott", "silverman"]
        if bw_method not in allowed_bw_methods:
            raise ValueError(
                f"When str, `bw_method` must be one of {allowed_bw_methods}."
            )
        self.bw_method = bw_method

    elif isinstance(bw_method, (int, float)):
        check_in_interval(float(bw_method), "(0, inf)")
        self.bw_method = torch.tensor(bw_method)

    elif torch.is_tensor(bw_method):
        check_in_interval(float(bw_method), "(0, inf)")
        self.bw_method = bw_method

    else:
        raise TypeError(
            f"Expected `bw_method` to be a str, float, or torch.Tensor, got "
            f"{type(bw_method)}."
        )

    if not isinstance(kt, (int, float)):
        raise TypeError("Energy unit `kt` must be a number.")
    kt = float(kt)
    check_in_interval(kt, "(0, inf)")
    self.kt = kt

    if not isinstance(eps, (int, float)):
        raise TypeError(f"Expected `eps` to be a number, got {type(eps).__name__}.")
    check_in_interval(eps, "[0, 1)")
    self.eps = eps

    if not isinstance(store_data, bool):
        raise TypeError(
            f"Expected `store_data` to be bool, got {type(store_data).__name__}."
        )
    self.store_data = store_data

    if not isinstance(grid_extent_margin, float):
        raise TypeError("Grid extent margin must be a float.")
    check_in_interval(grid_extent_margin, "(0, 1)")
    self.grid_extent_margin = grid_extent_margin

    if device is None:
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        self.device = torch.device(device)

    self.is_fitted = False

    self._data: torch.Tensor | None = None
    self._weights: torch.Tensor | None = None

Attributes

data: torch.Tensor | None property

Return training data.

weights: torch.Tensor | None property

Return training sample weights.

Functions

check_fitted()

Check if FreeEnergy is fitted.

Source code in spectre/compute/free_energy.py

def check_fitted(self):
    """Check if FreeEnergy is fitted."""
    if not self.is_fitted:
        raise ValueError("Not fitted. Call `fit()` first.")

fit(X: torch.Tensor, weights: torch.Tensor | None = None, log_det_diffusion: torch.Tensor | None = None, n_splits: int | None = None, bins: tuple[int, int] | int = 100) -> FreeEnergy

Calculates the standard free energy or invariant free energy (with respect to coordinates).

The standard free energy is computed as \(F(x) = -k_B T\log P(x)\), where \(k_B\) is the Boltzmann constant, \(T\) is the temperature, and \(P(x)\) is the probability density to be estimated.

The invariant free energy is computed differently as: \(F(x) = -k_B T \log \left[ P(x) \sqrt{ \det D(x) } \right]\), where \(D(x)\) is the coordinate-dependent diffusion tensor.

PARAMETER	DESCRIPTION
X	Dataset (1d or 2d) with one or two variables. TYPE: Tensor
weights	Sample weights (1d) when data is obtained from a biased distribution. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
log_det_diffusion	Log determinant of diffusion tensor for invariant free energy. If None, it computes the standard free energy. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
n_splits	Number of dataset splits, by default None. If more than 1, errors are estimated. TYPE: int | None, optional, by default None DEFAULT: None
bins	Number of bins, either a single integer for 1D data or a tuple of two integers for 2D data. TYPE: tuple[int, int] | int, optional, by default 100 DEFAULT: 100

RETURNS	DESCRIPTION
FreeEnergy	Self.

Source code in spectre/compute/free_energy.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    log_det_diffusion: torch.Tensor | None = None,
    n_splits: int | None = None,
    bins: tuple[int, int] | int = 100,
) -> "FreeEnergy":
    """
    Calculates the standard free energy or invariant free energy (with
    respect to coordinates).

    The standard free energy is computed as $F(x) = -k_B T\\log P(x)$,
    where $k_B$ is the Boltzmann constant, $T$ is the temperature, and
    $P(x)$ is the probability density to be estimated.

    The invariant free energy is computed differently as:
    $F(x) = -k_B T \\log \\left[ P(x) \\sqrt{ \\det D(x) } \\right]$,
    where $D(x)$ is the coordinate-dependent diffusion tensor.

    Parameters
    ----------
    X : torch.Tensor
        Dataset (1d or 2d) with one or two variables.

    weights : torch.Tensor | None, optional, by default None
        Sample weights (1d) when data is obtained from a biased
        distribution.

    log_det_diffusion : torch.Tensor | None, optional, by default None
        Log determinant of diffusion tensor for invariant free energy. If
        None, it computes the standard free energy.

    n_splits : int | None, optional, by default None
        Number of dataset splits, by default None. If more than 1, errors
        are estimated.

    bins : tuple[int, int] | int, optional, by default 100
        Number of bins, either a single integer for 1D data or a tuple of
        two integers for 2D data.

    Returns
    -------
    FreeEnergy
        Self.
    """
    if not torch.is_tensor(X):
        raise ValueError(f"`X` must be `torch.Tensor`, got {type(X)}.")

    X = X.squeeze()

    if X.ndim > 2:
        raise ValueError(f"Expected `X` to have ndim 1 or 2, got {X.ndim}.")
    elif X.ndim == 2 and X.shape[1] > 2:
        raise ValueError(
            f"Expected `X` to have maximally two variables, got {X.shape[1]}."
        )
    elif X.ndim == 0:
        raise ValueError("Data `X` has 0 ndim.")

    ndim = X.ndim
    X = X.reshape(-1, ndim)

    if weights is not None:
        if not torch.is_tensor(weights):
            raise ValueError(
                f"Wrong type for `weight`. Expected `torch.Tensor`, "
                f"got {type(weights)}."
            )
        check_1d(weights)
        check_same_shape(X, weights, axis=0)
        check_sample_weights(weights)

    else:
        weights = torch.ones(X.shape[0])

    if log_det_diffusion is not None:
        if not torch.is_tensor(log_det_diffusion):
            raise ValueError(
                f"Wrong type for `log_det_diffusion`. Expected `torch.Tensor`, "
                f"got {type(log_det_diffusion)}."
            )
        check_1d(log_det_diffusion)
        check_same_shape(X, log_det_diffusion, axis=0)

        # Combine weights with diffusion tensor contribution.
        diffusion_weights = torch.exp(0.5 * log_det_diffusion)
        weights = weights * diffusion_weights

    if isinstance(bins, (list, tuple)):
        if all([isinstance(bin, int) for bin in bins]):
            self.x_bins, self.y_bins = bins
    elif isinstance(bins, int):
        self.x_bins, self.y_bins = bins, bins
    else:
        raise ValueError(
            f"Expected `bins` to be `[int, int]` or `int`, got {type(bins)}."
        )

    # Grid must be calculate differently based on `ndim` of the input array.
    self.extent = self.grid_extent(X, index=0, margin=self.grid_extent_margin)
    self.x_grid = torch.linspace(self.extent[0], self.extent[1], self.x_bins)

    if X.shape[1] == 2:
        self.ndim = 2
        self.extent += self.grid_extent(X, index=1, margin=self.grid_extent_margin)
        self.y_grid = torch.linspace(self.extent[2], self.extent[3], self.y_bins)
        x, y = torch.meshgrid(self.x_grid, self.y_grid, indexing="ij")
        self.grid = torch.cat([x.reshape(-1, 1), y.reshape(-1, 1)], dim=1)
    else:
        self.ndim = 1
        self.grid = self.x_grid.reshape(-1, 1)

    n_splits = n_splits if not (n_splits is None or n_splits < 2) else 1
    estimate = torch.zeros((n_splits, self.x_bins, 1 if ndim == 1 else self.y_bins))
    w_estimate = torch.zeros(n_splits)

    # Calculate energy in batches from the input array to estimate errors.
    x_batch = torch.tensor_split(X, n_splits)
    w_batch = torch.tensor_split(weights, n_splits)

    for i in range(n_splits):
        self.kde = KernelDensity(bw_method=self.bw_method, device=self.device)
        self.kde.fit(x_batch[i], weights=w_batch[i])

        log_density = self.kde.predict(self.grid)
        log_density = log_density.reshape(
            self.x_bins, 1 if ndim == 1 else self.y_bins
        )

        # Normalize in log-space for numerical stability
        log_eps = torch.log(torch.tensor(self.eps, device=log_density.device))
        if log_density.max() > -torch.inf:
            # Shift so max is at 0, then add epsilon in log-space
            log_density = log_density - log_density.max()
            log_density = torch.logaddexp(log_density, log_eps)

        else:
            log_density = torch.full_like(log_density, log_eps)

        # Estimate energy as the negative logarithm of the density.
        log_density = -self.kt * log_density
        estimate[i] = log_density - log_density.min()
        w_estimate[i] = w_batch[i].sum()

    if ndim == 1:
        estimate = estimate.squeeze()

    if n_splits > 1:
        n_eff = w_estimate.sum() ** 2 / torch.sum(w_estimate**2)
        w_estimate = w_estimate / w_estimate.sum()
        # Weighted average across splits
        # estimate shape: (n_splits, x_bins) for 1D
        # or (n_splits, x_bins, y_bins)
        # for 2D w_estimate shape: (n_splits,)
        if ndim == 1:
            # For 1D: (n_splits, x_bins).T @ (n_splits,) = (x_bins,)
            log_density = estimate.T @ w_estimate
        else:
            # For 2D: (n_splits, x_bins, y_bins).T @ (n_splits,) =
            # (y_bins, x_bins)
            log_density = estimate.permute(2, 1, 0) @ w_estimate

        # Fix broadcasting for error calculation.
        if ndim == 1:
            variance = (
                n_eff
                / (n_eff - 1)
                * (((log_density - estimate) ** 2).T @ w_estimate)
            )
        else:
            estimate_flat = estimate.reshape(n_splits, -1)
            diff_sq = (log_density.flatten() - estimate_flat) ** 2
            variance = n_eff / (n_eff - 1) * (diff_sq.T @ w_estimate)
            variance = variance.reshape(log_density.shape)

        sem = torch.sqrt(variance / n_eff)

    else:
        # For single split, use transpose only for 2D data
        if ndim == 1:
            log_density = estimate.squeeze()
        else:
            log_density = estimate.squeeze().T
        sem = torch.zeros_like(log_density)

    self.log_density = log_density
    self.sem = sem

    if self.store_data:
        self._data = X.clone()
        self._weights = weights.clone()

    self.is_fitted = True

    return self

fit_transform(X: torch.Tensor, weights: torch.Tensor | None = None, log_det_diffusion: torch.Tensor | None = None, n_splits: int | None = None, bins: tuple[int, int] | int = 100) -> tuple[torch.Tensor, torch.Tensor]

Fit FreeEnergy and return free energy and errors.

PARAMETER	DESCRIPTION
X	Training and transformation data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples, ). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
log_det_diffusion	Log determinant of diffusion tensor for invariant free energy. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
n_splits	Number of dataset splits for error estimation. TYPE: int | None, optional, by default None DEFAULT: None
bins	Number of bins for the grid. TYPE: tuple[int, int] | int, optional, by default 100 DEFAULT: 100

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	Free energy and errors of shape (bins, bins) if X is 2D and (bins,) if X is 1D.

Source code in spectre/compute/free_energy.py

def fit_transform(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    log_det_diffusion: torch.Tensor | None = None,
    n_splits: int | None = None,
    bins: tuple[int, int] | int = 100,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Fit `FreeEnergy` and return free energy and errors.

    Parameters
    ----------
    X : torch.Tensor
        Training and transformation data of shape (n_samples, in_features).

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples, ).

    log_det_diffusion : torch.Tensor | None, optional, by default None
        Log determinant of diffusion tensor for invariant free energy.

    n_splits : int | None, optional, by default None
        Number of dataset splits for error estimation.

    bins : tuple[int, int] | int, optional, by default 100
        Number of bins for the grid.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Free energy and errors of shape (bins, bins) if X is 2D and
        (bins,) if X is 1D.
    """
    self.fit(
        X,
        weights=weights,
        log_det_diffusion=log_det_diffusion,
        n_splits=n_splits,
        bins=bins,
    )

    return self.log_density, self.sem

grid_extent(X: torch.Tensor, index: int | None = None, margin: float = 0.2) -> tuple[float, float]

Calculates grid margins for FreeEnergy for a given variable.

PARAMETER	DESCRIPTION
X	Dataset of samples of shape (n_samples, n_features). TYPE: Tensor
index	Index of variable for which grid extent is calculated. If None, the last feature is used. TYPE: int | None, optional, by default None DEFAULT: None
margin	Margin estimated as the fraction of difference between the maximal and minimal values. When data has zero range (all identical values), uses margin of 1.0. TYPE: float, optional, by default 0.2 DEFAULT: 0.2

RETURNS	DESCRIPTION
tuple[float, float]	Extent as (min, max).

Source code in spectre/compute/free_energy.py

def grid_extent(
    self, X: torch.Tensor, index: int | None = None, margin: float = 0.2
) -> tuple[float, float]:
    """
    Calculates grid margins for `FreeEnergy` for a given variable.

    Parameters
    ----------
    X : torch.Tensor
        Dataset of samples of shape (n_samples, n_features).

    index : int | None, optional, by default None
        Index of variable for which grid extent is calculated. If None, the
        last feature is used.

    margin : float, optional, by default 0.2
        Margin estimated as the fraction of difference between the maximal
        and minimal values. When data has zero range (all identical
        values), uses margin of 1.0.

    Returns
    -------
    tuple[float, float]
        Extent as (min, max).
    """
    if X.shape[0] == 0:
        raise ValueError("Cannot calculate grid extent from empty dataset.")

    index = index if index is not None else X.shape[1] - 1

    if index < 0 or index >= X.shape[1]:
        raise ValueError(
            f"`index` {index} out of bounds for tensor with {X.shape[1]} features."
        )

    X_var = X[:, index]
    X_min, X_max = float(X_var.min()), float(X_var.max())

    data_range = X_max - X_min
    # Use margin of 1.0 if data has no range (all identical values)
    margin_size = margin * data_range if data_range > self.eps else 1.0

    return (X_min - margin_size, X_max + margin_size)

gradient() -> torch.Tensor

Compute gradient of free energy landscape.

RETURNS	DESCRIPTION
Tensor	Gradient array. Shape (x_bins,) for 1D or (y_bins, x_bins, 2) for 2D.

Source code in spectre/compute/free_energy.py

def gradient(self) -> torch.Tensor:
    """
    Compute gradient of free energy landscape.

    Returns
    -------
    torch.Tensor
        Gradient array. Shape (x_bins,) for 1D or (y_bins, x_bins, 2) for
        2D.
    """
    self.check_fitted()

    if self.ndim == 1:
        dx = self.x_grid[1] - self.x_grid[0]

        return torch.gradient(self.log_density, spacing=float(dx))[0]

    else:
        dx = self.x_grid[1] - self.x_grid[0]
        dy = self.y_grid[1] - self.y_grid[0]

        grad_y, grad_x = torch.gradient(
            self.log_density, spacing=(float(dy), float(dx))
        )

        return torch.stack([grad_x, grad_y], dim=-1)

minimum_free_energy_path(start: torch.Tensor, end: torch.Tensor, n_samples: int = 50, lr: float = 0.05, max_iterations: int = 2000, tolerance: float = 1e-05, spacing_constant: float = 0.1, method: str = 'string', climb: bool = False, bw: float | None = None, k_nearest: int = 1) -> dict[str, torch.Tensor]

Find minimum free energy path using string methods.

This method finds a smooth path connecting two points by:

1. Using springs to maintain equal spacing along the path
2. Descending perpendicular to path (string)
3. Optionally using climbing image to find saddle points
4. Iteratively relaxing the path until convergence

For high energy barriers, use climb=True to prevent the path from sliding into nearby minima. As of now, this is experimental.

PARAMETER	DESCRIPTION
start	Starting coordinates, shape (ndim,). TYPE: Tensor
end	Ending coordinates, shape (ndim,). TYPE: Tensor
n_samples	Number of samples along the path for optimization. TYPE: int, optional, by default 50 DEFAULT: 50
lr	Step size for gradient descent. TYPE: float, optional, by default 0.05 DEFAULT: 0.05
max_iterations	Maximum number of optimization iterations. TYPE: int, optional, by default 2000 DEFAULT: 2000
tolerance	Convergence tolerance for path optimization. TYPE: float, optional, by default 1e-5 DEFAULT: 1e-05
spacing_constant	Constant for maintaining equal spacing. TYPE: float, optional, by default 0.1 DEFAULT: 0.1
method	Path optimization method: "string": Perpendicular gradient descent using grid gradients "fts": Finite temperature string using mean force from samples TYPE: str, optional, by default "string" DEFAULT: 'string'
climb	If True, the highest energy image climbs uphill to find the saddle point. This prevents the path from sliding into local minima. TYPE: bool, optional, by default False DEFAULT: False
bw	Bandwidth for mean force computation in FTS method. If None, uses the KDE bandwidth. Only used when method="fts". TYPE: float | None, optional, by default None DEFAULT: None
k_nearest	Number of nearest data samples to return for each path point. If k_nearest=1, returns shape (n_samples,). If k_nearest>1, returns shape (n_samples, k_nearest). Only included if store_data=True. TYPE: int, optional, by default 1 DEFAULT: 1

RETURNS

DESCRIPTION

dict[str, Tensor]

Dictionary with keys: 'path': Path coordinates (n_samples, ndim) 'energies': Path free energies (n_samples,) 'sample_indices': Indices of sampled points (n_samples,) 'sample_coords': Coordinates of sampled points (n_samples, ndim) 'sample_energies': Energies at sampled points (n_samples,) 'nearest_data_indices': Indices of k nearest data samples to sample_coords. Shape (n_samples,) if k_nearest=1, or (n_samples, k_nearest) if k_nearest>1. Use with original data: data[nearest_data_indices].

Source code in spectre/compute/free_energy.py

def minimum_free_energy_path(
    self,
    start: torch.Tensor,
    end: torch.Tensor,
    n_samples: int = 50,
    lr: float = 0.05,
    max_iterations: int = 2000,
    tolerance: float = 1e-5,
    spacing_constant: float = 0.1,
    method: str = "string",
    climb: bool = False,
    bw: float | None = None,
    k_nearest: int = 1,
) -> dict[str, torch.Tensor]:
    """
    Find minimum free energy path using string methods.

    This method finds a smooth path connecting two points by:

    1. Using springs to maintain equal spacing along the path
    2. Descending perpendicular to path (string)
    3. Optionally using climbing image to find saddle points
    4. Iteratively relaxing the path until convergence

    For high energy barriers, use `climb=True` to prevent the path from
    sliding into nearby minima. As of now, this is experimental.

    Parameters
    ----------
    start : torch.Tensor
        Starting coordinates, shape (ndim,).

    end : torch.Tensor
        Ending coordinates, shape (ndim,).

    n_samples : int, optional, by default 50
        Number of samples along the path for optimization.

    lr : float, optional, by default 0.05
        Step size for gradient descent.

    max_iterations : int, optional, by default 2000
        Maximum number of optimization iterations.

    tolerance : float, optional, by default 1e-5
        Convergence tolerance for path optimization.

    spacing_constant : float, optional, by default 0.1
        Constant for maintaining equal spacing.

    method : str, optional, by default "string"
        Path optimization method:

        - "string": Perpendicular gradient descent using grid gradients
        - "fts": Finite temperature string using mean force from samples

    climb : bool, optional, by default False
        If True, the highest energy image climbs uphill to find the saddle
        point. This prevents the path from sliding into local minima.

    bw : float | None, optional, by default None
        Bandwidth for mean force computation in FTS method. If None, uses
        the KDE bandwidth. Only used when `method="fts"`.

    k_nearest : int, optional, by default 1
        Number of nearest data samples to return for each path point. If
        k_nearest=1, returns shape (n_samples,). If k_nearest>1, returns
        shape (n_samples, k_nearest). Only included if `store_data=True`.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary with keys:

        - 'path': Path coordinates (n_samples, ndim)
        - 'energies': Path free energies (n_samples,)
        - 'sample_indices': Indices of sampled points (n_samples,)
        - 'sample_coords': Coordinates of sampled points (n_samples, ndim)
        - 'sample_energies': Energies at sampled points (n_samples,)
        - 'nearest_data_indices': Indices of k nearest data samples to
          sample_coords. Shape (n_samples,) if k_nearest=1, or
          (n_samples, k_nearest) if k_nearest>1. Use with original data:
          data[nearest_data_indices].
    """
    self.check_fitted()

    allowed_methods = ["string", "fts"]
    if method not in allowed_methods:
        raise ValueError(f"method must be one of {allowed_methods}, got '{method}'")

    if method == "fts" and not self.store_data:
        raise ValueError("FTS method requires samples to be stored.")

    if self.ndim == 1:
        # For 1D, path is simply interpolation between start and end
        path = torch.linspace(start, end, n_samples).unsqueeze(1)
        # Interpolate free energies on grid
        energies = torch.zeros(n_samples)

        for i, coord in enumerate(path):
            # Find nearest grid sample
            idx = torch.argmin(torch.abs(self.x_grid - coord))
            energies[i] = self.log_density[idx]

        # For 1D, use uniform spacing (since path is already linear)
        sample_indices = torch.linspace(
            0, n_samples - 1, n_samples, dtype=torch.long
        )
        sample_coords = path[sample_indices].squeeze()
        sample_energies = energies[sample_indices]

        result = {
            "path": path.squeeze(),
            "energies": energies,
            "sample_indices": sample_indices,
            "sample_coords": sample_coords,
            "sample_energies": sample_energies,
        }

        # Add nearest data indices if samples are stored
        if self.store_data:
            result["nearest_data_indices"] = self._find_nearest_samples(
                path, k=k_nearest
            )

        return result

    else:
        # For 2D, use string method with perpendicular gradient descent
        # Initialize path as linear interpolation
        path = torch.zeros(n_samples, 2, requires_grad=False)
        for i in range(n_samples):
            alpha = i / (n_samples - 1)
            path[i] = (1 - alpha) * start + alpha * end

        # Compute gradients at all grid points for interpolation
        # Note: gradient field is computed once and remains static during
        # optimization (except for FTS which computes on-the-fly).
        if method != "fts":
            grad_field = self.gradient()  # Shape: (y_bins, x_bins, 2)
        else:
            grad_field = None

        for iteration in range(max_iterations):
            prev_path = path.clone()

            # Compute energies if using climbing
            if climb:
                energies = self._interpolate_energy_2d(path[1:-1])
                max_energy_idx = torch.argmax(energies).item() + 1

            # Compute tangent directions for all interior points
            interior_points = path[1:-1]
            tangents = path[2:] - path[:-2]
            tangents = tangents / (
                torch.norm(tangents, dim=1, keepdim=True) + 1e-10
            )

            # Batch interpolate gradients for string method
            if method == "fts":
                grads = torch.zeros_like(interior_points)
                for i in range(len(interior_points)):
                    grads[i] = self._compute_mean_force_from_samples(
                        interior_points[i], bw=bw
                    )
            else:
                grads = self._interpolate_gradient_2d(interior_points, grad_field)

            # Normalize gradients to unit vectors to prevent instability.
            # Without normalization, gradient magnitude depends on grid
            # resolution and data scale, causing unstable step sizes that
            # make paths wander.
            grad_norms = torch.norm(grads, dim=1, keepdim=True)
            grads_normalized = grads / (grad_norms + 1e-10)

            # Compute perpendicular gradient components using normalized
            # gradients
            grad_parallel = (grads_normalized * tangents).sum(
                dim=1, keepdim=True
            ) * tangents
            grads_perp = grads_normalized - grad_parallel

            f_spring_next = path[2:] - interior_points
            f_spring_prev = path[:-2] - interior_points
            f_spring = spacing_constant * (f_spring_next + f_spring_prev)

            # Note: self.log_density stores F = -kT * log(P), so grads =
            # gradient(F). To descend the free energy landscape, move
            # opposite to gradient: -grads_perp
            forces = -grads_perp

            if climb:
                # Climbing image: remove perpendicular component, climb
                # uphill along tangent
                max_idx = max_energy_idx - 1  # Adjust for interior indexing
                forces[max_idx] = -grads_perp[max_idx] + 2 * grad_parallel[max_idx]

            # Update interior points
            path[1:-1] = interior_points + lr * forces + lr * f_spring

            # Clip path to stay within grid bounds
            path[:, 0] = torch.clamp(path[:, 0], self.extent[0], self.extent[1])
            path[:, 1] = torch.clamp(path[:, 1], self.extent[2], self.extent[3])  # ty: ignore[index-out-of-bounds]

            # Reparametrize to maintain equal arc-length spacing
            path = self._reparametrize_path(path)

            displacement = torch.norm(path - prev_path)
            if displacement < tolerance:
                break

        energies = self._interpolate_energy_2d(path)

        # Compute cumulative arc length along path
        segment_lengths = torch.norm(path[1:] - path[:-1], dim=1)
        arc_lengths = torch.cat(
            [torch.zeros(1), torch.cumsum(segment_lengths, dim=0)]
        )

        total_length = arc_lengths[-1]
        target_lengths = torch.linspace(0, total_length, n_samples)

        # Find indices closest to target arc lengths
        sample_indices = torch.zeros(n_samples, dtype=torch.long)
        sample_coords = torch.zeros(n_samples, path.shape[1])
        sample_energies = torch.zeros(n_samples)

        for i, target in enumerate(target_lengths):
            # Find the closest point on path by arc length
            idx = torch.argmin(torch.abs(arc_lengths - target))
            sample_indices[i] = idx
            sample_coords[i] = path[idx]
            sample_energies[i] = energies[idx]

        result = {
            "path": path,
            "energies": energies,
            "sample_indices": sample_indices,
            "sample_coords": sample_coords,
            "sample_energies": sample_energies,
        }

        if self.store_data:
            # Find nearest samples for the arc-length sampled coords, not
            # full path
            result["nearest_data_indices"] = self._find_nearest_samples(
                sample_coords, k=k_nearest
            )

        return result

compute_states(saddle_threshold: float | None = None, min_distance: int = 1) -> dict[str, torch.Tensor]

Perform state analysis using segmentation.

PARAMETER	DESCRIPTION
saddle_threshold	Energy threshold for saddle points. If None, uses median energy. TYPE: float | None, optional, by default None DEFAULT: None
min_distance	Minimum distance between states. TYPE: int, optional, by default 1 DEFAULT: 1

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary with keys: 'state_labels': State assignment for each grid point 'state_populations': Population fraction in each state 'n_states': Number of states identified

Source code in spectre/compute/free_energy.py

def compute_states(
    self,
    saddle_threshold: float | None = None,
    min_distance: int = 1,
) -> dict[str, torch.Tensor]:
    """
    Perform state analysis using segmentation.

    Parameters
    ----------
    saddle_threshold : float | None, optional, by default None
        Energy threshold for saddle points. If None, uses median energy.

    min_distance : int, optional, by default 1
        Minimum distance between states.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary with keys:

        - 'state_labels': State assignment for each grid point
        - 'state_populations': Population fraction in each state
        - 'n_states': Number of states identified
    """
    self.check_fitted()

    if saddle_threshold is None:
        saddle_threshold = float(self.log_density.median())

    # Find local minima
    minima_coords, minima_vals = self.local_extrema(
        mode="minima", min_distance=min_distance, threshold=saddle_threshold
    )

    n_states = len(minima_vals)

    if self.ndim == 1:
        # For 1D, assign each point to nearest minimum
        state_labels = torch.zeros(self.x_bins, dtype=torch.long)
        for i, x in enumerate(self.x_grid):
            if n_states > 0:
                distances = torch.abs(minima_coords - x)
                state_labels[i] = torch.argmin(distances)
            else:
                state_labels[i] = 0

        # Compute populations (using Boltzmann weights in log-space)
        log_populations = torch.zeros(max(n_states, 1))
        for state_id in range(max(n_states, 1)):
            mask = state_labels == state_id
            # Population proportional to exp(-F) via logsumexp
            log_populations[state_id] = torch.logsumexp(
                -self.log_density[mask], dim=0
            )

        # Normalize in log-space then convert
        log_Z_total = torch.logsumexp(log_populations, dim=0)
        populations = torch.exp(log_populations - log_Z_total)

    else:
        # For 2D, use vectorized distance-based assignment
        state_labels = torch.zeros((self.y_bins, self.x_bins), dtype=torch.long)

        if n_states > 0:
            # Create meshgrid of all points
            xx, yy = torch.meshgrid(self.x_grid, self.y_grid, indexing="ij")
            grid_points = torch.stack(
                [xx.flatten(), yy.flatten()], dim=1
            )  # (n_points, 2)

            # Broadcast: (n_points, 1, 2) - (1, n_minima, 2) -> (n_points,
            # n_minima)
            distances = torch.norm(
                grid_points.unsqueeze(1) - minima_coords.unsqueeze(0), dim=2
            )
            state_labels_flat = torch.argmin(distances, dim=1)
            state_labels = state_labels_flat.reshape(self.y_bins, self.x_bins).T

        log_populations = torch.zeros(max(n_states, 1))
        for state_id in range(max(n_states, 1)):
            mask = state_labels == state_id
            # Population proportional to exp(-F) via logsumexp
            log_populations[state_id] = torch.logsumexp(
                -self.log_density[mask], dim=0
            )

        log_Z_total = torch.logsumexp(log_populations, dim=0)
        populations = torch.exp(log_populations - log_Z_total)

    return {
        "state_labels": state_labels,
        "state_populations": populations,
        "n_states": n_states,
    }

local_extrema(mode: Literal['minima', 'maxima'] = 'minima', min_distance: int = 1, threshold: float | None = None) -> tuple[torch.Tensor, torch.Tensor]

Find local extrema (minima or maxima) in log-density calculated by FreeEnergy.

PARAMETER	DESCRIPTION
mode	Whether to find local minima or maxima. TYPE: Literal["minima", "maxima"], optional, by default "minima" DEFAULT: 'minima'
min_distance	Minimum distance between extrema. For 1D, this is in grid points. For 2D, defines the neighborhood size for local filtering. TYPE: int, optional, by default 1 DEFAULT: 1
threshold	Energy threshold for filtering results. For minima, returns values below threshold. For maxima, returns values above threshold. If None, all local extrema are returned. TYPE: float | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	For 1D: (coordinates, log_density_values) where coordinates are real-valued positions on the grid. For 2D: (coordinates, log_density_values) where coordinates is shape (n, 2) with (x, y) real-valued positions.

Source code in spectre/compute/free_energy.py

def local_extrema(
    self,
    mode: Literal["minima", "maxima"] = "minima",
    min_distance: int = 1,
    threshold: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Find local extrema (minima or maxima) in log-density calculated by
    `FreeEnergy`.

    Parameters
    ----------
    mode : Literal["minima", "maxima"], optional, by default "minima"
        Whether to find local minima or maxima.

    min_distance : int, optional, by default 1
        Minimum distance between extrema. For 1D, this is in grid points.
        For 2D, defines the neighborhood size for local filtering.

    threshold : float | None, optional, by default None
        Energy threshold for filtering results. For minima, returns values
        below threshold. For maxima, returns values above threshold. If
        None, all local extrema are returned.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        For 1D: (coordinates, log_density_values) where coordinates are
        real-valued positions on the grid.
        For 2D: (coordinates, log_density_values) where coordinates is
        shape (n, 2) with (x, y) real-valued positions.
    """
    self.check_fitted()

    return find_local_extrema(
        self.log_density,
        x_grid=self.x_grid,
        y_grid=self.y_grid if self.ndim == 2 else None,
        mode=mode,
        min_distance=min_distance,
        threshold=threshold,
    )

barrier_heights(minima_coords: torch.Tensor | None = None, min_distance: int = 1, threshold: float | None = None, n_samples: int = 50, lr: float = 0.05, max_iterations: int = 2000, tolerance: float = 1e-05, spacing_constant: float = 0.1) -> torch.Tensor

Calculate energy barriers between adjacent minima.

For each pair of adjacent minima, finds the saddle point (lowest maximum) along the minimum free energy path and returns barrier height.

PARAMETER	DESCRIPTION
minima_coords	Coordinates of minima. If None, automatically finds local minima. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
min_distance	Minimum distance between minima. TYPE: int, optional, by default 1 DEFAULT: 1
threshold	Energy threshold for filtering results. For minima, returns values below threshold. For maxima, returns values above threshold. If None, all local extrema are returned. TYPE: float | None, optional, by default None DEFAULT: None
n_samples	Number of samples along the path for optimization. TYPE: int, optional, by default 50 DEFAULT: 50
lr	Step size for gradient descent. TYPE: float, optional, by default 0.05 DEFAULT: 0.05
max_iterations	Maximum number of optimization iterations. TYPE: int, optional, by default 2000 DEFAULT: 2000
tolerance	Convergence tolerance for path optimization. TYPE: float, optional, by default 1e-5 DEFAULT: 1e-05
spacing_constant	Spring force constant for maintaining equal spacing. TYPE: float, optional, by default 0.1 DEFAULT: 0.1

RETURNS	DESCRIPTION
Tensor	Barrier heights for each pair of adjacent minima.

Source code in spectre/compute/free_energy.py

def barrier_heights(
    self,
    minima_coords: torch.Tensor | None = None,
    min_distance: int = 1,
    threshold: float | None = None,
    n_samples: int = 50,
    lr: float = 0.05,
    max_iterations: int = 2000,
    tolerance: float = 1e-5,
    spacing_constant: float = 0.1,
) -> torch.Tensor:
    """
    Calculate energy barriers between adjacent minima.

    For each pair of adjacent minima, finds the saddle point (lowest
    maximum) along the minimum free energy path and returns barrier height.

    Parameters
    ----------
    minima_coords : torch.Tensor | None, optional, by default None
        Coordinates of minima. If None, automatically finds local minima.

    min_distance : int, optional, by default 1
        Minimum distance between minima.

    threshold : float | None, optional, by default None
        Energy threshold for filtering results. For minima, returns values
        below threshold. For maxima, returns values above threshold. If
        None, all local extrema are returned.

    n_samples : int, optional, by default 50
        Number of samples along the path for optimization.

    lr : float, optional, by default 0.05
        Step size for gradient descent.

    max_iterations : int, optional, by default 2000
        Maximum number of optimization iterations.

    tolerance : float, optional, by default 1e-5
        Convergence tolerance for path optimization.

    spacing_constant : float, optional, by default 0.1
        Spring force constant for maintaining equal spacing.

    Returns
    -------
    torch.Tensor
        Barrier heights for each pair of adjacent minima.
    """
    self.check_fitted()

    if minima_coords is None:
        minima_coords, _ = self.local_extrema(
            mode="minima",
            min_distance=min_distance,
            threshold=threshold,
        )

    n_minima = len(minima_coords)
    if n_minima < 2:
        return torch.tensor([])

    if self.ndim == 1:
        # Sort minima by coordinate
        sorted_idx = torch.argsort(minima_coords)
        sorted_coords = minima_coords[sorted_idx]

        barriers = []
        for i in range(n_minima - 1):
            # Find maximum between consecutive minima
            start_idx = torch.argmin(torch.abs(self.x_grid - sorted_coords[i]))
            end_idx = torch.argmin(torch.abs(self.x_grid - sorted_coords[i + 1]))

            if start_idx > end_idx:
                start_idx, end_idx = end_idx, start_idx

            if start_idx == end_idx:
                # Single point segment
                barrier = self.log_density[start_idx].item()
            else:
                segment = self.log_density[start_idx : end_idx + 1]
                barrier = segment.max().item()
            barriers.append(barrier)

        return torch.tensor(barriers)
    else:
        # For 2D, compute barriers between all pairs using MFEP
        barriers = []
        for i in range(n_minima):
            for j in range(i + 1, n_minima):
                result = self.minimum_free_energy_path(
                    minima_coords[i],
                    minima_coords[j],
                    n_samples=n_samples,
                    lr=lr,
                    max_iterations=max_iterations,
                    tolerance=tolerance,
                    spacing_constant=spacing_constant,
                )

                barrier = result["energies"].max().item()
                barriers.append(barrier)

        return torch.tensor(barriers)

committor(state_A: torch.Tensor, state_B: torch.Tensor, n_iterations: int = 1000, tolerance: float = 1e-05) -> torch.Tensor

Compute committor function q(x) = P(reach B before A | start at x).

Solves the boundary value problem on the free energy landscape.

PARAMETER	DESCRIPTION
state_A	Coordinates defining state A boundary. TYPE: Tensor
state_B	Coordinates defining state B boundary. TYPE: Tensor
n_iterations	Maximum number of iterations for solving BVP. TYPE: int, optional, by default 1000 DEFAULT: 1000
tolerance	Convergence tolerance. TYPE: float, optional, by default 1e-5 DEFAULT: 1e-05

RETURNS	DESCRIPTION
Tensor	Committor values on the grid. Shape matches log_density.

Source code in spectre/compute/free_energy.py

def committor(
    self,
    state_A: torch.Tensor,
    state_B: torch.Tensor,
    n_iterations: int = 1000,
    tolerance: float = 1e-5,
) -> torch.Tensor:
    """
    Compute committor function q(x) = P(reach B before A | start at x).

    Solves the boundary value problem on the free energy landscape.

    Parameters
    ----------
    state_A : torch.Tensor
        Coordinates defining state A boundary.

    state_B : torch.Tensor
        Coordinates defining state B boundary.

    n_iterations : int, optional, by default 1000
        Maximum number of iterations for solving BVP.

    tolerance : float, optional, by default 1e-5
        Convergence tolerance.

    Returns
    -------
    torch.Tensor
        Committor values on the grid. Shape matches log_density.
    """
    self.check_fitted()

    if self.ndim == 1:
        if not torch.is_tensor(state_A) or state_A.numel() != 1:
            raise ValueError(
                "For 1D free energy, state_A must be a scalar tensor with 1 element"
            )
        if not torch.is_tensor(state_B) or state_B.numel() != 1:
            raise ValueError(
                "For 1D free energy, state_B must be a scalar tensor with 1 element"
            )

        state_A_scalar = state_A.item() if state_A.numel() == 1 else float(state_A)
        state_B_scalar = state_B.item() if state_B.numel() == 1 else float(state_B)

        q = torch.zeros_like(self.log_density)

        # Set boundary conditions
        # Find grid points in state A and B
        dist_A = torch.abs(self.x_grid - state_A_scalar)
        dist_B = torch.abs(self.x_grid - state_B_scalar)

        idx_A = torch.argmin(dist_A)
        idx_B = torch.argmin(dist_B)

        # Boundary conditions: q(A) = 0, q(B) = 1
        q[idx_A] = 0.0
        q[idx_B] = 1.0

        # Solve Laplace equation with Jacobi iteration
        for _ in range(n_iterations):
            q_prev = q.clone()

            for i in range(1, len(q) - 1):
                if i == idx_A or i == idx_B:
                    continue
                # Finite difference: d^2q/dx^2 = 0
                q[i] = 0.5 * (q[i - 1] + q[i + 1])

            if torch.norm(q - q_prev) < tolerance:
                break

        return q
    else:
        q = torch.zeros_like(self.log_density)

        # Set boundary conditions for 2D
        for i in range(self.x_bins):
            for j in range(self.y_bins):
                point = torch.tensor([self.x_grid[i], self.y_grid[j]])

                dist_A = torch.norm(point - state_A)
                dist_B = torch.norm(point - state_B)

                # Simple initial guess based on distance
                q[j, i] = dist_A / (dist_A + dist_B)

        for _ in range(n_iterations):
            q_prev = q.clone()

            for i in range(1, self.x_bins - 1):
                for j in range(1, self.y_bins - 1):
                    # Laplace equation: d^2q/dx^2 + d^2q/dy^2 = 0
                    q[j, i] = 0.25 * (
                        q[j, i - 1] + q[j, i + 1] + q[j - 1, i] + q[j + 1, i]
                    )

            if torch.norm(q - q_prev) < tolerance:
                break

        return q

free_energy_difference(state_A: torch.Tensor | None = None, state_B: torch.Tensor | None = None, state_labels: torch.Tensor | None = None, state_id_A: int = 0, state_id_B: int = 1, energy_threshold: float | None = None) -> dict[str, torch.Tensor]

Calculate free energy difference between two states.

Uses Boltzmann weighting to compute: \(\Delta F_{AB} = F_A - F_B = -kT \ln(Z_A/Z_B)\) where \(Z_i = \int_{state_i} \exp(-F(x)/kT) dx\)

PARAMETER	DESCRIPTION
state_A	Coordinates defining state A center. If None, uses state_labels. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
state_B	Coordinates defining state B center. If None, uses state_labels. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
state_labels	Pre-computed state labels from compute_states(). If None and state centers not provided, automatically computes state assignments. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
state_id_A	state ID for state A (when using state_labels). TYPE: int, optional, by default 0 DEFAULT: 0
state_id_B	state ID for state B (when using state_labels). TYPE: int, optional, by default 1 DEFAULT: 1
energy_threshold	Energy threshold relative to state minimum. Only grid points within this energy from the state minimum are included in the partition function calculation. If None, all points in the state are used. TYPE: float | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary containing: 'delta_F': Free energy difference \(\Delta F_{AB} = F_A - F_B\) 'Z_A': Partition function of state A 'Z_B': Partition function of state B 'prob_A': Probability of being in state A 'prob_B': Probability of being in state B 'population_ratio': \(Z_B / Z_A = exp(-\Delta F_{AB}/kT)\)

Examples:

Using state centers:

>>> fel = FreeEnergy()
>>> fel.fit(X, bins=(50, 50))
>>> result = fel.free_energy_difference(
...     state_A=torch.tensor([-2.0, -2.0]),
...     state_B=torch.tensor([2.0, 2.0]),
... )
>>> print(f"ΔF = {result['delta_F']:.2f}")

Using pre-computed state labels:

>>> states = fel.compute_states()
>>> result = fel.free_energy_difference(
...     state_labels=states["state_labels"], state_id_A=0, state_id_B=1
... )

Source code in spectre/compute/free_energy.py

def free_energy_difference(
    self,
    state_A: torch.Tensor | None = None,
    state_B: torch.Tensor | None = None,
    state_labels: torch.Tensor | None = None,
    state_id_A: int = 0,
    state_id_B: int = 1,
    energy_threshold: float | None = None,
) -> dict[str, torch.Tensor]:
    """
    Calculate free energy difference between two states.

    Uses Boltzmann weighting to compute:
    $\\Delta F_{AB} = F_A - F_B = -kT \\ln(Z_A/Z_B)$
    where $Z_i = \\int_{state_i} \\exp(-F(x)/kT) dx$

    Parameters
    ----------
    state_A : torch.Tensor | None, optional, by default None
        Coordinates defining state A center. If None, uses state_labels.

    state_B : torch.Tensor | None, optional, by default None
        Coordinates defining state B center. If None, uses state_labels.

    state_labels : torch.Tensor | None, optional, by default None
        Pre-computed state labels from compute_states(). If None and state
        centers not provided, automatically computes state assignments.

    state_id_A : int, optional, by default 0
        state ID for state A (when using state_labels).

    state_id_B : int, optional, by default 1
        state ID for state B (when using state_labels).

    energy_threshold : float | None, optional, by default None
        Energy threshold relative to state minimum. Only grid points within
        this energy from the state minimum are included in the partition
        function calculation. If None, all points in the state are used.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary containing:

        - 'delta_F': Free energy difference $\\Delta F_{AB} = F_A - F_B$
        - 'Z_A': Partition function of state A
        - 'Z_B': Partition function of state B
        - 'prob_A': Probability of being in state A
        - 'prob_B': Probability of being in state B
        - 'population_ratio': $Z_B / Z_A = exp(-\\Delta F_{AB}/kT)$

    Examples
    --------
    Using state centers:

    >>> fel = FreeEnergy()
    >>> fel.fit(X, bins=(50, 50))
    >>> result = fel.free_energy_difference(
    ...     state_A=torch.tensor([-2.0, -2.0]),
    ...     state_B=torch.tensor([2.0, 2.0]),
    ... )
    >>> print(f"ΔF = {result['delta_F']:.2f}")

    Using pre-computed state labels:

    >>> states = fel.compute_states()
    >>> result = fel.free_energy_difference(
    ...     state_labels=states["state_labels"], state_id_A=0, state_id_B=1
    ... )
    """
    self.check_fitted()

    if state_labels is None:
        if state_A is None or state_B is None:
            state_result = self.compute_states()
            state_labels = state_result["state_labels"]

            # Use first two states by default
            state_id_A = 0
            state_id_B = min(1, state_result["n_states"] - 1)

        else:
            # Assign each grid point to nearest state center
            state_labels = torch.zeros_like(self.log_density, dtype=torch.long)

            if self.ndim == 1:
                for i, x in enumerate(self.x_grid):
                    dist_A = abs(x - state_A)
                    dist_B = abs(x - state_B)
                    state_labels[i] = 0 if dist_A < dist_B else 1
                state_id_A = 0
                state_id_B = 1

            else:
                for i, x in enumerate(self.x_grid):
                    for j, y in enumerate(self.y_grid):
                        point = torch.tensor([x, y])
                        dist_A = torch.norm(point - state_A)
                        dist_B = torch.norm(point - state_B)
                        state_labels[j, i] = 0 if dist_A < dist_B else 1
                state_id_A = 0
                state_id_B = 1

    # Compute grid spacing for integration
    if self.ndim == 1:
        dx = self.x_grid[1] - self.x_grid[0]
        grid_volume = float(dx.item())
    else:
        dx = self.x_grid[1] - self.x_grid[0]
        dy = self.y_grid[1] - self.y_grid[0]
        grid_volume = float((dx * dy).item())

    # Extract regions for each state
    mask_A = state_labels == state_id_A
    mask_B = state_labels == state_id_B

    # Apply energy threshold if specified
    if energy_threshold is not None:
        # Find minimum energy in each state
        energy_A = self.log_density[mask_A]
        energy_B = self.log_density[mask_B]

        min_energy_A = energy_A.min()
        min_energy_B = energy_B.min()

        energy_mask_A = (
            self.log_density <= min_energy_A + energy_threshold
        ) & mask_A
        energy_mask_B = (
            self.log_density <= min_energy_B + energy_threshold
        ) & mask_B

        # Update masks to include only low-energy points
        mask_A = energy_mask_A
        mask_B = energy_mask_B

    # Z = int exp(-F/kT) dx = int p(x) dx -> log(Z) = log(int p(x) dx)
    log_boltzmann_A = -self.log_density[mask_A] / self.kt
    log_boltzmann_B = -self.log_density[mask_B] / self.kt

    log_grid_volume = torch.log(torch.tensor(grid_volume))

    log_Z_A = torch.logsumexp(log_boltzmann_A, dim=0) + log_grid_volume
    log_Z_B = torch.logsumexp(log_boltzmann_B, dim=0) + log_grid_volume

    if not torch.isfinite(log_Z_A) or not torch.isfinite(log_Z_B):
        raise ValueError(
            "Partition function calculation resulted in non-finite values. "
            "Check state definitions or energy values."
        )

    # Compute delta_F directly in log space
    # F_AB = F_A - F_B = -kT * ln(Z_A/Z_B)
    delta_F = -self.kt * (log_Z_A - log_Z_B)

    Z_A = torch.exp(log_Z_A)
    Z_B = torch.exp(log_Z_B)

    log_Z_total = torch.logaddexp(log_Z_A, log_Z_B)
    prob_A = torch.exp(log_Z_A - log_Z_total)
    prob_B = torch.exp(log_Z_B - log_Z_total)

    return {
        "delta_F": delta_F,
        "Z_A": Z_A,
        "Z_B": Z_B,
        "prob_A": prob_A,
        "prob_B": prob_B,
        "population_ratio": Z_B / Z_A,
    }

Functions

compute_kt(temperature: float = 300.0, units: str = 'kj/mol') -> float

Compute thermal energy \(k_B T\) for a given temperature \(T\) and units.

PARAMETER	DESCRIPTION
temperature	Temperature in Kelvin. TYPE: float, optional, by default 300.0 DEFAULT: 300.0
units	Energy units for the output. Options: "kj/mol": kilojoules per mole "kcal/mol": kilocalories per mole "ev": electron volts "j": joules (SI units) TYPE: str, optional, by default "kj/mol" DEFAULT: 'kj/mol'

RETURNS	DESCRIPTION
float	Thermal energy in the specified units.

Source code in spectre/compute/free_energy.py

def compute_kt(temperature: float = 300.0, units: str = "kj/mol") -> float:
    """
    Compute thermal energy $k_B T$ for a given temperature $T$ and units.

    Parameters
    ----------
    temperature : float, optional, by default 300.0
        Temperature in Kelvin.

    units : str, optional, by default "kj/mol"
        Energy units for the output. Options:

        - "kj/mol": kilojoules per mole
        - "kcal/mol": kilocalories per mole
        - "ev": electron volts
        - "j": joules (SI units)

    Returns
    -------
    float
        Thermal energy in the specified units.
    """
    units_lower = units.lower().replace("/", "").replace(" ", "")

    # Map unit strings to Boltzmann constants
    kb_values = {
        "kjmol": 0.0083144621,  # kJ/(mol·K)
        "kcalmol": 0.0019872041,  # kcal/(mol·K)
        "ev": 8.617333262e-5,  # eV/K
        "j": 1.380649e-23,  # J/K
    }

    if units_lower not in kb_values:
        raise ValueError(
            f"Unknown units '{units}'. Supported: {', '.join(kb_values.keys())}."
        )

    return kb_values[units_lower] * temperature

find_local_extrema(log_density: torch.Tensor, x_grid: torch.Tensor, y_grid: torch.Tensor | None = None, mode: Literal['minima', 'maxima'] = 'minima', min_distance: int = 1, threshold: float | None = None) -> tuple[torch.Tensor, torch.Tensor]

Find local extrema (minima or maxima).

PARAMETER	DESCRIPTION
mode	Whether to find local minima or maxima. TYPE: Literal["minima", "maxima"], optional, by default "minima" DEFAULT: 'minima'
x_grid	Grid on which free energy is computed along x axis TYPE: Tensor
y_grid	Grid on which free energy is computed along y axis if free energy is computed from a 2D array. For a 1D array is None. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
min_distance	Minimum distance between extrema. For 1D, this is in grid points. For 2D, defines the neighborhood size for local filtering. TYPE: int, optional, by default 1 DEFAULT: 1
threshold	Energy threshold for filtering results. For minima, returns values below threshold. For maxima, returns values above threshold. If None, all local extrema are returned. TYPE: float | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	For 1D: (coordinates, log_density_values) where coordinates are real-valued positions on the grid. For 2D: (coordinates, log_density_values) where coordinates is shape (n, 2) with (x, y) real-valued positions.

Source code in spectre/compute/free_energy.py

def find_local_extrema(
    log_density: torch.Tensor,
    x_grid: torch.Tensor,
    y_grid: torch.Tensor | None = None,
    mode: Literal["minima", "maxima"] = "minima",
    min_distance: int = 1,
    threshold: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Find local extrema (minima or maxima).

    Parameters
    ----------
    mode : Literal["minima", "maxima"], optional, by default "minima"
        Whether to find local minima or maxima.

    x_grid : torch.Tensor
        Grid on which free energy is computed along x axis

    y_grid : torch.Tensor | None, optional, by default None
        Grid on which free energy is computed along y axis if free energy
        is computed from a 2D array. For a 1D array is None.

    min_distance : int, optional, by default 1
        Minimum distance between extrema. For 1D, this is in grid points.
        For 2D, defines the neighborhood size for local filtering.

    threshold : float | None, optional, by default None
        Energy threshold for filtering results. For minima, returns values
        below threshold. For maxima, returns values above threshold. If None,
        all local extrema are returned.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        For 1D: (coordinates, log_density_values) where coordinates are
        real-valued positions on the grid.
        For 2D: (coordinates, log_density_values) where coordinates is shape
        (n, 2) with (x, y) real-valued positions.
    """
    if mode not in ["minima", "maxima"]:
        raise ValueError(f"Mode must be 'minima' or 'maxima', got '{mode}'.")

    sign = -1 if mode == "minima" else 1

    if log_density.ndim == 1:
        # For 1D, find peaks in signed log_density
        peaks = _find_peaks_1d(sign * log_density, distance=min_distance)

        if threshold is not None:
            if mode == "minima":
                mask = log_density[peaks] <= threshold
            else:
                mask = log_density[peaks] >= threshold
            peaks = peaks[mask]

        coords = x_grid[peaks]
        values = log_density[peaks]

        return coords, values

    elif log_density.ndim == 2:
        # For 2D, use local extremum filter
        footprint_size = 2 * min_distance + 1
        if mode == "minima":
            local_extremum = _local_extremum_filter_2d(
                log_density, size=footprint_size, mode="min"
            )
        else:
            local_extremum = _local_extremum_filter_2d(
                log_density, size=footprint_size, mode="max"
            )

        extrema_mask = log_density == local_extremum

        # Exclude boundary artifacts
        extrema_mask[0, :] = False
        extrema_mask[-1, :] = False
        extrema_mask[:, 0] = False
        extrema_mask[:, -1] = False

        # Get coordinates of extrema
        extrema_coords = torch.nonzero(extrema_mask, as_tuple=False)

        if threshold is not None:
            extrema_energies = log_density[extrema_coords[:, 0], extrema_coords[:, 1]]
            if mode == "minima":
                mask = extrema_energies <= threshold
            else:
                mask = extrema_energies >= threshold
            extrema_coords = extrema_coords[mask]

        extrema_energies = log_density[extrema_coords[:, 0], extrema_coords[:, 1]]

        # Convert indices to real coordinates
        x_coords = x_grid[extrema_coords[:, 1]]
        y_coords = y_grid[extrema_coords[:, 0]]  # ty: ignore
        coords = torch.stack([x_coords, y_coords], dim=1)

        return coords, extrema_energies

    else:
        raise ValueError(
            f"Expected log_density to have 1 or 2 dimensions, got {log_density.ndim}."
        )