Trails-MD

Trails-MD is a Python framework for adaptive molecular dynamics campaigns. It runs many short MD walkers, projects saved frames into a collective-variable or learned latent space, chooses informative restart frames, and repeats the cycle — continuing until a sufficiently converged space is reached.

The code is meant for method development and practical sampling workflows where you need to change engines, CVs, spawning policies, or analysis criteria without rewriting the whole pipeline.

Key features

• Engine-agnostic walkers — OpenMM, GROMACS, and Amber share the same adaptive loop.
• Fixed or learned sampling spaces — user-defined physical CVs, PCA, TICA, TVAE, and Deep-TICA, swappable at configuration time.
• Interchangeable spawning policies — density, Voronoi, local-outlier-factor, farthest-point, MSM-guided, and weighted-ensemble selection.
• Two modes — exploration (discover states and pathways) and kinetics (an unbiased weighted-ensemble rate via velocity inheritance + source→sink recycling).
• Lineage-aware exploration — every spawned frame stores its parent–child ancestry, so connected transition pathways can be reconstructed from otherwise disjoint exploration stages.
• Restartable campaigns — per-iteration checkpoints capture the adaptive model, feature history, sampling state, and walker coordinates.
• HPC scalability — run on a multi-GPU workstation or dispatch walkers as SLURM / PBS array jobs (execution.backend).

A runnable notebook tutorial with rendered plots walks through the whole workflow, and an annotated input file covering the available methods and hyperparameters is documented in docs/input_file.md.

See docs/ (full documentation & tutorials) and CHANGELOG.md. Build the docs site with pip install mkdocs-material && mkdocs serve.

Motivation

Long molecular transitions are often missed by a single continuous trajectory. Adaptive sampling attacks this by running short trajectory batches and restarting new walkers from frames that look under-sampled, unusual, far apart, or close to a target region.

Trails-MD focuses on three practical requirements:

• Modularity: OpenMM, GROMACS, and Amber-style engines share the same adaptive loop.
• Interpretable sampling spaces: fixed physical CVs, custom project-file projections, PCA, TICA, TVAE, and Deep-TICA can be swapped at configuration time.
• Reproducible provenance: every iteration stores projections, trajectory paths, spawn indices, frame lineage, and checkpoints so runs can be resumed and post-processed into connected paths.

The AIB9 and alanine dipeptide examples illustrate the intended scientific use: coverage in a projected space is not automatically a transition pathway. Trails-MD records lineage so basin discovery can be distinguished from a connected mechanistic path.

Repository Structure

trails_md/
  cli.py                    Command-line entry point for adaptive runs
  config.py                 Pydantic configuration schema
  core.py                   Main adaptive sampling controller
  engines/                  OpenMM, GROMACS, and Amber backends
  spaces/                   Feature extraction and adaptive latent spaces
  spawners/                 Density, Voronoi, LOF, FPS, WE, and MSM spawning
  binning/                  Regular-grid, Voronoi, adaptive (gradient/MAB/
                            eigenvector), and weighted-ensemble split/merge
  msm/                      In-loop MSM estimation and convergence monitoring
  checkpoints/              Checkpoint save/load logic
  paths.py, path_cli.py     Lineage-aware path reconstruction
  logs.py, log_cli.py       Run-log generation utilities

examples/
  alanine_dipeptide/        Self-contained CPU hello-world (OpenMM, no GPU/GROMACS)
  AlaD/                     Alanine dipeptide GROMACS-topology examples
  AIB9/                     AIB9 fixed and learned CV examples

Installation

Create the conda environment from env.yml. It installs Trails-MD in editable mode and uses the project metadata in pyproject.toml.

conda env create -f env.yml
conda activate trails-md

Optional Deep-TICA extras:

python -m pip install -e ".[deep-tica]"

External engine executables must also be installed separately if you use those backends:

• GROMACS executable for engine.md_engine: gromacs
• Amber/pmemd executable for engine.md_engine: amber

Quick Start

