API reference

Auto-generated from the source docstrings. These are the modules a user or referee is most likely to call directly; the full package has more.

Configuration

trails_md.config.TrailsMDConfig

Bases: BaseModel

MSM estimation & convergence

trails_md.msm.estimator

Markov State Model estimation over the Trails-MD CV / latent space.

MSMEstimator wraps :mod:deeptime.markov into a small, testable API that the adaptive loop can call once per iteration:

estimator = MSMEstimator(lagtime=10, n_microstates=100)
result = estimator.fit(trajs)          # trajs: list of (n_frames_i, n_cv)

It performs: clustering of the (continuous, per-walker) projections into microstates -> sliding-window transition counts -> restriction to the largest connected set -> maximum-likelihood (or Bayesian) MSM -> implied timescales, VAMP-2 score and PCCA+ metastable decomposition. The output is a serialisable :class:~trails_md.msm.diagnostics.MSMResult.

deeptime is imported lazily so that importing this module never hard-requires it; a clear error is raised only when estimation is actually attempted.

MSMEstimator

Estimate an MSM from a list of continuous CV trajectories.

Parameters:

Name	Type	Description	Default
lagtime	int	Lag time (in saved frames) used for the production MSM.	10
n_microstates	int	Number of clusters (microstates) used to discretise the CV space.	100
cluster_method	str	"kmeans" or "regspace" (regular-space clustering).	'kmeans'
estimator	str	"mle" (maximum likelihood) or "bayesian" (adds posterior error bars on the slow timescales via :class:deeptime.markov.msm.BayesianMSM).	'mle'
n_metastable	int | None	If set, run PCCA+ to coarse-grain into this many metastable states.	None
n_timescales	int	Number of slow processes (implied timescales) to track.	3
lagtimes	Sequence[int] | None	Optional lag-time ladder for an implied-timescale sweep; when provided, :meth:fit attaches an :class:ITSResult to the output.	None
n_bayesian_samples	int	Posterior sample count for the Bayesian estimator.	50
seed	int	Random seed forwarded to the clustering for reproducibility.	42

cluster(trajs)

Discretise CV trajectories into microstate index trajectories.

Returns (dtrajs, cluster_model) where dtrajs is a list of int arrays (one per input trajectory) and cluster_model exposes cluster_centers.

fit(trajs, iteration=None)

Cluster, estimate the MSM and return a serialisable result.

implied_timescales(dtrajs, lagtimes)

Estimate implied timescales across a ladder of lag times.

MSMEstimatorFactory

Registry for MSM estimator variants, mirroring SpawnerFactory/EngineFactory.

trails_md.msm.convergence

MSM-based convergence detection for the adaptive sampling loop.

A :class:ConvergenceMonitor holds a list of pluggable :class:ConvergenceCriterion objects. Each iteration it is fed the latest :class:~trails_md.msm.diagnostics.MSMResult; it records per-criterion state and reports convergence once the configured combination of criteria ("all" / "any") has held for patience consecutive iterations.

The criteria here operate on quantities that are invariant to microstate relabelling (slow implied timescales, VAMP-2 score, sorted metastable populations), so they remain valid even though the clustering changes from one iteration to the next.

ConvergenceCriterion

Bases: ABC

Base class for a single convergence test fed one MSMResult per call.

update(result) abstractmethod

Record result and report whether the test is currently satisfied.

reset()

Clear accumulated history (used when restarting a monitor).

ImpliedTimescaleCriterion

Bases: ConvergenceCriterion

Satisfied when the slowest k implied timescales stop changing.

Compares the current slow timescales against the previous iteration; the test passes when the maximum relative change is below tol.

VAMP2Criterion

Bases: ConvergenceCriterion

Satisfied when the VAMP-2 score plateaus (relative change below tol).

StationaryDistributionCriterion

Bases: ConvergenceCriterion

Satisfied when metastable populations stabilise across iterations.

