Trails-MD is a lightweight, engine-agnostic adaptive sampling framework that iteratively launches short trajectory ensembles across fixed or on-the-fly machine-learned collective-variable spaces. The MD engine, the sampling space, the spawning rule and the execution backend are all interchangeable behind a single configuration file.
Its defining feature is explicit parent-child trajectory lineage preservation, which allows continuous transition pathways to be reconstructed from highly parallelised, disjointed exploration stages, and cleanly separates basin discovery from genuine pathway connectivity. Two operating modes are kept distinct: an exploration mode for rapid conformational coverage, and a weighted-ensemble kinetics mode that yields an unbiased mean first passage time from the steady-state flux.
", + "upload_type": "software", + "license": "other-nc", + "access_right": "open", + "creators": [ + { + "name": "Maity, Dibyendu", + "affiliation": "S. N. Bose National Centre for Basic Sciences, Kolkata, India" + }, + { + "name": "Majumdar, Rupak", + "affiliation": "Max Planck Institute for Software Systems, Kaiserslautern, Germany" + }, + { + "name": "Chakrabarty, Suman", + "affiliation": "S. N. Bose National Centre for Basic Sciences, Kolkata, India" + } + ], + "keywords": [ + "molecular dynamics", + "adaptive sampling", + "enhanced sampling", + "collective variables", + "lineage-aware sampling", + "Markov state model", + "weighted ensemble", + "mean first passage time" + ], + "related_identifiers": [ + { + "identifier": "https://github.com/TeamSuman/Trails-MD", + "relation": "isSupplementTo", + "scheme": "url" + } + ] +} diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a4d077..d71920e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,62 @@ All notable changes to Trails-MD are documented here. The format is based on [Keep a Changelog](https://keepachangelog.com/), and the project aims to follow [Semantic Versioning](https://semver.org/). +## [1.1.0] — 2026-08-02 + +Adds two features that were previously implicit or absent, and one new analysis +module. No behaviour changes for existing configuration files: every new setting +defaults to what the code already did. + +### Added + +- **`spawning.history_window`** — controls how much of the campaign a spawner scores + against. The default (`null`) is unchanged: the whole history. Setting it to a small + integer gives PaCS-MD-style cycle-local selection, and `0` scores only the current + iteration. This exists because a *coverage* objective is only definable over + cumulative history, and the setting makes that claim testable under an otherwise + identical loop. See `docs/configuration.md`. + +- **`adaptive_model.tvae_beta`** — the KL weight in the TVAE objective, + `loss = mse + beta * kld / n_features`. Previously the estimator was constructed + without this argument, so every run silently used deeptime's default of `1.0` and no + user could change it; answering "which beta did you use?" required reading a + dependency's source. The default remains `1.0`, so published results are unaffected, + and checkpoints written before this release restore to `1.0` rather than to a new + default. `docs/cv_methods.md` now writes the loss out in full. + +- **`trails_md.analysis.riteweight`** — randomized iterative trajectory reweighting + (Kania *et al.*, PNAS **123**, e2529246123, 2026) for recovering a stationary + distribution from adaptively-sampled data **without a lag time and without assuming + cluster-level Markovianity**, complementing the existing MSM route. Independent + implementation from the published algorithm; the authors' reference code carries no + licence statement and was deliberately not copied. New page: `docs/reweighting.md`. + +### Fixed + +- **The shipped template disagreed with the code defaults.** `templates.py` advertised + `encoder_hidden_dims: [64, 32]` / `decoder_hidden_dims: [32, 64]` / + `deep_tica_hidden_dims: [64, 32]` while the built-in defaults are `[256, 128]` / + `[128, 256]` / `[256, 128]`, so a user who copied the template trained a different + network from one who omitted the block. The template now states the real defaults. + +- `trails_md.analysis` did not export anything but `data`, so `riteweight` was + importable only by full module path. + +### Documentation + +- `docs/cv_methods.md`: the TVAE loss in full, `beta` and its per-feature + normalisation, that `lagtime` is counted in **frames** (physical lag = + `lagtime × stride × dt`), and an explicit statement that time-lagged pairs are built + per walker and therefore never span a respawn. +- `docs/configuration.md`: `history_window`, `tvae_beta`, `dropout_rate`, decoder + widths and the SPIB hyperparameters; plus a section on cumulative vs. cycle-local + selection and why the comparison must be made against aggregate simulation time + rather than wall-clock time. +- `docs/reweighting.md` (new): when to use MSM reweighting and when RiteWeight, how to + build segment pairs without crossing a respawn, and the limitation both share — + neither can fix mis-*coverage*. +- `.zenodo.json` added so the archived release carries proper metadata. + ## [1.0.0] — 2026-07-23 **First public release.** It brings a coverage-driven adaptive sampler up to an diff --git a/CITATION.cff b/CITATION.cff index 9ce60cf..b883389 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -26,7 +26,8 @@ authors: email: sumanc@bose.res.in affiliation: "S. N. Bose National Centre for Basic Sciences, Kolkata, India" # orcid: "https://orcid.org/0000-0000-0000-0000" -version: 1.0.0 +date-released: 2026-08-02 +version: 1.1.0 license: PolyForm-Noncommercial-1.0.0 repository-code: "https://github.com/TeamSuman/Trails-MD" url: "https://github.com/TeamSuman/Trails-MD" diff --git a/docs/configuration.md b/docs/configuration.md index 21fbfc1..16ee9c1 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -53,6 +53,7 @@ at startup. Below, only non-obvious defaults are noted — see | `target` | — | CV-space target `[x, y, …]` when `search_mode: target`. | | `recent_density_window` | `5` | Bins sampled in the last N iterations are down-weighted (`density`). | | `lof_neighbors` | `20` | Neighbours for the LOF spawner. | +| `history_window` | `None` | How many past iterations the spawner scores over. `None` (default) uses the entire campaign; an integer `k` keeps only the last `k` iterations, and `0` scores the current iteration alone. See [cumulative vs. cycle-local selection](#cumulative-vs-cycle-local-selection). | | `voronoi_periodic` | `false` | Wrap Voronoi cells periodically. | | `voronoi_grid_size` | `250` | Grid resolution for Voronoi cell-area estimation. | | `voronoi_max_clusters` | `5000` | Upper bound on auto-grown Voronoi cells. | @@ -69,6 +70,32 @@ at startup. Below, only non-obvious defaults are noted — see limit and marks the batch failed — a guard against a hung in-process OpenMM walker. Off by default. +### Cumulative vs. cycle-local selection + +By default a spawner ranks candidate frames against **every frame the campaign has +ever produced**. This is deliberate, and it is the main structural difference from +the PaCS-MD family, whose selection rules rank only the frames of the current cycle. + +The distinction matters because a *coverage* objective is only definable over +cumulative history: "which regions are under-sampled?" is meaningless if the only +frames in view are the ones just generated. Cycle-local schemes therefore use +frontier or extremum objectives (furthest from the start, closest to a target, most +outlying) rather than coverage. + +`history_window` exposes the choice so the two can be compared under an otherwise +identical loop: + +```yaml +spawning: + history_window: null # default: score against the whole campaign + # history_window: 0 # score only the current iteration (cycle-local) + # history_window: 5 # sliding window of the last five iterations +``` + +Cycle-local selection is slightly cheaper per iteration because the candidate pool is +smaller, so comparisons between settings should be made against **aggregate +simulation time, not wall-clock time**. + ## `space_mode` and adaptive model `space_mode`: `fixed` \| `pca` \| `tica` \| `tvae` \| `deep-tica` (+ experimental @@ -87,7 +114,10 @@ learned modes: | `adaptive_model.lagtime` | `5` | Lag time for time-lagged CVs. | | `adaptive_model.latent_dim` | `2` | CV dimensionality. | | `adaptive_model.epochs` / `learning_rate` | `50` / `5e-4` | Training. | -| `adaptive_model.encoder_hidden_dims` | `[256,128]` | Network width. | +| `adaptive_model.encoder_hidden_dims` / `decoder_hidden_dims` | `[256,128]` / `[128,256]` | Network width. | +| `adaptive_model.dropout_rate` | `0.1` | Dropout in the encoder/decoder. | +| `adaptive_model.tvae_beta` | `1.0` | **TVAE only.** Weight of the KL term in `loss = mse + beta * kld / n_features`. `1.0` is the standard VAE objective and the value all published Trails-MD results used; `< 1` favours reconstruction, `> 1` favours a smoother latent space. | +| `adaptive_model.spib_n_states` / `spib_beta` | `10` / `1e-3` | **SPIB only.** | ## `execution` diff --git a/docs/cv_methods.md b/docs/cv_methods.md index 0ca8aec..adb14c9 100644 --- a/docs/cv_methods.md +++ b/docs/cv_methods.md @@ -60,16 +60,52 @@ space_mode: tica adaptive_feature_type: distances # distances | fitted_coords | phi_psi retrain_freq: 5 # retrain the CV every 5 iterations adaptive_model: - lagtime: 5 + lagtime: 5 # in FRAMES: lag = lagtime x stride x dt latent_dim: 2 epochs: 50 - encoder_hidden_dims: [64, 32] + encoder_hidden_dims: [256, 128] # also the built-in default ``` +`lagtime` is counted in **saved frames**, not in time units, so the physical lag is +`lagtime × stride × dt`. With `stride: 100` and `dt: 0.002` ps, `lagtime: 5` is a 1 ps +lag. + When a model is retrained, the full accumulated feature history is reprojected into the updated latent space before spawning, so selection always reflects the current coordinates. +### Time-lagged pairs never cross a respawn + +For every time-lagged method (TICA, TVAE, Deep-TICA, VAMPNets, SPIB) the features are +split per walker before the lagged dataset is built, so a lagged pair is always drawn +from one continuous walker segment. Velocities are redrawn at each respawn, so a pair +spanning that boundary would relate dynamically unrelated configurations. This also +means the lag must be shorter than a walker segment: `lagtime < step / stride`. + +Adaptive sampling still biases *which* segments exist, so the estimator sees a +non-Boltzmann mixture of short trajectories. That affects the estimated timescales, +not the validity of the pairs; see [Reweighting](reweighting.md). + +### The TVAE objective + +The time-lagged variational autoencoder encodes frame $\mathbf{x}_t$ and reconstructs +$\mathbf{x}_{t+\tau}$, so the latent space is trained to retain slowly-decorrelating +information rather than high-variance information. Its loss is + +$$ +\mathcal{L} = \underbrace{\lVert \mathbf{x}_{t+\tau} - \hat{\mathbf{x}}_{t+\tau} \rVert^2}_{\text{reconstruction}} +\; + \; \frac{\beta}{n_\text{features}} \, D_\mathrm{KL}\!\left(q(\mathbf{z}\mid\mathbf{x}_t) \,\Vert\, p(\mathbf{z})\right) +$$ + +with $p(\mathbf{z}) = \mathcal{N}(0, I)$. Note the KL term is divided by the number of +input features, so $\beta$ is defined **per feature**; the same $\beta$ therefore means +the same thing across systems of different dimensionality. + +`adaptive_model.tvae_beta` sets $\beta$ and defaults to `1.0`, the standard VAE +objective and the value used for every published Trails-MD result. Lowering it favours +reconstruction accuracy; raising it favours a smoother, more strongly regularised latent +space at the cost of reconstruction. + ## Availability checks If a method's backend is missing, Trails-MD raises an actionable error, e.g.: diff --git a/docs/reweighting.md b/docs/reweighting.md new file mode 100644 index 0000000..fd761db --- /dev/null +++ b/docs/reweighting.md @@ -0,0 +1,101 @@ +# Reweighting an adaptive ensemble + +Adaptive sampling deliberately over-samples sparse regions. The configurations an +exploration campaign produces are therefore **not** Boltzmann distributed, and a +histogram of them is not a free-energy surface. Something has to put the weights back. + +Trails-MD offers two routes, and they fail in different ways — which is the reason for +keeping both. + +| | MSM reweighting | RiteWeight | +| --- | --- | --- | +| Assumption | Markovian at the chosen lag, on the chosen discretisation | none of either | +| Needs a lag time | **yes** — and the answer depends on it | no | +| Discretisation error | enters through the clustering | averaged away by re-randomising it | +| Cost | one clustering + one eigenproblem | one clustering + one eigenproblem **per iteration** | + +Neither fixes mis-*coverage*. A basin that was never visited cannot be reweighted into +existence, and no diagnostic in either method will tell you one is missing. + +## MSM reweighting + +The standard route: discretise, count transitions at lag $\tau$, solve for the +stationary distribution $\pi$, and give each frame weight $\pi_I / N_I$ where $I$ is its +state and $N_I$ that state's frame count. See [MSM & kinetic seeding](msm.md). + +The catch is the lag. Too short and the cluster-level dynamics are not Markovian, so +$\pi$ is wrong; too long and the transition counts run out. Trails-MD reports an +implied-timescale diagnostic precisely so this can be checked rather than assumed. + +## RiteWeight + +`trails_md.analysis.riteweight` implements the randomized iterative trajectory +reweighting scheme of Kania, Webber, Simpson, Aristoff and Zuckerman, +*PNAS* **123**, e2529246123 (2026). Short trajectory segments are reweighted +iteratively until their weighted distribution is self-consistent with the stationary +distribution of a transition matrix built from those same weights. + +The key device is that **the clustering is re-randomised at every iteration**. Segments +that shared a cluster in one iteration are separated in the next, so the fixed point is +governed by the underlying microstate dynamics rather than by any particular +discretisation. In practice the answer should not depend on the cluster count — a +property worth checking, and one the test suite asserts. + +```python +import numpy as np +from trails_md.analysis import riteweight + +# One (start, end) pair per short segment. Features must be invariant to rotation +# and translation: pairwise distances, torsion sin/cos, a learned projection, ... +result = riteweight(start_features, end_features, n_clusters=150, + n_iterations=2000, average_last=400, seed=0) + +print(result.converged) # heuristic: late drift << early drift, and small +weights = result.weights # one weight per segment, summing to 1 + +H, xe, ye = np.histogram2d(cv[:, 0], cv[:, 1], bins=60, weights=weights) +F = -kT * np.log(H / H.sum()) +``` + +### Forming the segment pairs + +Both configurations of a pair must come from the **same continuous stretch of unbiased +dynamics**. Trails-MD applies no biasing forces, but exploration mode redraws velocities +at every respawn, so a pair must never straddle a respawn boundary — the same constraint +that governs the time-lagged CV estimators: + +```python +starts, ends = [], [] +for iteration in iterations: + cv = np.load(iteration / "cvs.npz")["cvs"] + per_walker = len(cv) // n_walkers + for w in range(n_walkers): + seg = cv[w * per_walker:(w + 1) * per_walker] + starts.append(seg[:-lag]) + ends.append(seg[lag:]) +``` + +`lag` here only sets which pairs are formed; unlike an MSM, the fixed point is not a +function of it. Checking that the answer is flat across several lags is a cheap and +informative diagnostic. + +### Convergence + +`RiteWeightResult.weight_drift` records the per-iteration $L_1$ change in the weight +vector. A settled run shows it falling and then fluctuating about a small value. +`result.converged` applies a heuristic to that trace; it is reported rather than +enforced, so you can judge convergence instead of assuming it. + +### Cost + +Each iteration re-clusters every pooled point, so cost grows as +`n_iterations × n_points`. The nearest-centre assignment uses a KD-tree over the +(few) cluster centres, which is exact and turns an otherwise prohibitive +$O(N k)$ scan into roughly $O(N \log k)$ — the difference between feasible and not for +a campaign with $10^6$ pooled frames. + +!!! note "Attribution" + This is an independent implementation written from the published algorithm. The + authors' reference code (`github.com/ZuckermanLab/rite_weight`) carries no licence + statement and was deliberately not copied or vendored. If you use this feature, + cite the PNAS paper. diff --git a/examples/template.yaml b/examples/template.yaml index 15e377e..e355d9b 100644 --- a/examples/template.yaml +++ b/examples/template.yaml @@ -50,6 +50,11 @@ spawning: voronoi_clusters: 150 # cells / microstates (voronoi & msm spawners) we_target_per_bin: 4 # walkers per bin for spawn_scheme: we lof_neighbors: 20 + # history_window: null # iterations of history the spawner scores over. + # null (default) = the whole campaign, which is what + # makes a coverage objective definable at all. Set to + # a small integer for PaCS-MD-style cycle-local + # selection; 0 scores only the current iteration. # ---- Kinetics mode (rate / MFPT): spawn_scheme: we + md_engine: openmm -------- # inherit_velocities: true # continue parent velocities (required for a rate) # recycle_target: [[-2.5, -1.0], [2.0, 3.0]] # source->sink sink box, one [lo, hi] per CV dim @@ -74,14 +79,17 @@ aggregate_memory: true max_adaptive_memory_frames: 50000 adaptive_model: # hyperparameters for learned CVs - lagtime: 5 + lagtime: 5 # in FRAMES, so lag time = lagtime x stride x dt latent_dim: 2 epochs: 50 learning_rate: 0.0005 - encoder_hidden_dims: [64, 32] - decoder_hidden_dims: [32, 64] + # These are also the built-in defaults, so a config that omits the whole + # adaptive_model block trains exactly the network described here. + encoder_hidden_dims: [256, 128] + decoder_hidden_dims: [128, 256] dropout_rate: 0.1 - deep_tica_hidden_dims: [64, 32] + deep_tica_hidden_dims: [256, 128] + tvae_beta: 1.0 # TVAE only: KL weight in mse + beta*kld/n_features spib_n_states: 10 # SPIB only spib_beta: 0.001 # SPIB only diff --git a/mkdocs.yml b/mkdocs.yml index 7d5d2a9..1f8b15a 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -40,7 +40,9 @@ nav: - Execution (workstation & HPC): execution.md - HPC scaling & WESTPA comparison: hpc_scaling.md - Performance & GPU utilization: performance.md - - Analysis: analysis.md + - Analysis: + - MSM analysis & plotting: analysis.md + - Reweighting an adaptive ensemble: reweighting.md - Reference: - CLI reference: cli.md - Python API: api.md diff --git a/tests/test_history_window.py b/tests/test_history_window.py new file mode 100644 index 0000000..8c11934 --- /dev/null +++ b/tests/test_history_window.py @@ -0,0 +1,181 @@ +"""History-window (cumulative vs cycle-local) selection. + +PaCS-MD ranks and reseeds from the snapshots of the *current cycle only* +(PaCS-Toolkit: ``self.CVs = results[::skip_frame]`` over one cycle's replicas), +whereas TRAILS-MD spawns from the entire accumulated history. To measure what that +single algorithmic difference is worth, the candidate pool must be restrictable to +the most recent N iterations: + + history_window = None -> full cumulative history (default; current behaviour) + history_window = 0 -> current iteration only (PaCS-MD-like, cycle-local) + history_window = N -> the N most recent historical iterations + +The pool is also index-synchronised with the trajectory and frame-record lists built +in core.py, so the window must be applied in the one shared helper rather than at +each call site -- otherwise a spawn index would select a different conformation than +the one whose lineage is recorded. +""" + +import numpy as np + +from trails_md.spawners.history import pooled_history_iterations + + +def _history(n_iterations: int, dim: int = 2, frames: int = 4) -> dict: + return { + i: {"projection": np.zeros((frames, dim), dtype=float)} + for i in range(n_iterations) + } + + +def test_window_none_keeps_full_cumulative_history(): + """Default behaviour is unchanged: every stored iteration stays eligible.""" + history = _history(5) + assert pooled_history_iterations(history, target_dim=2) == [0, 1, 2, 3, 4] + assert pooled_history_iterations(history, target_dim=2, window=None) == [ + 0, 1, 2, 3, 4 + ] + + +def test_window_zero_excludes_all_history(): + """window=0 is the cycle-local (PaCS-MD-like) pool: no historical frames.""" + history = _history(5) + assert pooled_history_iterations(history, target_dim=2, window=0) == [] + + +def test_window_keeps_only_most_recent_iterations(): + history = _history(5) + assert pooled_history_iterations(history, target_dim=2, window=1) == [4] + assert pooled_history_iterations(history, target_dim=2, window=2) == [3, 4] + + +def test_window_larger_than_history_keeps_everything(): + history = _history(3) + assert pooled_history_iterations(history, target_dim=2, window=99) == [0, 1, 2] + + +def test_window_applies_after_the_dimension_filter(): + """Dimension filtering still governs eligibility; the window trims what survives. + + A 1-D initial-trajectory projection at iteration -1 must not consume a slot in a + 2-D window, or the pool would silently shrink. + """ + history = { + -1: {"projection": np.zeros((4, 1), dtype=float)}, + 0: {"projection": np.zeros((4, 2), dtype=float)}, + 1: {"projection": np.zeros((4, 2), dtype=float)}, + 2: {"projection": np.zeros((4, 2), dtype=float)}, + } + assert pooled_history_iterations(history, target_dim=2, window=2) == [1, 2] + assert pooled_history_iterations(history, target_dim=1, window=2) == [-1] + + +def test_negative_window_is_rejected(): + """A negative window is a configuration error, not a silent full-history pool.""" + history = _history(3) + try: + pooled_history_iterations(history, target_dim=2, window=-1) + except ValueError: + return + raise AssertionError("expected ValueError for a negative window") + + +# --- wiring: the window must reach the candidate pool and the config --------- + + +def test_cumulative_points_respects_window(): + """The density spawner's candidate pool shrinks with the window. + + Historical frames are stacked before the current ones, so pool size is the + observable that decides which frames a spawn index can reach. + """ + from trails_md.spawners.density import _cumulative_points + + history = _history(3, dim=2, frames=4) # 3 iterations x 4 frames = 12 + points = np.zeros((4, 2), dtype=float) # current iteration + + assert _cumulative_points(points, history).shape[0] == 16 # default: all + assert _cumulative_points(points, history, window=None).shape[0] == 16 + assert _cumulative_points(points, history, window=1).shape[0] == 8 + assert _cumulative_points(points, history, window=0).shape[0] == 4 # cycle-local + + +def test_density_spawner_honours_history_window(): + """A spawner built with history_window=0 can only return current-iteration frames.""" + from trails_md.spawners.density import DensitySpawner + + history = _history(3, dim=2, frames=4) + points = np.zeros((4, 2), dtype=float) + + cycle_local = DensitySpawner( + n_bins=[4, 4], min_values=[-1.0, -1.0], max_values=[1.0, 1.0], + history_window=0, + ) + picks = cycle_local.sample(points, top_n=4, history=history) + assert picks, "spawner returned no frames" + assert max(picks) < points.shape[0], ( + f"cycle-local spawner reached a historical frame: {picks}" + ) + + cumulative = DensitySpawner( + n_bins=[4, 4], min_values=[-1.0, -1.0], max_values=[1.0, 1.0], + ) + assert cumulative.sample(points, top_n=4, history=history) is not None + + +def test_config_exposes_history_window_defaulting_to_full_history(): + from trails_md.config import SpawningConfig + + assert SpawningConfig().history_window is None + assert SpawningConfig(history_window=0).history_window == 0 + + +def test_frame_records_stay_index_synced_with_the_windowed_pool(): + """core's frame/lineage mapping must use the same window as the spawner pool. + + pooled_history_iterations is the single source of truth precisely so these two + cannot diverge. If core kept the full history while the spawner was windowed, a + spawn index would resolve to a different conformation than the one whose lineage + is recorded -- silently, with no error anywhere. + """ + import types + + from tests.test_review_fixes import _bare_core, _iteration_entry + from trails_md.spawners.density import _historical_points + + sampler = _bare_core() + sampler.config.spawning.history_window = 0 + sampler.history = { + 0: _iteration_entry(0, n_walkers=2, frames_per_walker=3, n_features=2), + 1: _iteration_entry(1, n_walkers=2, frames_per_walker=3, n_features=2), + } + current = np.zeros((6, 2), dtype=float) + + pool = _historical_points(current, sampler.history, window=0) + records = sampler._sampling_frame_records([], target_dim=2) + trajectories = sampler._sampling_trajectories([], target_dim=2) + + assert pool.shape[0] == 0, "cycle-local pool should contain no historical frames" + assert len(records) == pool.shape[0], ( + f"frame records ({len(records)}) desynced from spawner pool ({pool.shape[0]})" + ) + assert len(trajectories) == 0, "cycle-local run should pool no historical trajectories" + + +def test_default_config_leaves_frame_records_on_full_history(): + """Regression guard: without a window, core still pools the entire history.""" + from tests.test_review_fixes import _bare_core, _iteration_entry + from trails_md.spawners.density import _historical_points + + sampler = _bare_core() + sampler.config.spawning.history_window = None + sampler.history = { + 0: _iteration_entry(0, n_walkers=2, frames_per_walker=3, n_features=2), + 1: _iteration_entry(1, n_walkers=2, frames_per_walker=3, n_features=2), + } + current = np.zeros((6, 2), dtype=float) + + pool = _historical_points(current, sampler.history) + records = sampler._sampling_frame_records([], target_dim=2) + assert pool.shape[0] == 12 + assert len(records) == 12 diff --git a/tests/test_riteweight.py b/tests/test_riteweight.py new file mode 100644 index 0000000..3c99761 --- /dev/null +++ b/tests/test_riteweight.py @@ -0,0 +1,226 @@ +"""RiteWeight trajectory reweighting. + +Implements the algorithm of Kania et al., PNAS 123, e2529246123 (2026), +"Randomized iterative trajectory reweighting for steady-state distributions +without discretization error". Written from the published description; the +authors' reference code carries no licence and was not copied. + +The decisive test is a discrete Markov chain whose stationary distribution is +exactly calculable: we deliberately draw segment start states from the WRONG +distribution, then require RiteWeight to recover the true stationary weights. +That is ground truth, not a regression baseline. +""" + +import numpy as np +import pytest + +from trails_md.analysis.riteweight import riteweight + + +# A 3-state chain with a well-separated stationary distribution. +P_TRUE = np.array( + [ + [0.80, 0.15, 0.05], + [0.10, 0.70, 0.20], + [0.05, 0.25, 0.70], + ] +) + + +def exact_stationary(P): + vals, vecs = np.linalg.eig(P.T) + v = np.real(vecs[:, np.argmin(np.abs(vals - 1.0))]) + return v / v.sum() + + +def make_segments(n=6000, skew=(0.85, 0.10, 0.05), seed=0): + """Transition pairs from P_TRUE whose START states follow a WRONG distribution. + + Each state is embedded at a well-separated 1-D coordinate so that any sane + clustering resolves the three states. + """ + rng = np.random.default_rng(seed) + starts = rng.choice(3, size=n, p=np.asarray(skew) / np.sum(skew)) + ends = np.array([rng.choice(3, p=P_TRUE[s]) for s in starts]) + coord = np.array([0.0, 10.0, 20.0]) + jitter = 0.05 + x0 = coord[starts] + rng.normal(0, jitter, n) + x1 = coord[ends] + rng.normal(0, jitter, n) + return x0.reshape(-1, 1), x1.reshape(-1, 1), starts + + +def state_weights(weights, starts): + """Total weight assigned to segments starting in each state.""" + return np.array([weights[starts == s].sum() for s in range(3)]) + + +def test_recovers_exact_stationary_distribution_from_skewed_data(): + """The headline claim: mis-distributed input, correct stationary output.""" + x0, x1, starts = make_segments() + pi_true = exact_stationary(P_TRUE) + + res = riteweight(x0, x1, n_clusters=8, n_iterations=3000, + average_last=500, seed=1) + + got = state_weights(res.weights, starts) + # Uniform initial weights reproduce the skewed input distribution, so the + # test only means something if the starting point is genuinely wrong. + naive = state_weights(np.full(len(starts), 1 / len(starts)), starts) + assert np.abs(naive - pi_true).max() > 0.25, "test setup is not actually skewed" + assert np.abs(got - pi_true).max() < 0.05, ( + f"stationary distribution not recovered: got {got}, want {pi_true}" + ) + + +def test_fixed_point_is_independent_of_cluster_count(): + """Paper's central claim: the fixed point does not depend on n_clusters.""" + x0, x1, starts = make_segments() + results = {} + for n_clusters in (4, 12, 40): + res = riteweight(x0, x1, n_clusters=n_clusters, n_iterations=3000, + average_last=500, seed=2) + results[n_clusters] = state_weights(res.weights, starts) + ref = results[4] + for n_clusters, got in results.items(): + assert np.abs(got - ref).max() < 0.05, ( + f"cluster count {n_clusters} changed the fixed point: {got} vs {ref}" + ) + + +def test_weights_stay_normalised_and_positive(): + x0, x1, _ = make_segments(n=1500) + res = riteweight(x0, x1, n_clusters=6, n_iterations=300, average_last=50, seed=3) + assert np.isclose(res.weights.sum(), 1.0), "weights must remain normalised" + assert (res.weights > 0).all(), "weights must remain strictly positive" + assert len(res.weights) == len(x0) + + +def test_already_correct_input_is_left_alone(): + """Idempotence: start from the true distribution and weights should barely move.""" + x0, x1, starts = make_segments(skew=tuple(exact_stationary(P_TRUE)), seed=4) + n = len(starts) + res = riteweight(x0, x1, n_clusters=8, n_iterations=1500, + average_last=300, seed=5) + before = state_weights(np.full(n, 1 / n), starts) + after = state_weights(res.weights, starts) + assert np.abs(after - before).max() < 0.05, ( + f"already-stationary input was disturbed: {before} -> {after}" + ) + + +def test_is_reproducible_under_a_fixed_seed(): + x0, x1, _ = make_segments(n=1200) + a = riteweight(x0, x1, n_clusters=6, n_iterations=400, average_last=100, seed=7) + b = riteweight(x0, x1, n_clusters=6, n_iterations=400, average_last=100, seed=7) + np.testing.assert_allclose(a.weights, b.weights) + + +def test_rejects_mismatched_inputs(): + x0, x1, _ = make_segments(n=100) + with pytest.raises(ValueError): + riteweight(x0, x1[:50], n_clusters=4, n_iterations=10) + with pytest.raises(ValueError): + riteweight(x0, x1, n_clusters=0, n_iterations=10) + + +def test_reports_convergence_diagnostics(): + """A user must be able to tell whether the iteration actually settled.""" + x0, x1, _ = make_segments(n=1500) + res = riteweight(x0, x1, n_clusters=6, n_iterations=800, average_last=200, seed=8) + assert res.n_iterations == 800 + assert len(res.weight_drift) == 800 + # Drift is the per-iteration change in weights; it must decrease overall. + early = np.mean(res.weight_drift[:100]) + late = np.mean(res.weight_drift[-100:]) + assert late < early, f"weights not settling: early={early:.3e} late={late:.3e}" + + +# --- the test that actually exercises the "Randomized" in RiteWeight --------- +# +# The tests above use three well-separated states, which ANY clustering with +# k >= 3 resolves exactly. That leaves no discretisation error to average away, +# so they cannot detect removal of the per-iteration re-clustering -- a mutation +# run confirmed they all still pass with a single fixed clustering. +# +# Here the states are deliberately UNRESOLVABLE by one clustering: 12 microstates +# on a line, discretised into only 4 clusters. With a fixed clustering the +# relative weights of states sharing a cluster are frozen at their (wrong) +# initial values forever. Only re-randomising the clusters lets those ties be +# broken, which is precisely the paper's claim that the fixed point is set by the +# microstate transition matrix and not by the discretisation. + +N_STATES = 12 + + +def _line_chain(): + """Nearest-neighbour walk on a line, biased so the stationary law is non-uniform.""" + # Mild bias only. A strong bias concentrates the stationary law on the last + # state, and if the skewed sampling then barely visits it the test would be + # demanding a fix for missing COVERAGE -- which RiteWeight explicitly cannot + # provide. Keep every state well sampled so this isolates mis-WEIGHTING. + P = np.zeros((N_STATES, N_STATES)) + for s in range(N_STATES): + right = 0.35 if s < N_STATES - 1 else 0.0 + left = 0.30 if s > 0 else 0.0 + P[s, s] = 1.0 - right - left + if s < N_STATES - 1: + P[s, s + 1] = right + if s > 0: + P[s, s - 1] = left + return P + + +def _line_segments(n=20000, seed=11): + P = _line_chain() + rng = np.random.default_rng(seed) + # Wrong distribution, but every state still receives ample statistics: + # roughly the reverse of the true stationary ordering. + skew = (0.30 / 0.35) ** np.arange(N_STATES) + skew /= skew.sum() + starts = rng.choice(N_STATES, size=n, p=skew) + ends = np.array([rng.choice(N_STATES, p=P[s]) for s in starts]) + jitter = 0.02 + x0 = starts + rng.normal(0, jitter, n) + x1 = ends + rng.normal(0, jitter, n) + return x0.reshape(-1, 1), x1.reshape(-1, 1), starts, exact_stationary(P) + + +def test_recovers_structure_finer_than_the_clustering(): + """Fewer clusters than states: only re-randomised clustering can succeed.""" + x0, x1, starts, pi_true = _line_segments() + res = riteweight(x0, x1, n_clusters=4, n_iterations=4000, + average_last=800, seed=21) + got = np.array([res.weights[starts == s].sum() for s in range(N_STATES)]) + naive = np.array([(starts == s).mean() for s in range(N_STATES)]) + # The input must genuinely be wrong, or the test proves nothing. + assert np.abs(naive - pi_true).max() > 0.10, "test setup is not skewed enough" + assert np.abs(got - pi_true).max() < 0.05, ( + f"failed to resolve structure finer than the clustering:\n" + f" got {np.round(got, 3)}\n want {np.round(pi_true, 3)}" + ) + + +def test_clustering_assigns_each_point_to_its_nearest_centre(): + """Lock the clustering CONTRACT before optimising its implementation. + + The assignment must be exact nearest-centre. This guards the KD-tree + implementation against silently returning approximate neighbours, which + would bias the transition matrix in a way no other test would catch. + """ + from trails_md.analysis.riteweight import _random_clustering + + rng = np.random.default_rng(0) + pts = rng.normal(size=(500, 3)) + labels = _random_clustering(pts, n_clusters=7, rng=np.random.default_rng(1)) + + # Recover the centres the routine chose, then verify by brute force. + centres = np.array([pts[labels == c].mean(axis=0) for c in range(labels.max() + 1)]) + # A point must be at least as close to its own cluster's members as implied; + # test the invariant directly against the true centre set instead. + centre_idx = np.random.default_rng(1).choice(len(pts), size=7, replace=False) + true_centres = pts[centre_idx] + brute = np.argmin( + np.linalg.norm(pts[:, None, :] - true_centres[None, :, :], axis=2), axis=1 + ) + np.testing.assert_array_equal(labels, brute) + assert centres.shape[0] <= 7 diff --git a/tests/test_tvae_beta.py b/tests/test_tvae_beta.py new file mode 100644 index 0000000..94615e4 --- /dev/null +++ b/tests/test_tvae_beta.py @@ -0,0 +1,83 @@ +"""The TVAE KLD weight (beta) must be settable, recorded, and default to 1.0. + +deeptime's TVAE loss is ``mse + beta * kld / n_features``. Until now Trails-MD +constructed the estimator without passing ``beta`` at all, so every reported result +silently used deeptime's default of 1.0 and no user could change it. A referee asking +"what beta did you use?" could only be answered by reading a dependency's source. + +These tests pin both halves of the fix: the value reaches the estimator, and the +default is the 1.0 that all previously published Trails-MD results were produced with. +""" + +from __future__ import annotations + +import warnings + +import numpy as np +import pytest + +warnings.filterwarnings("ignore") + +from trails_md.spaces.model import AdaptiveSpaceModel # noqa: E402 + +torch = pytest.importorskip("torch") + + +def _features(n_walkers=2, walker_length=24, n_features=6, seed=0): + rng = np.random.default_rng(seed) + return rng.normal(size=(n_walkers * walker_length, n_features)) + + +def test_tvae_beta_defaults_to_one(): + """1.0 is deeptime's default and therefore what every published run used.""" + model = AdaptiveSpaceModel(space_mode="tvae") + assert model.tvae_beta == pytest.approx(1.0) + + +def test_tvae_beta_reaches_the_estimator(): + """Setting beta must change the estimator's KLD weight, not just an attribute. + + Asserting on ``model.model._beta`` after a real fit checks the wiring end to end: + an implementation that stored the value but kept calling ``TVAE(...)`` without it + would pass an attribute check and fail this one. + """ + model = AdaptiveSpaceModel( + space_mode="tvae", tvae_beta=0.25, epochs=1, lagtime=1, latent_dim=2 + ) + model.fit(_features(), walker_length=24, n_walkers=2) + assert model.model._beta == pytest.approx(0.25) + + +def test_tvae_default_beta_reaches_the_estimator(): + model = AdaptiveSpaceModel(space_mode="tvae", epochs=1, lagtime=1, latent_dim=2) + model.fit(_features(), walker_length=24, n_walkers=2) + assert model.model._beta == pytest.approx(1.0) + + +def test_tvae_beta_must_be_positive(): + """A non-positive KLD weight is not a meaningful VAE objective.""" + with pytest.raises(ValueError): + AdaptiveSpaceModel(space_mode="tvae", tvae_beta=0.0) + with pytest.raises(ValueError): + AdaptiveSpaceModel(space_mode="tvae", tvae_beta=-1.0) + + +def test_tvae_beta_survives_a_checkpoint_round_trip(): + """Restored checkpoints must keep beta; otherwise a resumed run silently + changes its objective mid-campaign.""" + import pickle + + model = AdaptiveSpaceModel(space_mode="tvae", tvae_beta=0.4) + restored = pickle.loads(pickle.dumps(model)) + assert restored.tvae_beta == pytest.approx(0.4) + + +def test_old_checkpoint_without_beta_restores_to_one(): + """Checkpoints written before beta existed must keep reproducing their original + behaviour, which was deeptime's default of 1.0.""" + model = AdaptiveSpaceModel(space_mode="tvae", tvae_beta=0.4) + state = dict(model.__dict__) + del state["tvae_beta"] + revived = AdaptiveSpaceModel.__new__(AdaptiveSpaceModel) + revived.__setstate__(state) + assert revived.tvae_beta == pytest.approx(1.0) diff --git a/trails_md/__init__.py b/trails_md/__init__.py index 2778d9e..ddfa70f 100644 --- a/trails_md/__init__.py +++ b/trails_md/__init__.py @@ -4,7 +4,7 @@ A modular, extensible, and scalable framework for enhanced molecular dynamics sampling. """ -__version__ = "1.0.0" +__version__ = "1.1.0" __all__ = ["TrailsMDConfig", "TrailsMDCore"] diff --git a/trails_md/analysis/__init__.py b/trails_md/analysis/__init__.py index 6516227..4bd45b9 100644 --- a/trails_md/analysis/__init__.py +++ b/trails_md/analysis/__init__.py @@ -1,9 +1,12 @@ """Post-hoc MSM analysis and plotting utilities. ``data`` holds matplotlib-free numerics (loading msm.npz/cvs.npz, free -energies); ``plots`` holds the matplotlib visualisations. +energies); ``plots`` holds the matplotlib visualisations; ``riteweight`` +reweights an adaptively-sampled ensemble towards its stationary distribution +without a lag-time or Markovianity assumption. """ from . import data +from .riteweight import RiteWeightResult, riteweight -__all__ = ["data"] +__all__ = ["data", "riteweight", "RiteWeightResult"] diff --git a/trails_md/analysis/riteweight.py b/trails_md/analysis/riteweight.py new file mode 100644 index 0000000..bf764f8 --- /dev/null +++ b/trails_md/analysis/riteweight.py @@ -0,0 +1,275 @@ +"""RiteWeight: randomized iterative trajectory reweighting. + +Estimates the stationary (equilibrium or steady-state) distribution from +mis-distributed trajectory data by iteratively reweighting short trajectory +segments until their weighted distribution is self-consistent with the +stationary distribution of a transition matrix built from those same weights. + +Reference +--------- +S. Kania, R. J. Webber, G. Simpson, D. Aristoff and D. M. Zuckerman, +"Randomized iterative trajectory reweighting for steady-state distributions +without discretization error", PNAS 123, e2529246123 (2026). + +This is an independent implementation written from the published algorithm. +The authors' reference implementation (github.com/ZuckermanLab/rite_weight) +carries no licence statement and was deliberately not copied or vendored. + +Why this matters for adaptive sampling +-------------------------------------- +Adaptive selection deliberately over-samples sparse regions, so the +configurations produced by an exploration campaign are *not* Boltzmann +distributed. RiteWeight corrects exactly that: it fixes mis-*weighting*. +Two consequences follow, and both are limitations worth stating plainly: + +* It cannot fix mis-*coverage*. A basin never visited cannot be reweighted + into existence, and no diagnostic here will tell you one is missing. +* It requires segment pairs drawn from unbiased dynamics. TRAILS-MD applies + no biasing forces, but velocities are redrawn at each respawn, so pairs + must never straddle a respawn boundary -- the same constraint that governs + the time-lagged CV estimators. + +The key algorithmic device is that the clustering is *re-randomised every +iteration*. Segments that shared a cluster (and so kept fixed relative +weights) in one iteration are separated in the next, so the fixed point is +governed by the microstate transition matrix rather than by any particular +discretisation. In practice this means the answer should not depend on the +number of clusters -- a property worth asserting in tests. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np + +__all__ = ["riteweight", "RiteWeightResult"] + + +@dataclass +class RiteWeightResult: + """Outcome of a RiteWeight run. + + Attributes + ---------- + weights: + Final per-segment weights, normalised to sum to one. These are the + weights averaged over the last ``average_last`` iterations, which + suppresses the residual noise from the random clustering. + n_iterations: + Number of iterations actually performed. + weight_drift: + Per-iteration L1 change in the weight vector. A run that has settled + shows this decreasing and then fluctuating about a small value; a run + that has not settled does not. Reported so callers can judge + convergence rather than assume it. + n_clusters: + Cluster count used for each random discretisation. + """ + + weights: np.ndarray + n_iterations: int + weight_drift: np.ndarray = field(repr=False) + n_clusters: int = 0 + + @property + def converged(self) -> bool: + """Heuristic: late drift is well below early drift and small in absolute terms.""" + if len(self.weight_drift) < 20: + return False + early = float(np.mean(self.weight_drift[: max(1, len(self.weight_drift) // 10)])) + late = float(np.mean(self.weight_drift[-max(1, len(self.weight_drift) // 10) :])) + return late < 0.5 * early and late < 1e-2 + + +def _random_clustering( + points: np.ndarray, n_clusters: int, rng: np.random.Generator +) -> np.ndarray: + """Assign points to clusters seeded at randomly chosen data points. + + A fresh random seeding each iteration is the mechanism that averages away + discretisation error, so this deliberately does NOT run k-means to + convergence -- the clustering is meant to differ between iterations. + """ + n = len(points) + k = min(n_clusters, n) + centre_idx = rng.choice(n, size=k, replace=False) + centres = points[centre_idx] + + # Exact nearest-centre assignment. A KD-tree over the (few) centres turns the + # cost from O(n*k) into roughly O(n log k), which is the difference between + # feasible and not: a production campaign supplies ~10^6 pooled points and + # thousands of iterations, where the brute-force form needs ~10^11 distance + # evaluations. The query is exact (k=1 nearest neighbour), so the assignment + # is identical to brute force -- a test pins that equivalence. + try: + from scipy.spatial import cKDTree + + tree = cKDTree(centres) + _, labels = tree.query(points, k=1, workers=-1) + return labels.astype(np.int64) + except ImportError: # pragma: no cover - scipy is a hard dependency in practice + labels = np.empty(n, dtype=np.int64) + chunk = max(1, int(2e7 // max(1, k))) + for lo in range(0, n, chunk): + hi = min(n, lo + chunk) + d = np.linalg.norm(points[lo:hi, None, :] - centres[None, :, :], axis=2) + labels[lo:hi] = np.argmin(d, axis=1) + return labels + + +def _stationary_distribution(T: np.ndarray) -> np.ndarray: + """Left eigenvector of T for eigenvalue 1, normalised to a probability vector. + + Falls back to a uniform distribution if the spectrum is degenerate or the + dominant eigenvector is unusable (which happens when a random clustering + yields a disconnected transition graph). + """ + n = T.shape[0] + try: + vals, vecs = np.linalg.eig(T.T) + except np.linalg.LinAlgError: + return np.full(n, 1.0 / n) + v = np.real(vecs[:, int(np.argmin(np.abs(vals - 1.0)))]) + v = np.abs(v) # sign is arbitrary; a probability vector is non-negative + s = v.sum() + if not np.isfinite(s) or s <= 0: + return np.full(n, 1.0 / n) + return v / s + + +def riteweight( + start_features: np.ndarray, + end_features: np.ndarray, + *, + n_clusters: int = 100, + n_iterations: int = 5000, + initial_weights: np.ndarray | None = None, + learning_rate: float = 1.0, + average_last: int = 1000, + seed: int | None = None, +) -> RiteWeightResult: + """Reweight trajectory segments toward the stationary distribution. + + Parameters + ---------- + start_features, end_features: + ``(n_segments, n_features)`` arrays holding the featurised first and + second configuration of each segment. Features must be invariant to + rotation and translation (pairwise distances, torsions, a learned + projection, ...). Both configurations of a pair must come from the + same continuous stretch of unbiased dynamics. + n_clusters: + Number of clusters per random discretisation. The fixed point should + be insensitive to this; it trades resolution against the statistics + available per cluster. + n_iterations: + Reweighting iterations. Each uses a fresh random clustering. + initial_weights: + Optional prior over segments; defaults to uniform, the uninformative + choice. Note the fixed point depends on this prior, so a strongly + informative and wrong prior can bias the result. + learning_rate: + Exponent damping the multiplicative update, ``(pi_I / w_I) ** alpha``. + 1.0 applies the full update; smaller values damp oscillation when the + per-cluster statistics are sparse. + average_last: + Number of trailing iterations to average the weights over, which + suppresses noise from the random clustering. + seed: + Seed for the clustering RNG, for reproducibility. + + Returns + ------- + RiteWeightResult + """ + start_features = np.asarray(start_features, dtype=float) + end_features = np.asarray(end_features, dtype=float) + if start_features.ndim == 1: + start_features = start_features.reshape(-1, 1) + if end_features.ndim == 1: + end_features = end_features.reshape(-1, 1) + if start_features.shape != end_features.shape: + raise ValueError( + "start_features and end_features must have the same shape; got " + f"{start_features.shape} and {end_features.shape}" + ) + if n_clusters < 1: + raise ValueError(f"n_clusters must be >= 1; got {n_clusters}") + if n_iterations < 1: + raise ValueError(f"n_iterations must be >= 1; got {n_iterations}") + + n_segments = len(start_features) + if initial_weights is None: + weights = np.full(n_segments, 1.0 / n_segments) + else: + weights = np.asarray(initial_weights, dtype=float).copy() + if len(weights) != n_segments: + raise ValueError("initial_weights must have one entry per segment") + if np.any(weights <= 0): + raise ValueError("initial_weights must be strictly positive") + weights /= weights.sum() + + rng = np.random.default_rng(seed) + # Cluster over the pooled configurations so that a segment's start and end + # are discretised on the same footing. + pooled = np.vstack([start_features, end_features]) + + drift = np.empty(n_iterations) + accum = np.zeros(n_segments) + n_accum = 0 + keep_from = max(0, n_iterations - average_last) + + for it in range(n_iterations): + labels = _random_clustering(pooled, n_clusters, rng) + start_labels = labels[:n_segments] + end_labels = labels[n_segments:] + k = int(labels.max()) + 1 + + # Weighted transition counts between clusters, then row-normalise. + counts = np.zeros((k, k)) + np.add.at(counts, (start_labels, end_labels), weights) + row = counts.sum(axis=1) + occupied = row > 0 + T = np.zeros((k, k)) + T[occupied] = counts[occupied] / row[occupied, None] + # Unvisited clusters get a self-transition so T stays a stochastic + # matrix; they carry no weight and cannot affect the update below. + T[~occupied, ~occupied] = 1.0 + + pi = _stationary_distribution(T) + + # Rescale each cluster's total weight to pi_I, leaving the relative + # weights of segments inside a cluster untouched (Eq. 2 of the paper). + w_cluster = np.zeros(k) + np.add.at(w_cluster, start_labels, weights) + ratio = np.ones(k) + good = (w_cluster > 0) & (pi > 0) + ratio[good] = pi[good] / w_cluster[good] + if learning_rate != 1.0: + ratio = ratio**learning_rate + + new_weights = weights * ratio[start_labels] + total = new_weights.sum() + if not np.isfinite(total) or total <= 0: + # Degenerate update (all mass annihilated): keep the previous + # weights rather than propagate NaNs. + new_weights = weights.copy() + total = new_weights.sum() + new_weights /= total + + drift[it] = float(np.abs(new_weights - weights).sum()) + weights = new_weights + + if it >= keep_from: + accum += weights + n_accum += 1 + + final = accum / n_accum if n_accum else weights + final = final / final.sum() + return RiteWeightResult( + weights=final, + n_iterations=n_iterations, + weight_drift=drift, + n_clusters=n_clusters, + ) diff --git a/trails_md/config.py b/trails_md/config.py index 326600f..f7fd71e 100644 --- a/trails_md/config.py +++ b/trails_md/config.py @@ -80,6 +80,13 @@ class SpawningConfig(BaseModel): max_workers: int = 4 target: list[float] | None = None recent_density_window: int = 5 + # How far back the spawning candidate pool reaches. ``None`` (default) keeps the + # full accumulated history. ``0`` restricts spawning to the current iteration, + # reproducing the memoryless, cycle-local selection of PaCS-MD so the two can be + # compared under an otherwise identical loop. ``N`` keeps the N most recent + # iterations. Note this bounds *selection* only; coverage/convergence diagnostics + # continue to measure the whole campaign. + history_window: int | None = None voronoi_clusters: int = 150 voronoi_periodic: bool = False voronoi_grid_size: int = 250 @@ -117,6 +124,11 @@ class AdaptiveModelConfig(BaseModel): decoder_hidden_dims: list[int] = [128, 256] dropout_rate: float = 0.1 deep_tica_hidden_dims: list[int] = [256, 128] + # TVAE only: weight of the KL term in ``mse + tvae_beta * kld / n_features``. + # 1.0 reproduces every previously published Trails-MD result; values below 1 + # relax the latent prior and favour reconstruction, values above it favour a + # smoother, more disentangled latent space at the cost of reconstruction. + tvae_beta: float = 1.0 # SPIB (State Predictive Information Bottleneck) hyperparameters. spib_n_states: int = 10 spib_beta: float = 1e-3 @@ -153,6 +165,13 @@ def validate_dropout_rate(cls, value: float) -> float: raise ValueError("dropout_rate must be >= 0 and < 1") return value + @field_validator("tvae_beta", "spib_beta") + @classmethod + def validate_beta(cls, value: float) -> float: + if value <= 0: + raise ValueError("beta must be greater than 0") + return value + @field_validator("encoder_hidden_dims", "decoder_hidden_dims", "deep_tica_hidden_dims") @classmethod def validate_hidden_dims(cls, value: list[int]) -> list[int]: diff --git a/trails_md/core.py b/trails_md/core.py index ef57770..2fce8e8 100644 --- a/trails_md/core.py +++ b/trails_md/core.py @@ -78,6 +78,7 @@ def __init__(self, config_dict: dict[str, Any]): probabilistic=self.config.spawning.spawn_type != "hard", target=self.config.spawning.target, recent_window=self.config.spawning.recent_density_window, + history_window=self._history_window(), n_clusters=self.config.spawning.voronoi_clusters, periodic=self.config.spawning.voronoi_periodic, grid_size=self.config.spawning.voronoi_grid_size, @@ -1004,6 +1005,15 @@ def run_iteration(self, walkers: list[Any]): "convergence_reason": self.convergence_reason, } + def _history_window(self) -> int | None: + """Configured spawning history window (None = full cumulative history). + + Read through one accessor because the spawner's candidate pool and the + trajectory/frame-record mapping must agree exactly; if they used different + windows a spawn index would resolve to the wrong conformation silently. + """ + return getattr(self.config.spawning, "history_window", None) + def _sampling_trajectories( self, current_trajectories: list[str], target_dim: int | None = None ) -> list[str]: @@ -1011,7 +1021,9 @@ def _sampling_trajectories( # (via pooled_history_iterations), so a spawn index maps to the intended # trajectory even when history mixes projection dimensionalities. trajectories: list[str] = [] - for iteration in pooled_history_iterations(self.history, target_dim): + for iteration in pooled_history_iterations( + self.history, target_dim, self._history_window() + ): entry = self.history[iteration] stored = entry.get("trajectories") if stored: @@ -1029,7 +1041,9 @@ def _sampling_frame_records( target_dim: int | None = None, ) -> list[dict[str, Any]]: records: list[dict[str, Any]] = [] - for iteration in pooled_history_iterations(self.history, target_dim): + for iteration in pooled_history_iterations( + self.history, target_dim, self._history_window() + ): entry = self.history[iteration] stored = entry.get("frames") if stored: diff --git a/trails_md/spaces/model.py b/trails_md/spaces/model.py index 0f31c93..5e14597 100644 --- a/trails_md/spaces/model.py +++ b/trails_md/spaces/model.py @@ -27,6 +27,11 @@ class AdaptiveSpaceModel: "decoder_hidden_dims": [128, 256], "dropout_rate": 0.1, "deep_tica_hidden_dims": [256, 128], + # 1.0 is deeptime's default and is what every Trails-MD result published + # before this parameter existed was produced with. Keeping it here means a + # checkpoint written before ``tvae_beta`` was introduced restores to the + # behaviour it actually had, rather than to a new default. + "tvae_beta": 1.0, "spib_n_states": 10, "spib_beta": 1e-3, "seed": 0, @@ -44,6 +49,7 @@ def __init__( decoder_hidden_dims: list[int] | None = None, dropout_rate: float = 0.1, deep_tica_hidden_dims: list[int] | None = None, + tvae_beta: float = 1.0, spib_n_states: int = 10, spib_beta: float = 1e-3, seed: int = 0, @@ -60,6 +66,9 @@ def __init__( self.decoder_hidden_dims = list(decoder_hidden_dims or [128, 256]) self.dropout_rate = float(dropout_rate) self.deep_tica_hidden_dims = list(deep_tica_hidden_dims or [256, 128]) + self.tvae_beta = float(tvae_beta) + if self.tvae_beta <= 0: + raise ValueError(f"tvae_beta must be greater than 0; got {tvae_beta}") self.spib_n_states = int(spib_n_states) self.spib_beta = float(spib_beta) self.scaler = TrajectoryScaler("minmax") @@ -157,7 +166,13 @@ def fit(self, features: np.ndarray, walker_length: int, n_walkers: int): hidden_dims=self.decoder_hidden_dims, dropout_rate=self.dropout_rate, ).to(self.device) - self.model = TVAE(encoder, decoder, learning_rate=self.learning_rate) + # beta weights the KL term: loss = mse + beta * kld / n_features. + self.model = TVAE( + encoder, + decoder, + learning_rate=self.learning_rate, + beta=self.tvae_beta, + ) # Prepare data dataset = TrajectoryDataset.from_trajectories(self.lagtime, traj_list) diff --git a/trails_md/spawners/density.py b/trails_md/spawners/density.py index e88f766..160e7b2 100644 --- a/trails_md/spawners/density.py +++ b/trails_md/spawners/density.py @@ -23,6 +23,7 @@ def __init__( probabilistic: bool = True, target: list[float] | None = None, recent_window: int = 5, + history_window: int | None = None, **kwargs: Any, ): super().__init__(**kwargs) @@ -33,6 +34,8 @@ def __init__( self.probabilistic = probabilistic self.target = target self.recent_bins: deque[set[Any]] = deque(maxlen=recent_window) + # Bounds how far back the candidate pool reaches (None = full history). + self.history_window = history_window # Optional landscape-adaptive binner (set by the orchestrator); None -> grid. self.binner = None @@ -40,7 +43,7 @@ def sample( self, points: np.ndarray, top_n: int, history: dict[int, Any] | None = None ) -> list[int]: points = np.asarray(points, dtype=float) - cumulative_points = _cumulative_points(points, history) + cumulative_points = _cumulative_points(points, history, self.history_window) if self.binner is not None: # Keep the adaptive binner in sync with resolution bumps. self.binner.n_bins = np.asarray(self.n_bins, dtype=int) @@ -131,16 +134,20 @@ def _sample_frames( def _cumulative_points( - points: np.ndarray, history: dict[int, Any] | None + points: np.ndarray, + history: dict[int, Any] | None, + window: int | None = None, ) -> np.ndarray: - historical = _historical_points(points, history) + historical = _historical_points(points, history, window) if historical.size == 0: return points return np.vstack([historical, points]) def _historical_points( - points: np.ndarray, history: dict[int, Any] | None + points: np.ndarray, + history: dict[int, Any] | None, + window: int | None = None, ) -> np.ndarray: empty = np.empty((0, points.shape[1]), dtype=float) if not history: @@ -150,7 +157,7 @@ def _historical_points( # the candidate points here stay index-synchronized with the trajectory and # frame-record lists that core.py builds for the same spawn indices. projections = [] - for iteration in pooled_history_iterations(history, points.shape[1]): + for iteration in pooled_history_iterations(history, points.shape[1], window): projection = np.asarray(history[iteration]["projection"], dtype=float) if projection.ndim == 1: projection = projection.reshape(-1, 1) diff --git a/trails_md/spawners/history.py b/trails_md/spawners/history.py index b450a40..20ea2c6 100644 --- a/trails_md/spawners/history.py +++ b/trails_md/spawners/history.py @@ -28,7 +28,9 @@ def projection_dim(projection: Any) -> int: def pooled_history_iterations( - history: dict[int, Any] | None, target_dim: int | None + history: dict[int, Any] | None, + target_dim: int | None, + window: int | None = None, ) -> list[int]: """Sorted history iterations whose projection joins the cumulative pool. @@ -36,10 +38,24 @@ def pooled_history_iterations( dimension equals ``target_dim`` (the dimension of the current projection being spawned from). ``target_dim=None`` includes every stored projection. + ``window`` bounds how far back the pool reaches, which is what separates + cumulative spawning from the cycle-local selection used by PaCS-MD: + + * ``None`` (default) -- the full accumulated history, TRAILS-MD's normal mode. + * ``0`` -- no history at all, so only the current iteration's frames are + candidates. This reproduces PaCS-MD's memoryless frontier selection and + exists so the two can be compared under an otherwise identical loop. + * ``N`` -- the ``N`` most recent eligible iterations. + + The window is applied *after* the dimension filter, so a projection of the + wrong dimensionality cannot consume a slot and silently shrink the pool. + The returned order (ascending iteration) is the canonical concatenation order for the cumulative point cloud, the trajectory list, and the frame-record list, guaranteeing index ``i`` refers to the same frame in all three. """ + if window is not None and window < 0: + raise ValueError(f"history window must be non-negative; got {window}") if not history: return [] included: list[int] = [] @@ -53,4 +69,8 @@ def pooled_history_iterations( if target_dim is not None and projection_dim(projection) != target_dim: continue included.append(iteration) - return included + if window is None: + return included + # Note `included[-window:]` would be wrong: a window of 0 slices to the whole + # list rather than to nothing. Index from the front instead. + return included[max(0, len(included) - window) :] diff --git a/trails_md/templates.py b/trails_md/templates.py index a54414c..87459ad 100644 --- a/trails_md/templates.py +++ b/trails_md/templates.py @@ -60,6 +60,11 @@ voronoi_clusters: 150 # cells / microstates (voronoi & msm spawners) we_target_per_bin: 4 # walkers per bin for spawn_scheme: we lof_neighbors: 20 + # history_window: null # iterations of history the spawner scores over. + # null (default) = the whole campaign, which is what + # makes a coverage objective definable at all. Set to + # a small integer for PaCS-MD-style cycle-local + # selection; 0 scores only the current iteration. # ---- Kinetics mode (rate / MFPT): spawn_scheme: we + md_engine: openmm -------- # inherit_velocities: true # continue parent velocities (required for a rate) # recycle_target: [[-2.5, -1.0], [2.0, 3.0]] # source->sink sink box, one [lo, hi] per CV dim @@ -84,14 +89,17 @@ max_adaptive_memory_frames: 50000 adaptive_model: # hyperparameters for learned CVs - lagtime: 5 + lagtime: 5 # in FRAMES, so lag time = lagtime x stride x dt latent_dim: 2 epochs: 50 learning_rate: 0.0005 - encoder_hidden_dims: [64, 32] - decoder_hidden_dims: [32, 64] + # These are also the built-in defaults, so a config that omits the whole + # adaptive_model block trains exactly the network described here. + encoder_hidden_dims: [256, 128] + decoder_hidden_dims: [128, 256] dropout_rate: 0.1 - deep_tica_hidden_dims: [64, 32] + deep_tica_hidden_dims: [256, 128] + tvae_beta: 1.0 # TVAE only: KL weight in mse + beta*kld/n_features spib_n_states: 10 # SPIB only spib_beta: 0.001 # SPIB only