The examples/alanine_dipeptide/ config is a self-contained, CPU-only hello-world: it builds its OpenMM system from bundled files and needs no GPU and no external force-field paths. (The examples/AlaD/ configs are richer but use a GROMACS topology, so they require GROMACS installed and engine.gromacs_include_dir set — see the tutorial.)

Validate a configuration before running MD:

trails-md --config examples/alanine_dipeptide/config.yaml --check

Run an adaptive campaign:

trails-md --config examples/alanine_dipeptide/config.yaml --iterations 20

Resume from the latest checkpoint:

trails-md --config examples/alanine_dipeptide/config.yaml --resume --iterations 20

Resume from a specific checkpoint:

trails-md --config examples/alanine_dipeptide/config.yaml --resume 10 --iterations 20

Generate a post-hoc exploration log for a completed run:

trails-md-log \
  --run-dir examples/alanine_dipeptide/runs/alanine_dipeptide_hello \
  --config examples/alanine_dipeptide/config.yaml

Reconstruct a connected lineage path between two CV-space points (phi/psi in radians; adjust the endpoints to two basins your run actually sampled):

trails-md-path \
  --run-dir examples/alanine_dipeptide/runs/alanine_dipeptide_hello \
  --topology examples/alanine_dipeptide/structure.pdb \
  --start=-1.4,2.6 \
  --end=-1.4,-0.7 \
  --output alad_path.xtc

For batch path extraction, use --pairs-file and --output-dir.

CLI Help

Main adaptive runner:

usage: trails-md [-h] [--config CONFIG] [--iterations ITERATIONS]
                   [--resume [RESUME]] [--ignore-missing-history] [--check]
                   [--log-level {CRITICAL,ERROR,WARNING,INFO,DEBUG}]

options:
  --config CONFIG       YAML config path. Relative paths inside it are resolved
                        from this file.
  --iterations N        Number of adaptive iterations to run (default 1).
  --resume [RESUME]     Resume from latest checkpoint, or from checkpoints/iter_N.
  --ignore-missing-history
                        On resume, tolerate a broken delta-checkpoint history
                        chain instead of aborting (use with care).
  --check               Validate inputs and executables, then exit before MD.
  --log-level LEVEL     Python logging verbosity (default WARNING).

Connected path post-processing:

usage: trails-md-path --run-dir RUN_DIR --topology TOPOLOGY
                        [--start START] [--end END] [--output OUTPUT]
                        [--pairs-file PAIRS_FILE] [--output-dir OUTPUT_DIR]
                        [--metadata METADATA] [--checkpoint CHECKPOINT]
                        [--ignore-missing-history]

Exploration log generation:

usage: trails-md-log --run-dir RUN_DIR [--config CONFIG]
                       [--output OUTPUT] [--n-bins N_BINS]
                       [--min-values MIN_VALUES] [--max-values MAX_VALUES]
                       [--append]

Configuration Overview

A minimal OpenMM configuration looks like:

system:
    conf_file: structure.pdb
    top_file: structure.pdb
    project_file: project_cvs.py

engine:
    md_engine: openmm
    platform_name: CUDA

spawning:
    spawn_scheme: density
    walker: 10
    step: 10000
    stride: 100
    max_workers: 2

space_mode: fixed
n_bins: [30, 30]
min_values: [-3.14159, -3.14159]
max_values: [3.14159, 3.14159]
outdir: runs/my_sampling_run

When space_mode: fixed, system.project_file should define:

def extract_cvs(trajectories, top_file, conf_file):
    ...
    return cvs  # shape: (n_frames, n_cvs)

See examples/AlaD/project_phi_psi.py and examples/AIB9/project_phi_psi.py for concrete examples.

Spawning Strategies

Trails-MD currently supports:

• density: regular-grid low-population bin selection.
• voronoi: KMeans-backed Voronoi cells with exact clipped polygon areas.
• lof: local-outlier-factor based frame selection.
• fps: farthest-point sampling for geometric spread.
• we: weighted-ensemble split/merge resampling (experimental).
• msm: MSM-guided least-counts spawning (experimental; see docs/msm.md).