Uses the PCCA+ metastable populations (sorted, so the test is invariant to macrostate relabelling) and measures L1 drift between consecutive iterations. Falls back to inactivity when no metastable decomposition is available.

StatisticalErrorCriterion

Bases: ConvergenceCriterion

Satisfied when the Bayesian relative error on the slowest timescale is low.

Requires estimator: bayesian so that timescale_errors is populated.

TransitionMatrixCriterion

Bases: ConvergenceCriterion

Satisfied when every significant transition probability is well determined.

Uses the connected count matrix to get an analytic Dirichlet uncertainty on each transition probability, σ(T_ij) = sqrt(T_ij(1−T_ij)/(c_i+1)) (no bootstrap). The test passes when the largest flux-weighted relative uncertainty over significant transitions is below tol: max_{(i,j): π_i T_ij > min_flux·max} σ(T_ij)/T_ij < tol. Flux weighting avoids being dominated by noise in tiny, irrelevant entries. This is a within-iteration absolute statistical-convergence test; combine it with the spectral criteria under mode: all to require both kinetic resolution and statistical convergence of the microstate transition matrix.

Requires count_matrix on the result (always populated by MSMEstimator).

ConvergenceMonitor

Aggregate pluggable criteria into a single converged / not-converged signal.

Parameters:

Name	Type	Description	Default
criteria	list[ConvergenceCriterion]	List of :class:ConvergenceCriterion instances.	required
mode	str	"all" (default) requires every criterion satisfied simultaneously; "any" requires at least one.	'all'
patience	int	Number of consecutive iterations the combination must hold before the monitor reports convergence.	2

update(result)

Feed a new MSMResult; return whether convergence is now declared.

Collective-variable feature selection

trails_md.spaces.feature_selection

VAMP-2 based input-feature selection and optimisation.

The quality of an MSM/CV is bounded by the input features fed to it. VAMP-2 is a variational score for how well a feature set captures the slow dynamics: higher is better, and it can be compared across different feature sets on the same trajectories (Wu & Noé, 2017; Scherer et al., 2019).

This module provides:

• :func:vamp2_score — a dependency-light VAMP-2 score from time-lagged covariances of feature trajectories.
• :func:rank_candidates — rank named candidate feature sets by VAMP-2.
• :func:greedy_vamp_selection — greedy forward selection of the feature columns (or column groups) that maximise VAMP-2, i.e. an optimisation protocol that picks the best subset of features.
• :class:FeatureSelector — thin orchestrator used by the adaptive loop to choose and periodically update the input features.

All functions operate on plain arrays, so they are testable without MD inputs.

FeatureSelection dataclass

Outcome of a feature-selection step (serialisable for checkpoints).

FeatureSelector

Choose the best input-feature columns by VAMP-2 optimisation.

Used by the adaptive loop when feature_selection.enabled is set. Operates on a feature matrix reshaped into per-walker trajectories.

vamp2_score(trajs, lagtime, dim=None, epsilon=1e-06)

VAMP-2 score of feature trajs at lagtime.

Defined as the sum of squared singular values of the whitened time-lagged correlation (Koopman) matrix C00^{-1/2} C0t C11^{-1/2} on mean-free features. Larger means the features resolve more, slower kinetic variance. dim optionally caps the number of singular values retained.

rank_candidates(candidates, lagtime, dim=None)

Rank named candidate feature sets by VAMP-2 (best first).

greedy_vamp_selection(trajs, lagtime, groups=None, max_groups=None, dim=None, min_gain=0.0001)

Greedy forward selection of feature columns maximising VAMP-2.

Starting from an empty set, repeatedly add the column group whose inclusion most increases the VAMP-2 score, stopping when no group improves the score by more than min_gain (or max_groups are selected). Returns the sorted list of selected column indices.

Adaptive binning

trails_md.binning.adaptive

Landscape-adaptive binning schemes.

