AutoSampler: MSM-convergence engine, HPC scaling, feature optimisation, analysis & UX

.github/workflows/ci.yml

@@ -0,0 +1,37 @@
+name: CI
+
+on:
+  push:
+    branches: [main, devel]
+  pull_request:
+
+jobs:
+  test:
+    name: Tests (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install test dependencies
+        run: |
+          python -m pip install --upgrade pip
+          # Lightweight stack covering the MSM and CV-method tests. The heavy MD
+          # backends (openmm, MDAnalysis) are not required: tests that need them
+          # skip automatically via pytest.importorskip.
+          pip install numpy scipy scikit-learn pydantic pyyaml deeptime pytest ruff
+          pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+      - name: Lint (ruff)
+        run: ruff check autosampler/msm autosampler/execution autosampler/analysis autosampler/analysis_cli.py autosampler/init_cli.py autosampler/templates.py autosampler/binning/we.py autosampler/binning/adaptive.py autosampler/spaces/registry.py autosampler/spaces/spib.py autosampler/spaces/feature_selection.py autosampler/spaces/retraining.py autosampler/utils/seeds.py autosampler/spawners/we.py autosampler/spawners/msm.py autosampler/reporting.py autosampler/engines/base.py tests
+
+      - name: Run test suite
+        run: pytest -q

.pre-commit-config.yaml

@@ -0,0 +1,18 @@
+# Pre-commit hooks for AutoSampler. Install with:
+#   pip install pre-commit && pre-commit install
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.6.9
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+      - id: check-yaml
+      - id: check-added-large-files
+        args: [--maxkb=1024]

CHANGELOG.md

@@ -0,0 +1,139 @@
+# Changelog
+
+All notable changes to AutoSampler 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/).
+
+## [Unreleased] — `devel`
+
+This cycle turns AutoSampler from a coverage-driven adaptive sampler into an
+**MSM-convergence-driven** framework, hardens the engineering foundation, and
+adds first-class **HPC scalability** and **VAMP-2 feature optimisation**. It also
+adds a flux-weighted **transition-matrix convergence** gate with
+**uncertainty-guided spawning**, and opt-in **landscape-adaptive binning**.
+
+### Added
+
+#### MSM convergence engine (Phase 1)
+- New `autosampler/msm/` subsystem built on `deeptime`:
+  - `MSMEstimator` — clustering (k-means / regular-space) on the CV/latent
+    space → transition counts → MLE **or** Bayesian MSM → implied timescales,
+    VAMP-2 score, PCCA+ metastable states, stationary distribution.
+  - `diagnostics.py` — serialisable implied-timescale / Chapman-Kolmogorov /
+    VAMP results.
+  - `ConvergenceMonitor` with composable, pluggable criteria: implied-timescale
+    stability, VAMP-2 plateau, stationary-distribution drift, Bayesian
+    statistical-error thresholds, and a **flux-weighted transition-matrix**
+    criterion (analytic Dirichlet `T_ij` uncertainty) — combined with
+    `all` / `any` + patience.
+- `MSMSpawner` (`spawn_scheme: msm`) — **uncertainty × leverage × flux**
+  microstate seeding (`π_i · |ψ_i| · σ_out,i + α/√c_i`) on the estimator's shared
+  clustering, throwing runs at the transitions whose in/out rates are uncertain
+  and important; least-counts fallback before the first MSM. Knobs
+  `msm.spawn_alpha` / `spawn_leverage` / `spawn_uncertainty`; `msm.stable_clustering`
+  keeps microstate IDs comparable across iterations.
+- **Weighted-ensemble resampling** — a real `WeightedEnsemble` split/merge core
+  (Huber & Kim, weight-conserving) and `WESpawner` (`spawn_scheme: we`,
+  `we_target_per_bin`), replacing the former placeholder.
+- `MSMConfig` — all MSM behaviour is **opt-in** (`msm.enabled: false` by
+  default), so existing configs are unaffected.
+- Per-iteration MSM results are written to the run directory and checkpointed
+  for resume.
+
+#### MSM analysis & plotting
+- `autosampler/analysis/` — matplotlib-free data utilities (`load_msm_series`,
+  `load_latest_msm`, `load_cv_points`, free energies, free-energy surface) plus
+  `plots` (implied timescales, VAMP-2 / timescale convergence, free-energy
+  surface, metastable free energies, MSM network) and a one-command
+  `autosampler-analyze` CLI producing a multi-panel convergence report.
+- `msm.npz` now also stores the implied-timescale sweep and metastable
+  populations for plotting.
+
+#### Extensible ML collective variables (Phase 1)
+- `autosampler/spaces/registry.py` — single source of truth for CV methods,
+  their backends, and availability. Adds **VAMPNet** and **SPIB** (State
+  Predictive Information Bottleneck) alongside TICA, TVAE, PCA, deep-TICA, and
+  deep-LDA, with actionable errors when an optional backend is missing.
+
+#### HPC execution backends (Phase 3)
+- `autosampler/execution/` — pluggable `ExecutionBackend` (factory pattern):
+  - `local` — multiprocessing across CPU/GPU slots on one node (multi-GPU
+    workstation); preserves the original GPU-slot scheduling.
+  - `slurm` / `pbs` — one **array job per iteration**, with completion driven by
+    filesystem result markers and automatic **resubmission** of failed walkers
+    (`execution.max_retries`).
+- `ExecutionConfig` — backend selection plus scheduler resources (partition /
+  queue, account, walltime, cpus/gpus per task, memory, module loads, extra
+  directives). Defaults to `local`.
+
+#### VAMP-2 input-feature selection
+- `autosampler/spaces/feature_selection.py` — `vamp2_score`, `rank_candidates`,
+  `greedy_vamp_selection`, and `FeatureSelector`: choose and **adaptively
+  update** the input features that best resolve the slow dynamics.
+- `FeatureSelectionConfig` (`feature_selection.enabled`, opt-in) — re-selects
+  feature columns every `cadence` iterations; selection persisted for resume.
+
+#### Landscape-adaptive binning
+- `autosampler/binning/adaptive.py` — `AdaptiveBinner` + `BinnerFactory` with
+  `gradient` (equi-resistance: fine bins where the density is low / barriers),
+  `mab` (Minimal-Adaptive-Binning style front footholds), and `eigenvector` (bin
+  along the leading slow CV coordinate) schemes alongside `uniform`. Selected via
+  `binning.scheme`; wired into the density and weighted-ensemble spawners; opt-in,
+  default `uniform` reproduces the constant-width grid exactly.
+
+#### Adaptive CV quality & reproducibility (Phase 4)
+- **VAMP-2-driven adaptive retraining** (`retrain_policy: vamp_adaptive`):
+  `RetrainController` retrains the CV only when its VAMP-2 score on fresh data
+  drops by more than `vamp_retrain_tol` below its reference (with
+  `retrain_min_interval` / `retrain_max_interval` bounds). Reference score is
+  checkpointed. The default `fixed` policy preserves the legacy schedule.
+- **Reproducibility:** `SeedManager` now also seeds PyTorch Lightning
+  (`seed_everything`, used by deep-TICA/LDA) and documents determinism limits.
+- **Feature-type selection:** `feature_selection.candidate_feature_types` ranks
+  whole feature types (`distances` / `fitted_coords` / `phi_psi`) by VAMP-2 and
+  switches the loop to the best one (re-running column selection on a change).
+
+#### End-user input file & tutorial
+- **`autosampler-init`** writes a fully-annotated starter input file
+  (`autosampler/templates.py`, mirrored to `examples/template.yaml`) covering
+  every section, method choice, and hyperparameter. Documented in
+  `docs/input_file.md`.
+- **Jupyter notebook tutorial** with rendered plots
+  (`examples/notebooks/adaptive_msm_tutorial.ipynb`): input file, VAMP-2 feature
+  selection, MSM estimation, convergence, weighted ensemble, and the analysis
+  report — all on fast synthetic data.
+
+#### Tooling & docs
+- Test suite (`pytest`) covering MSM, CV methods, execution backends, feature
+  selection, config, spawners, and hardening — **49 tests**.
+- GitHub Actions CI (Python 3.10 / 3.11) running ruff + pytest.
+- `ruff` / `black` / `isort` config, `.pre-commit-config.yaml`, `CONTRIBUTING.md`.
+- `docs/` site (MkDocs) and tutorials; `CHANGELOG.md`; example run scripts for
+  local, SLURM, and PBS.
+
+### Changed (Phase 2 — engineering foundation)
+- Refactored the per-iteration UI out of `core.py` into
+  `autosampler/reporting.py` (`IterationReporter`); de-duplicated project-file
+  CV loading.
+- Migrated configuration to **Pydantic v2** (`field_validator` /
+  `model_validator`, `model_dump()`); pinned `pydantic>=2.0`.
+- Added **checkpoint format versioning** and generalised torch-encoder
+  snapshots to tvae / vampnet / spib.
+- Made `autosampler.spaces` import lazily so lightweight modules (e.g. the CV
+  registry) import without MDAnalysis / torch.
+
+### Fixed (Phase 2)
+- Renamed `AdaptiveSpaceModel.fited` → `fitted` (with a backwards-compatible
+  loader for old checkpoints).
+- Added MD subprocess timeouts (`AUTOSAMPLER_MD_TIMEOUT`) for GROMACS / Amber.
+- Validate trajectory files exist and are non-empty before CV extraction.
+- Replaced hardcoded `/tmp` with `tempfile.gettempdir()`; narrowed an
+  over-broad exception handler.
+- Removed the dead `WEResampler` stub.
+
+## [2.0.0] — baseline
+
+Modular adaptive sampling framework: OpenMM / GROMACS / Amber engines, fixed or
+learned (TICA / TVAE / PCA / deep-TICA) CV spaces, density / Voronoi / LOF / FPS
+spawners, bin-occupancy convergence, checkpoint/restart, and lineage-aware path
+reconstruction.