Voronoi note: voronoi_clusters controls the restart-selection partition. The regular n_bins grid is still used for run-log coverage diagnostics.

Collective-Variable (CV) Methods

The sampling space is chosen with space_mode. Beyond fixed user CVs, several learned CV methods are available through a single registry (trails_md/spaces/registry.py), so new methods can be added in one place:

space_mode	Method	Backend	Notes
fixed	User CVs via a project file	—	e.g. AlaD phi/psi
pca	Principal component analysis	scikit-learn	linear baseline
tica	Time-lagged ICA	deeptime	linear, dynamics-aware
tvae	Time-lagged VAE	deeptime + torch	nonlinear bottleneck
deep-tica	Deep (nonlinear) TICA	mlcolvar (optional)	pip install "trails-md[deep-tica]"

Three further methods are available as experimental extension points through the same registry: vampnet (deeptime + torch), spib (torch), and deep-lda (mlcolvar). See docs/cv_methods.md for all eight.

When a model is retrained, the full feature history is reprojected into the updated latent space before spawning, so selection always reflects the current coordinates. Optional methods raise a clear, actionable error if their backend is missing.

Rates and Kinetic Seeding

Trails-MD runs in one of two modes (see docs/modes.md):

• Exploration mode (default). Walkers are short and their velocities are redrawn from a Maxwell–Boltzmann distribution at each spawn point, so the adaptive trajectories are for exploration, not an unbiased kinetic ensemble. To get a rate afterwards, discretize the explored space, select representative structures to seed longer unbiased production runs, and build a Markov State Model (MSM) from those.
• Kinetics mode. With spawn_scheme: we, inherit_velocities: true, and a source→sink recycle_target, walkers inherit their parent's velocities and weighted ensemble targets an unbiased mean-first-passage time directly during sampling (MFPT = 1/flux, the Hill relation) — no separate production stage needed.

Typical Workflow

1. Define the scientific question. Decide whether you need basin discovery, endpoint structures, or a connected transition path. These are different claims.

2. Choose a sampling space. Start with physical CVs when possible, for example AlaD phi/psi. Use learned spaces when physical CVs are unclear, then validate the results with interpretable observables.

3. Write or select a YAML config. Start from one of the YAML files under examples/. Keep path values relative to the config file when possible.

4. Preflight the run.

   trails-md --config config.yaml --check

5. Run adaptive sampling.

   trails-md --config config.yaml --iterations 100

6. Monitor output.log. The log records per-iteration timings, successful walkers, cumulative frames, diagnostic occupied bins, exploration fraction, spawn indices, and checkpoint paths.

7. Resume if needed.

   trails-md --config config.yaml --resume --iterations 50

8. Analyze coverage and endpoints. Use cvs.npz, output.log, and example notebooks/scripts to inspect the sampled CV space.

9. Check lineage before claiming a pathway. Use trails-md-path to reconstruct connected trajectories. Endpoint proximity alone does not prove that a transition was sampled.

10. Archive configs, logs, checkpoints, and analysis. A reproducible campaign should retain the YAML config, project CV file, output.log, checkpoints, and any scripts used to classify states.

Examples

Alanine dipeptide:

cd examples/AlaD
trails-md --config config.yaml --check
trails-md --config config.yaml --iterations 2

AlaD Voronoi smoke test:

cd examples/AlaD
trails-md --config config_voronoi.yaml --iterations 1

AIB9 fixed phi/psi CV:

cd examples/AIB9
trails-md --config config_fixed_phi_psi.yaml --check

Current Status

This is an active research codebase. The examples are useful starting points, but production scientific claims should be made only after checking CV definitions, sampling bounds, lineage connectivity, and system-specific validation.

License

Trails-MD is licensed under the PolyForm Noncommercial License 1.0.0. It's free to use for research, academic, educational, and other noncommercial purposes; commercial use requires a separate license from the copyright holder. See the LICENSE file for the full terms.