The default :class:~trails_md.binning.spatial.RegularBinner is a uniform grid: constant bin width everywhere. Near a steep free-energy barrier a wide bin lets a walker slide back before it can reach the next bin within the lag time, so the WE flux across the barrier stalls; in flat basins fine bins waste replicas. These schemes make the bins landscape-adaptive — finer where the landscape is steep / sparse, coarser where it is flat — and are recomputed every iteration.

All binners share the :class:~trails_md.binning.spatial.RegularBinner API (fit(points) -> BinTable), so the density / WE spawners consume them interchangeably. uniform maps straight to RegularBinner for exact backwards compatibility.

Schemes

• uniform : constant-width grid (RegularBinner).
• gradient : equi-resistance edges — boundaries at equal increments of ∫ exp(βF) dx ∝ ∫ 1/P dx so bins concentrate where the sampled density is low (barriers / steep regions).
• mab : Minimal-Adaptive-Binning-style — uniform bins between the occupied extremes plus narrow "foothold" bins at the moving fronts.
• eigenvector : bin uniformly along the leading (slowest) CV coordinate only; for a learned CV / committor proxy this is automatically fine across the barrier and coarse in basins.

AdaptiveBinner

Bases: ABC

Base class: per-axis adaptive edges over a (padded) bounding box.

GradientBinner

Bases: AdaptiveBinner

Equi-resistance edges: dense where the sampled density is low (barriers).

Boundaries are placed at equal increments of the "resistance" ∫ 1/P(x) dx ∝ ∫ exp(βF(x)) dx, so bins bunch up where the sampled density is low — i.e. across barriers — and spread out in well-sampled basins.

Two safeguards matter in practice and were added after benchmarking:

• Edges are laid out across the occupied range only. Over the full configured domain the unvisited region has zero density, so 1/P diverges there and swallows the entire edge budget: the occupied region then collapses into one or two enormous bins and the frontier is diluted rather than resolved — the exact opposite of the intent.
• The resistance is clipped to max_resistance times its median, so a single near-empty slice cannot monopolise the edges either.

MABinner

Bases: AdaptiveBinner

Minimal Adaptive Binning (Torrillo, Bogetti & Chong, J. Phys. Chem. A 2021).

Three ingredients, recomputed every iteration:

1. Dedicated boundary bins. The single leading (front-most) and trailing (rear-most) frames each get their own narrow bin. A bin holding one frame has the maximum possible density weight 1/n_b, so the frontier is always eligible for respawning and can never be diluted into a populated neighbour.
2. Dedicated bottleneck bins. Bottlenecks are detected with the MAB objective Z_i = log(n_i) - log(sum of n_j ahead of i): a slice that still holds population while almost nothing lies beyond it is the uphill face of a barrier. n_bottleneck such slices (per direction) get their own narrow bins, which concentrates respawning exactly where flux is being lost.
3. Evenly spaced bins in between, spanning only the occupied range rather than the full configured domain, so resolution follows the walkers.

Note that fine bins at the frontier are necessary but not sufficient: with hard density spawning the frontier bin is only one of walker selected bins, so only ~1/walker of the effort attacks the barrier. Pair this binner with the WE spawner (spawn_scheme: we, which replicates we_target_per_bin walkers into each occupied bin) to obtain a genuine ratchet.

EigenvectorBinner

Bases: AdaptiveBinner

Bin uniformly along the leading (slowest) CV coordinate only.

For a learned CV / committor-proxy the slow coordinate compresses basins and stretches the barrier, so uniform bins in it are automatically fine across the barrier and coarse in basins. coordinate selects the column (default 0).

make_binner(scheme, *, n_bins, min_values=None, max_values=None, target=None, n_fine=100, smoothing=3, n_bottleneck=1)

Construct the configured binner. uniform → RegularBinner (exact).

Weighted ensemble & kinetics

trails_md.spawners.we

Weighted-ensemble spawner (Huber & Kim, 1996).