CONTRIBUTING.md

@@ -0,0 +1,58 @@
+# Contributing to AutoSampler
+
+Thanks for your interest in improving AutoSampler. This guide covers the
+development workflow, coding standards, and how to add new components.
+
+## Development setup
+
+```bash
+conda env create -f env.yml
+conda activate autosampler
+pip install -e ".[all]"        # runtime + deep-tica + examples + test extras
+pip install pre-commit && pre-commit install
+```
+
+For a lightweight setup that runs the MSM and CV-method tests without the heavy
+MD backends (OpenMM, MDAnalysis):
+
+```bash
+pip install numpy scipy scikit-learn pydantic pyyaml deeptime pytest torch
+```
+
+## Running tests and linters
+
+```bash
+pytest -q                 # full test suite (MD-dependent tests self-skip)
+ruff check .              # lint
+ruff format .             # format (or `black .`)
+```
+
+CI runs the test suite on Python 3.10 and 3.11 (see `.github/workflows/ci.yml`).
+Please add or update tests for any behaviour change; aim to keep new modules
+covered.
+
+## Coding standards
+
+- Target Python 3.10+. Use type hints on public functions.
+- Formatting and linting are enforced by ruff/black (line length 88).
+- Keep presentation/IO out of core logic (see `autosampler/reporting.py`).
+- Prefer the existing factory/registry patterns when adding components.
+
+## Adding components
+
+AutoSampler is built around small registries so new methods slot in cleanly:
+
+- **MD engine** — subclass `MDEngine` and call `EngineFactory.register(...)`
+  in `autosampler/engines/`.
+- **Spawner** — subclass `Spawner` and call `SpawnerFactory.register(...)`
+  in `autosampler/spawners/`.
+- **CV method** — add a `CVMethod` to `autosampler/spaces/registry.py` and a
+  branch in `AdaptiveSpaceModel.fit` / `.project`.
+- **MSM convergence criterion** — subclass `ConvergenceCriterion` and register
+  it in `autosampler/msm/convergence.py`.
+
+## Commit / PR guidelines
+
+- Write focused commits with descriptive messages.
+- Ensure `pytest` and `ruff check` pass locally before opening a PR.
+- Describe the scientific or engineering motivation in the PR description.

