boltzkit.evaluation.sample_based.torsion_marginals

Functions

get_free_energy_difference(phis[, phi_range])	Estimate the free energy difference between two regions of backbone torsion phi angles.
get_torsion_angles(samples, topology)	Extract backbone φ and ψ torsion angles in the range (-pi, pi) from a batch of molecular coordinates.
get_torsion_marginal_hists(phis, psis, **kwargs)	Compute 1D and 2D histograms of torsion angles for all φ/ψ pairs.
get_torus_wasserstein_2(phis_psis_true, ...)	Compute the toroidal 2-Wasserstein distance between precomputed φ/ψ angles given in range (-pi, pi).
visualize_torsion_marginals_all(...)
visualize_torsion_marginals_dual(...)

boltzkit.evaluation.sample_based.torsion_marginals.get_torsion_angles(samples: ndarray, topology: Topology)

Extract backbone φ and ψ torsion angles in the range (-pi, pi) from a batch of molecular coordinates.

Parameters:

• samples (np.ndarray) – Cartesian coordinates of shape (batch, n_atoms, 3) or (batch, n_atoms*3), where batch is the number of frames and the last dimension corresponds to (x, y, z).

• topology (md.Topology) – MDTraj topology describing the molecular system.

Returns:

Tuple of arrays (phis, psis): - phis : np.ndarray of shape (batch, n_torsions) - psis : np.ndarray of shape (batch, n_torsions) Each entry contains the backbone φ or ψ angles in radians.

Return type:

tuple[np.ndarray, np.ndarray]

boltzkit.evaluation.sample_based.torsion_marginals.get_torsion_marginal_hists(phis: ndarray, psis: ndarray, **kwargs) → tuple[list[Histogram2D], list[Histogram1D], list[Histogram1D]]

Compute 1D and 2D histograms of torsion angles for all φ/ψ pairs.

Parameters:

• phis (np.ndarray) – Array of shape (batch, n_torsions) with φ angles in radians in range (-pi, pi).

• psis (np.ndarray) – Array of shape (batch, n_torsions) with ψ angles in radians in range (-pi, pi).

• **kwargs – Additional keyword arguments forwarded to get_histogram_1d and get_histogram_2d.

Returns:

A tuple of equal-size lists (one list element per torsion angle pair). - “phi_psi”: 2D Histogram2D of the joint φ/ψ angles - “phi”: 1D Histogram1D of φ angles - “psi”: 1D Histogram1D of ψ angles

Return type:

tuple[Histogram2D, Histogram1D, Histogram1D]

boltzkit.evaluation.sample_based.torsion_marginals.get_free_energy_difference(phis: ndarray, phi_range: tuple[float, float] = (0, 2))

Estimate the free energy difference between two regions of backbone torsion phi angles.

The function partitions the input values into an “inner” region defined by phi_range and an “outer” region (everything outside that interval), then computes a free energy difference based on the relative populations of these regions:

ΔF = - (log(N_inner) - log(N_outer)) = log(N_outer / N_inner)

Parameters:

• phis (np.ndarray) – Array of sampled values. Can be of shape (batch,) or (batch, n_torsions). All values are flattened implicitly when counting.

• phi_range (tuple of float, optional) – Tuple (phi_min, phi_max) defining the inclusive bounds of the “inner” region. Default is (0, 2).

Returns:

Estimated free energy difference ΔF in $kB T$ (i.e., no multiplication with kB*T) between the outer and inner regions. Returns np.nan if either region contains zero samples.

Return type:

float

boltzkit.evaluation.sample_based.torsion_marginals.get_torus_wasserstein_2(phis_psis_true: tuple[ndarray, ndarray], phis_psis_pred: tuple[ndarray, ndarray])

Compute the toroidal 2-Wasserstein distance between precomputed φ/ψ angles given in range (-pi, pi).

Parameters:

• phis_psis_true (tuple[np.ndarray, np.ndarray]) – Tuple (phis, psis) for the reference ensemble: - phis: (batch, n_torsions) - psis: (batch, n_torsions)

• phis_psis_pred (tuple[np.ndarray, np.ndarray]) – Tuple (phis, psis) for the predicted ensemble: - phis: (batch, n_torsions) - psis: (batch, n_torsions)

Returns:

Toroidal 2-Wasserstein distance between the joint φ/ψ distributions.

Return type:

float

Notes

• The joint angular distribution is treated as living on a torus, therefore respecting the periodicity.

boltzkit.evaluation.sample_based.torsion_marginals.visualize_torsion_marginals_all(torsion_marginals: tuple[list[~boltzkit.utils.histogram.Histogram2D], list[~boltzkit.utils.histogram.Histogram1D], list[~boltzkit.utils.histogram.Histogram1D]], vis_mode: ~boltzkit.utils.histogram.VisualizationMode = <function plot_as_log_density>, show: bool = False, **kwargs)

boltzkit.evaluation.sample_based.torsion_marginals.visualize_torsion_marginals_dual(torsion_marginals_true: tuple[list[~boltzkit.utils.histogram.Histogram2D], list[~boltzkit.utils.histogram.Histogram1D], list[~boltzkit.utils.histogram.Histogram1D]], torsion_marginals_pred: tuple[list[~boltzkit.utils.histogram.Histogram2D], list[~boltzkit.utils.histogram.Histogram1D], list[~boltzkit.utils.histogram.Histogram1D]], vis_mode: ~boltzkit.utils.histogram.VisualizationMode = <function plot_as_log_density>, show: bool = False, cmap: str | None = None, **kwargs)