Resamples the live walker ensemble by split/merge, keeping we_target_per_bin walkers in every occupied bin while conserving total statistical weight. Unlike the exploration-oriented spawners (density / FPS / LOF / MSM least-counts), the weights it carries are rigorous: a walker split c ways yields c children of weight w/c, merges sum weight, and sum(weights) == 1 every iteration. That is what makes an unbiased rate (MFPT) recoverable from a WE run, and it is the whole reason to prefer spawn_scheme: we over the cheaper exploration schemes.

Implements the standard sample(points, top_n, history) contract, returning indices into the cumulative point cloud (repeats indicate split walkers), so it is a drop-in spawn_scheme: we option. Walker weights are carried across iterations in the spawner instance (and exposed via state_dict).

Two invariants are easy to get backwards and both are load-bearing; see sample and _resample_to_budget for why:

• CPU is allocated across bins, never in proportion to weight.
• The ensemble is the live walkers, never an arbitrary frame from history.

WESpawner

Bases: Spawner

sample(points, top_n, history=None)

One Huber-Kim weighted-ensemble resampling step.

Two things make this rigorous WE rather than "adaptive sampling that also tracks some numbers":

The ensemble is the current walkers, not the whole history. WE is a Markov resampling of the live ensemble: each walker runs for tau, and the endpoints are then split/merged. Restarting from an arbitrary frame drawn out of the cumulative cloud has no well-defined statistical weight -- the frame's weight was already spent in the iteration that produced it -- and reusing it silently double-counts probability. A rate computed from such weights is wrong, which is worse than having no rate. So the ensemble here is exactly the endpoint of each current walker.

The walker budget is the resampling target. Slots are allocated across occupied bins first (equal share; ties to the sparsest bin, which is the frontier), and each bin is then split/merged to exactly its slot count. Every resampled walker is therefore actually run, so total weight is conserved: sum(w) == 1 every iteration, and a child of a walker split c ways carries w/c. That is what makes an unbiased MFPT recoverable downstream.

mfpt(tau_ps, discard_fraction=0.5)

Steady-state MFPT (ns) from the recycled-flux series, via the Hill relation.

The early iterations are a transient: the ensemble has not yet reached the non-equilibrium steady state, and the flux during that ramp-up systematically UNDERESTIMATES the rate. So the leading discard_fraction of the series is dropped before averaging -- reporting the un-discarded average is the single most common way a WE rate is wrong. Returns None if nothing has been recycled yet (no flux -> no rate, rather than an infinite one).

MFPTResult dataclass

Steady-state MFPT (Hill relation) plus the diagnostics needed to trust it.

converged property

Heuristic: the retained flux has plateaued (within 20%) and has events.

steady_state_mfpt(flux_history, tau_ps, discard_fraction=0.5)

MFPT (ns) from a recycled-flux series via the Hill relation MFPT = tau / flux.

Drops the leading discard_fraction of the series (the pre-steady-state transient systematically underestimates the rate) and averages the remaining tail. Also reports a plateau ratio (second half of the tail vs first half) so a caller can tell a converged steady state from a still-drifting one. mfpt_ns is None when nothing has been recycled yet -- no flux means no rate, not an infinite one.

trails_md.binning.we

Weighted-ensemble (WE) resampling.

Implements the split/merge resampling of Huber & Kim (1996): walkers carry statistical weights and are kept at a target count per bin while total weight is conserved. Under-represented bins gain walkers by splitting high-weight walkers (weight divided among copies); over-represented bins lose walkers by merging low-weight walkers (weights summed, one survivor chosen with probability proportional to weight). This focuses sampling on bins/regions without biasing the estimated probabilities.

The core operates on plain arrays (weights + bin labels), so it is independent of the binning implementation and fully unit-testable.

ResampleResult dataclass

Outcome of one WE resampling step.

parents are indices into the input ensemble (a value may repeat when a walker was split); weights are the matching statistical weights. Total weight equals the input total (up to floating point).

WeightedEnsemble

Split/merge resampler that conserves probability weight.

Parameters:

Name	Type	Description	Default
target_per_bin	int	Desired number of walkers in each occupied bin after resampling.	4

