AutoSampler: MSM-convergence engine, HPC scaling, feature optimisation, analysis & UX#1
Merged
Conversation
Introduce a real Markov State Model subsystem so adaptive sampling can run until the MSM converges, plus a registry of cutting-edge collective-variable methods. All new behaviour is opt-in; existing configs and examples are unaffected. MSM subsystem (autosampler/msm/): - MSMEstimator: deeptime-based clustering -> transition counts -> MLE/Bayesian MSM -> implied timescales, VAMP-2 score, PCCA+ metastable states. - ConvergenceMonitor with pluggable criteria (implied-timescale stability, VAMP-2 plateau, stationary-distribution drift, Bayesian statistical error) combined with all/any and a patience window. - MSMResult/ITSResult serialisable diagnostics; per-iteration iter_*/msm.npz. - Wired into AutoSamplerCore (opt-in via msm.enabled), with checkpoint save/restore of monitor state. MSM-guided spawning: - New `msm` spawner: microstate least-counts to target MSM statistical error. CV-method registry (autosampler/spaces/registry.py): - Single source of truth for adaptive CV methods + backend availability; centralises the previously hard-coded adaptive-mode lists in core.py. - New methods: VAMPNet (deeptime) and SPIB (built-in torch implementation), both working out of the box; deep-tica/deep-lda gated on optional mlcolvar with clear, actionable errors. Config: - MSMConfig (opt-in) and space_mode validation; SPIB hyperparameters. Tests + CI: - 25 tests covering MSM recovery on a toy 3-state system, convergence logic, the MSM spawner, the registry, and VAMPNet/SPIB training; core integration tests for the MSM wiring. Minimal GitHub Actions workflow. Docs: - README sections on CV methods and MSM-based convergence; new examples/AIB9/config_msm_vampnet.yaml. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Engineering-foundation pass to stabilise the codebase before deeper work. Correctness & hardening: - Rename AdaptiveSpaceModel.fited -> fitted with a backwards-compatible __setstate__ shim so existing pickled checkpoints still load. - Add MD subprocess timeouts (AUTOSAMPLER_MD_TIMEOUT) with TimeoutExpired handling in the GROMACS and Amber engines. - Validate trajectory files exist and are non-empty before CV extraction. - Replace hardcoded /tmp with tempfile.gettempdir() in the CLI. - Narrow an over-broad except in the physical-CV evaluation path. - Add checkpoint format versioning (format_version file + load-time check) and generalise the torch-encoder snapshot to tvae/vampnet/spib. - Remove the dead WEResampler pass-through stub and its wiring. Refactor (core.py): - Extract the ANSI/tabular per-iteration UI into autosampler/reporting.py (IterationReporter). - De-duplicate the project-file CV loading into _extract_physical_cvs. Pydantic v2 modernisation: - Migrate validators to field_validator/model_validator, drop the v1/v2 shims, switch .dict() -> model_dump(), and pin pydantic>=2.0. Quality gates: - ruff/black/isort config, .pre-commit-config.yaml, CONTRIBUTING.md, and a ruff lint step in CI. Tests: - New tests/test_hardening.py (reporter, timeout helper, checkpoint version, fited->fitted shim, trajectory validation, WEResampler removal). 31 tests pass; ruff clean on new/touched modules. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
CI was failing on every run because collecting tests/test_cv_methods.py imported autosampler.spaces.registry, which triggered the spaces package __init__ to eagerly import features -> MDAnalysis. The lightweight CI stack intentionally omits MDAnalysis, so collection errored and all jobs failed. - Make autosampler/spaces/__init__.py import lazily (PEP 562 __getattr__) so registry and other light modules import without MDAnalysis/torch/deeptime. The public API (FeatureExtractor, AdaptiveSpaceModel, ...) still resolves on first access. - Narrow CI triggers to push on [main, devel] plus pull_request, eliminating duplicate runs (and notification spam) on every feature-branch push. 31 tests pass; ruff clean. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Adds an execution-backend abstraction so walker MD jobs can run on a multi-GPU workstation or be dispatched to CPU/GPU HPC clusters, with fault tolerance. - autosampler/execution/: ExecutionBackend ABC + factory (mirrors the engine/ spawner factories), WalkerTask payload, and build_walker_tasks/run_walker_task helpers. - LocalProcessBackend: the original ProcessPoolExecutor + dynamic GPU-slot scheduling, moved behind the new interface (behaviour preserved). - SchedulerBackend + SlurmBackend + PBSBackend: one array job per iteration. Completion is driven by filesystem result markers (portable, testable); failed/missing walkers are resubmitted up to execution.max_retries. The job entrypoint is autosampler/execution/run_task.py. - ExecutionConfig (config.py): backend selection + scheduler resources (partition/queue, account, walltime, cpus/gpus per task, memory, module loads, extra directives). Defaults to local, so existing configs are unchanged. - workflows/parallel.py reduced to a thin facade dispatching to the configured backend; core.py passes config.execution through. Tests: tests/test_execution.py drives the full submit -> poll -> collect -> retry state machine with a fake in-process scheduler, plus script rendering, job-id parsing, retry, and make_backend dispatch. 42 tests pass; ruff clean. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Previously the input features for the CV/MSM were fixed for the whole run
(adaptive_feature_type + feature_selection atom mask). This adds an opt-in
protocol to select — and adaptively update — the best input features by VAMP-2,
the variational score for how well a feature set resolves the slow dynamics.
- autosampler/spaces/feature_selection.py:
- vamp2_score(): dependency-light VAMP-2 from time-lagged covariances.
- rank_candidates(): rank named candidate feature sets by VAMP-2.
- greedy_vamp_selection(): greedy forward selection of feature columns/groups
maximising VAMP-2 (parsimony via a min_gain threshold).
- FeatureSelector / FeatureSelection: orchestration + serialisable result.
- FeatureSelectionConfig (config.py): enabled, method, lagtime, cadence,
max_features, dim, min_gain. Off by default.
- core.py: when enabled, (re)selects feature columns every `cadence` iterations
and applies the selection consistently at CV fit, projection, and history
re-projection; selection persisted in the checkpoint state for resume.
Tests: tests/test_feature_selection.py (VAMP-2 ranking, greedy parsimony,
serialisation, config validation). 49 tests pass; ruff clean on new modules.
https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Documents all features implemented so far (Phases 1-3 + VAMP-2 feature selection) and makes the project approachable. - docs/ (MkDocs + material): index, quickstart, concepts, full configuration reference, CV methods, VAMP-2 feature selection, MSM & convergence, execution (workstation/SLURM/PBS), and an end-to-end adaptive-MSM tutorial. mkdocs.yml added; `docs` extra in pyproject. - CHANGELOG.md: detailed, Keep-a-Changelog-style log of the MSM engine, CV methods, execution backends, feature selection, and Phase 2 hardening. - examples/: run_local.sh, slurm_submit.sh, pbs_submit.sh driver scripts and examples/AIB9/config_msm_feature_selection.yaml combining MSM + VAMP-2 feature selection + SLURM dispatch. - README: "What's new" summary and links to docs/changelog. - pyproject: add `docs` extra and ruff to the `test` extra. 49 tests pass. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Tightens the CV<->MSM coupling and improves run reproducibility. - autosampler/spaces/retraining.py — RetrainController decides when to (re)train the CV. The new `vamp_adaptive` policy retrains only when the CV's VAMP-2 score on fresh data drops by more than `vamp_retrain_tol` below its post-training reference (bounded by retrain_min_interval / retrain_max_interval). The default `fixed` policy reproduces the legacy retrain_freq schedule. - core.py: scores the current CV by VAMP-2 (latent projection) each iteration, drives retraining through the controller, and checkpoints its reference score for resume. - config.py: retrain_policy, vamp_retrain_tol, retrain_min_interval, retrain_max_interval (validated; default keeps legacy behaviour). - utils/seeds.py: SeedManager now also seeds PyTorch Lightning (seed_everything) and documents determinism limits. Tests: tests/test_retraining.py covers the fixed/adaptive policies, min/max intervals, state round-trip, config validation, and deterministic seeding. 57 tests pass; ruff clean on new modules. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Completes the feature-selection story: beyond choosing the best column subset, the loop can now rank whole feature *types* by VAMP-2 and use the best one. - FeatureSelectionConfig.candidate_feature_types (validated subset of distances / fitted_coords / phi_psi). Empty -> use adaptive_feature_type. - core.py: _extract_adaptive_features extracts each candidate type, ranks them via rank_candidates (VAMP-2) every `cadence` iterations, and switches to the winner; a type change resets the column selection. Chosen type checkpointed. - Docs (feature_selection / configuration) and CHANGELOG updated; tests for the config validation and type ranking. 59 tests pass; ruff clean on in-scope modules. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Replaces the removed WEResampler placeholder with a correct, weight-conserving implementation. - autosampler/binning/we.py — WeightedEnsemble split/merge core (Huber & Kim): splits high-weight walkers in under-populated bins and merges low-weight walkers in over-populated bins, keeping a target count per bin while total statistical weight is exactly conserved. Pure/array-based and unit-tested. - autosampler/spawners/we.py — WESpawner (spawn_scheme: we): carries per-frame weights across the cumulative cloud, bins, runs WE resampling, and draws the next walkers proportional to weight (splits => repeated indices). Weights tracked across iterations with state_dict/load_state_dict. - config: spawning.we_target_per_bin; core passes it and the seed to the spawner factory. Tests: tests/test_weighted_ensemble.py covers split, merge, multi-bin weight conservation, no-op-at-target, validation, and the spawner. 66 tests pass. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Turns the per-iteration msm.npz diagnostics into figures and a one-command report. - autosampler/analysis/data.py — matplotlib-free numerics: load_msm_series, load_latest_msm, load_cv_points, free_energy_from_populations, free_energy_surface. - autosampler/analysis/plots.py — implied timescales, VAMP-2 and timescale convergence, free-energy surface, metastable free energies, and an MSM network (circular layout, no networkx); plot_convergence_report assembles a multi-panel summary. matplotlib imported lazily with an actionable error. - autosampler-analyze CLI (analysis_cli.py) -> convergence_report.png. - core: msm.npz now also stores the implied-timescale sweep and metastable populations. - docs/analysis.md + nav/quickstart; CHANGELOG updated. Tests: tests/test_analysis.py (data loaders, free energies, plotting smoke, CLI guard). 73 tests pass; ruff clean. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
…utorial End-user-facing UX: a single, fully-documented input file and a runnable tutorial. Input file structure: - autosampler/templates.py (DEFAULT_TEMPLATE) — one annotated YAML covering every section: system, engine, spawning, CV method + hyperparameters, feature_selection, msm, execution, and run-level settings, with the available method choices inline. Mirrored to examples/template.yaml. - autosampler-init CLI writes the starter file (-o/--force); registered entry point. docs/input_file.md explains the structure for end users (linked from nav, README, quickstart). Notebook tutorial: - examples/notebooks/adaptive_msm_tutorial.ipynb — executed with rendered plots (4 figures), covering the input file, VAMP-2 feature selection, MSM estimation, convergence detection, weighted ensemble, and the analysis report, all on fast synthetic data. docs/notebook_tutorial.md added. Docs: mkdocs nav + snippets include of the template; CHANGELOG updated. Tests: tests/test_input_template.py (template parses, matches examples copy, init CLI writes/guards). 76 tests pass; ruff clean. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
namus
changed the title
- AutoSampler_Changelog.pdf: structured 4-page summary of all new features and improvements over the v2.0.0 baseline (headline features, detailed changes by area, and a capability comparison table). - tools/generate_changelog_pdf.py: reportlab generator so the PDF can be regenerated (`pip install reportlab && python tools/generate_changelog_pdf.py`). https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Ties MSM convergence diagnosis and adaptive spawning to one coherent, error-aware model on a shared clustering. All opt-in; defaults reproduce the prior least-counts + spectral behaviour. - msm/diagnostics.py: MSMResult gains count_matrix, eigenvectors (slow right eigenvectors), and state_symbols (cluster id -> active-state map). - msm/estimator.py: populate those from the connected count model / MSM; add stable_clustering (seed k-means from previous centres so T_ij and microstate IDs stay comparable across iterations). - msm/convergence.py: new TransitionMatrixCriterion — converged when the largest flux-weighted relative Dirichlet uncertainty of the transition probabilities (sigma(T_ij)/T_ij, sigma^2 = T(1-T)/(c_i+1)) is below tol; min_flux ignores negligible transitions. Augments the spectral criteria (require both via mode: all). Registered + exported. - spawners/msm.py: MSM-guided scoring pi_i * |psi_i| * (sigma_out/mean) + alpha/sqrt(c_i) on the estimator's shared clustering (set by core each iteration); frontier / disconnected microstates get the exploration weight; least-counts fallback when no MSM yet. Knobs alpha / leverage / uncertainty. - config.py: MSMConfig.stable_clustering, spawn_alpha, spawn_leverage, spawn_uncertainty. - core.py: pass stable_clustering to the estimator and the spawner knobs to the factory; set spawner.msm_result + cluster_model (previous iteration's, mutually consistent) before sampling; persist count_matrix/eigenvectors to iter_*/msm.npz. - Docs (msm/configuration), CHANGELOG, and the annotated template updated. Tests: tests/test_msm_tmatrix.py (criterion well/sparse/flux-mask/None; spawner concentration + least-counts fallback; estimator field population + stable clustering). 85 tests pass; ruff clean. https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Bins for the density / weighted-ensemble spawners were a uniform grid (constant
width everywhere) — suboptimal at barriers (wide bins prevent WE flux across
steep regions) and wasteful in flat basins. This makes bins landscape-adaptive,
recomputed each iteration. Opt-in; default `uniform` reproduces RegularBinner.
- autosampler/binning/adaptive.py: AdaptiveBinner + BinnerFactory + make_binner.
- gradient: equi-resistance edges (∫exp(βF) ∝ ∫1/P) — fine bins where density
is low (barriers / steep regions), coarse in basins.
- mab: Minimal-Adaptive-Binning-style front footholds.
- eigenvector: uniform bins along the leading slow CV coordinate only (committor
proxy; fine across the barrier, coarse in basins; clean for many CVs).
All produce a RegularBinner-compatible BinTable via a shared variable-width grid.
- config.py: BinningConfig (scheme, n_fine, smoothing); added to AutoSamplerConfig.
- spawners/density.py, spawners/we.py: use an injected `binner` when set (synced
with resolution bumps), else the uniform grid. WE per-frame weights need no
remap across re-binning.
- core.py: build the binner from config.binning and attach it to the spawner.
- Docs (binning.md + nav + configuration), CHANGELOG, annotated template.
Tests: tests/test_adaptive_binning.py (gradient finer at the barrier, uniform ==
RegularBinner, BinTable consistency for all schemes, eigenvector 1-D, config
validation, density/WE spawner integration). 96 tests pass; ruff clean.
https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
…adaptive binning Bring the user-facing material in line with the two new features (transition-matrix convergence + uncertainty-guided spawning, and landscape-adaptive binning): - CHANGELOG.md: highlight both features in the cycle summary. - README.md: add "what's new" bullets for the T_ij convergence gate, uncertainty x leverage x flux spawning, and adaptive binning. - docs/index.md: add binning to "why" + nav. - docs/tutorials/adaptive_msm.md: show the transition_matrix criterion, stable_clustering / spawn_uncertainty, and an adaptive-binning block. - examples/AIB9/config_msm_vampnet.yaml: demonstrate the transition_matrix criterion, uncertainty-guided spawning knobs, and a gradient binning block. - tools/generate_changelog_pdf.py: add a "Landscape-adaptive binning" headline section, refresh the MSM bullets, add a Binning comparison row and a "new configuration options" reference table; bump test count 76 -> 96. - AutoSampler_Changelog.pdf: regenerated structured PDF (6 pages). https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
dmighty007
reviewed
Jun 30, 2026
dmighty007
left a comment
dmighty007
left a comment
Contributor
There was a problem hiding this comment.
Working, storage was overflowing since checkpoint writter was writing historical information. MSM workflow(efficiency and spawing) needs to be tested.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Integration branch
devel→main. This PR collects the full development effort that turns AutoSampler from a coverage-driven sampler into an autonomous, MSM-convergence-driven framework, with HPC scalability, VAMP-2 feature optimisation, analysis tooling, and an end-user input-file workflow. All new behaviour is opt-in; existing configs are unaffected.Status: 76 tests pass; ruff clean on all new/touched modules; CI green on Python 3.10 & 3.11.
Phase 1 — MSM convergence engine
autosampler/msm/:MSMEstimator(deeptime) — clustering → transition counts → MLE/Bayesian MSM → implied timescales, VAMP-2, PCCA+, stationary distribution; serialisablediagnostics;ConvergenceMonitorwith composable criteria (ITS stability, VAMP-2 plateau, stationary-distribution drift, statistical error).MSMSpawner(spawn_scheme: msm, least-counts/uncertainty seeding).spaces/registry.py— extensible CV methods, adding VAMPNet and SPIB beside TICA/TVAE/PCA/deep-TICA/deep-LDA.MSMConfig; per-iterationmsm.npz+ checkpointing.Phase 2 — Engineering foundation
fited→fitted(back-compatible checkpoint loader); MD subprocess timeouts (AUTOSAMPLER_MD_TIMEOUT); trajectory-file validation;tempfile.gettempdir(); narrowed excepts; checkpoint format versioning; removed deadWEResamplerstub.core.pyrefactor:IterationReporter, de-duplicated project-file CV loading.CONTRIBUTING.md, CI (lint + tests); lazyspacesimports.Phase 3 — Execution backends (workstation & HPC)
autosampler/execution/:ExecutionBackendfactory —local(multi-GPU),slurm,pbs(one array job per iteration, marker-driven completion, automatic resubmission).ExecutionConfig.Phase 4 — Adaptive CV quality
retrain_policy: vamp_adaptive).SeedManageralso seeds Lightning.Feature selection & optimisation (VAMP-2)
vamp2_score,rank_candidates, greedy column optimisation, and best feature-type selection — select and adaptively update the input features that best resolve the slow dynamics (feature_selection.*).Weighted ensemble
WeightedEnsemblesplit/merge (weight-conserving) +WESpawner(spawn_scheme: we), replacing the former placeholder.Analysis & plotting
autosampler/analysis/(data utils + plots: implied timescales, VAMP-2/timescale convergence, free-energy surface, metastable free energies, MSM network) and theautosampler-analyzeCLI producing a multi-panel report.msm.npzextended with the ITS sweep + metastable populations.End-user UX: input file, docs & tutorials
autosampler-initwrites a single, fully-annotated input file (autosampler/templates.py↔examples/template.yaml) exposing every method, feature, and hyperparameter;docs/input_file.mdexplains it.docs/): quickstart, concepts, configuration reference, CV methods, feature selection, MSM & convergence, analysis, execution, and tutorials;CHANGELOG.md; local/SLURM/PBS run scripts.examples/notebooks/adaptive_msm_tutorial.ipynb).Notes / follow-ups
core.py::run_iterationinto stage classes was intentionally left incremental.https://claude.ai/code/session_01MjnMX5xKm184y4sLXhq4pT
Generated by Claude Code