README.md

@@ -3,12 +3,39 @@
 AutoSampler 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 with checkpointed provenance.
+cycle — continuing until a **Markov State Model (MSM)** built on the sampled
+data has **converged**.
 
 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.
 
+## What's new
+
+- **MSM-convergence engine** — build an MSM each iteration (deeptime) and stop
+  automatically on implied-timescale / VAMP-2 convergence *and* a flux-weighted
+  **transition-matrix** statistical-error gate (`spawn_scheme: msm`, `msm.enabled`).
+- **Uncertainty-guided spawning** — seed walkers by **uncertainty × leverage ×
+  flux** to drive the MSM toward convergence fastest (`msm.spawn_uncertainty`).
+- **Landscape-adaptive binning** — place bins finer across barriers and coarser
+  in basins (`binning.scheme: gradient | mab | eigenvector`), recomputed each
+  iteration; defaults to the uniform grid.
+- **More learned CVs** — VAMPNet and SPIB alongside TICA / TVAE / PCA / deep-TICA.
+- **VAMP-2 feature selection** — optionally select and adaptively update the
+  input features that best resolve the slow dynamics (`feature_selection.enabled`).
+- **HPC scalability** — run on a multi-GPU workstation or dispatch walkers as
+  **SLURM** / **PBS** array jobs (`execution.backend`).
+
+Get started in one command — `autosampler-init` writes an annotated input file
+covering every method, feature, and hyperparameter (see
+[`docs/input_file.md`](docs/input_file.md)). A runnable
+[notebook tutorial](examples/notebooks/adaptive_msm_tutorial.ipynb) with rendered
+plots walks through the whole workflow.
+
+See **[`docs/`](docs/index.md)** (full documentation & tutorials) and
+**[`CHANGELOG.md`](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.
@@ -205,10 +232,44 @@ AutoSampler currently supports:
 - `voronoi`: KMeans-backed Voronoi cells with exact clipped polygon areas.
 - `lof`: local-outlier-factor based frame selection.
 - `fps`: farthest-point sampling for geometric spread.
+- `msm`: MSM least-counts — restart from sparsely-sampled microstates to
+  directly reduce the statistical error of the Markov State Model.
 
 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
+(`autosampler/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 |
+| `vampnet`    | VAMPNet | deeptime + torch | trained with the VAMP-2 score |
+| `spib`       | State Predictive Information Bottleneck | built-in (torch) | Wang & Tiwary 2021 |
+| `deep-tica`  | Deep (nonlinear) TICA | mlcolvar (optional) | `pip install "autosampler[deep-tica]"` |
+| `deep-lda`   | Deep LDA (supervised) | mlcolvar (optional) | needs state labels |
+
+`vampnet` and `spib` work out of the box (only deeptime/torch). Optional methods
+raise a clear, actionable error if their backend is missing.
+
+## MSM-Based Convergence
+
+With `msm.enabled: true`, AutoSampler builds a Markov State Model over the CV
+space each iteration (clustering → transition counts → MLE/Bayesian MSM →
+implied timescales, VAMP-2 score, PCCA+ metastable states) and stops sampling
+once the MSM has **converged**. Convergence is decided by a composable
+`ConvergenceMonitor` with pluggable criteria — implied-timescale stability,
+VAMP-2 plateau, stationary-distribution drift, and Bayesian statistical error —
+combined with `all`/`any` and a patience window. Per-iteration results are
+written to `iter_*/msm.npz`. See `examples/AIB9/config_msm_vampnet.yaml` for a
+complete MSM-driven adaptive-sampling configuration.
+
 ## Typical Workflow
 
 1. **Define the scientific question.**

autosampler/analysis/__init__.py

@@ -0,0 +1,9 @@
+"""Post-hoc MSM analysis and plotting utilities.
+
+``data`` holds matplotlib-free numerics (loading msm.npz/cvs.npz, free
+energies); ``plots`` holds the matplotlib visualisations.
+"""
+
+from . import data
+
+__all__ = ["data"]