resample(weights, bin_labels, target_per_bin=None, rng=None)

Resample walkers to target_per_bin per occupied bin.

weights[i] and bin_labels[i] describe walker i. Returns the post-resampling ensemble as parent indices + weights.

Analysis & plotting

trails_md.analysis.data

Data utilities for post-hoc MSM analysis.

Pure NumPy helpers (no matplotlib) that load the per-iteration msm.npz / cvs.npz files written during a run and derive quantities for plotting: convergence series, free energies, and free-energy surfaces. Kept separate from :mod:trails_md.analysis.plots so the numerics are testable without a plotting backend.

load_msm_series(run_dir)

Collect per-iteration MSM scalars into aligned arrays.

Returns a dict with iterations, vamp2 and timescales (shape (n_iters, max_processes), NaN-padded). Iterations without an msm.npz are skipped.

load_latest_msm(run_dir)

Return the arrays of the most recent msm.npz (or None if absent).

load_run_meta(run_dir)

Parse the # key=value header of output.log into a dict.

Numeric values (step, dt, walker, ...) are returned as floats. Used to recover tau_ps = step * dt for the weighted-ensemble MFPT without re-loading the config. Returns {} if the log is absent.

load_flux_history(run_dir)

Recycled-flux series from the latest checkpoint's spawner state.

Empty unless the run used weighted ensemble with source->sink recycling (kinetics mode). Reads the most recent readable checkpoints/iter_*/ sampler_state.pkl and returns its spawner flux_history.

load_cv_points(run_dir)

Stack all per-iteration CV projections (cvs.npz) into one array.

free_energy_from_populations(populations, temperature=300.0)

Relative free energy -kT ln(p) (kJ/mol), shifted so the min is 0.

free_energy_surface(points, bins=60, temperature=300.0)

2D free-energy surface F(x,y) = -kT ln P(x,y) from CV points.

Returns (F, xedges, yedges) with F shifted to a 0 minimum and unsampled cells set to NaN. Requires at least 2D points.

trails_md.analysis.plots

Matplotlib plotting utilities for MSM analysis.

Each function accepts an optional Axes and returns it, so plots compose into custom figures; :func:plot_convergence_report assembles a standard multi-panel summary for a run. matplotlib is an optional dependency (trails-md[examples]); it is imported lazily with an actionable error if missing.

plot_implied_timescales(lagtimes, timescales, ax=None)

Implied timescales vs lag time (the ITS convergence plot).

plot_timescale_convergence(series, ax=None)

Slowest implied timescales vs iteration.

plot_vamp2_convergence(series, ax=None)

VAMP-2 score vs iteration.

plot_free_energy_surface(points, bins=60, temperature=300.0, ax=None)

Free-energy surface over the first two CV dimensions.

plot_metastable_free_energy(populations, temperature=300.0, ax=None)

Bar chart of metastable-state free energies (from PCCA+ populations).

plot_msm_network(transition_matrix, stationary=None, ax=None, threshold=0.01)

Draw the MSM as a network: node size ~ stationary weight, edges ~ T_ij.

Uses a circular layout (no networkx dependency).

plot_convergence_report(run_dir, outfile=None, temperature=300.0)

Assemble a standard multi-panel summary figure for a run.

Panels: VAMP-2 convergence, timescale convergence, free-energy surface, and (when available) the latest implied-timescale sweep / MSM network. Saves to outfile (default <run_dir>/analysis/convergence_report.png) and returns the saved path.

plot_flux_convergence(flux_history, tau_ps, discard_fraction=0.5, ax=None)

Weighted-ensemble kinetics: recycled flux and the running MFPT vs iteration.

Two overlaid series on twin axes: the per-iteration recycled flux (left) and the cumulative MFPT estimate over the retained tail (right). The shaded region is the discarded pre-steady-state transient. A flat flux tail = a trustworthy rate.

save_flux_convergence(flux_history, tau_ps, outfile, discard_fraction=0.5)

Render :func:plot_flux_convergence to outfile and return its path.