Release 1.1.0: TVAE KL weight, bounded spawner history, RiteWeight reweighting

.zenodo.json

@@ -0,0 +1,38 @@
+{
+  "title": "TRAILS-MD: Mapping Complex Conformational Landscapes and Transition Pathways via Lightweight Lineage-Aware Adaptive Sampling",
+  "description": "<p>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.</p><p>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.</p>",
+  "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"
+    }
+  ]
+}

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

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"

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`
 

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.:

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.

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
 

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