Production & publication readiness: hardening, reproducibility, packaging, examples, docs, tests

.github/workflows/ci.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12"]
     steps:
       - uses: actions/checkout@v4
 
@@ -27,11 +27,29 @@ jobs:
           # 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 numpy scipy scikit-learn shapely 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: Lint (ruff, whole tree)
+        run: ruff check autosampler tests
 
       - name: Run test suite
         run: pytest -q
+
+  docs:
+    name: Docs build
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install docs dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install mkdocs-material "mkdocstrings[python]"
+          # Make the package importable for mkdocstrings/griffe without pulling
+          # the heavy runtime deps (griffe analyses the source statically).
+          pip install -e . --no-deps
+      - name: Build site
+        run: mkdocs build

.github/workflows/release.yml

@@ -0,0 +1,46 @@
+name: Release
+
+# Publishes to PyPI when a version tag (v*) is pushed.
+#
+# Uses PyPI Trusted Publishing (OIDC) — no API token/secret required. Configure a
+# trusted publisher for this repository + workflow on PyPI before the first
+# release: https://docs.pypi.org/trusted-publishers/
+on:
+  push:
+    tags: ["v*"]
+
+jobs:
+  build:
+    name: Build distributions
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Build sdist and wheel
+        run: |
+          python -m pip install --upgrade pip build
+          python -m build
+      - name: Check metadata
+        run: |
+          pip install twine
+          twine check dist/*
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    name: Publish to PyPI
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write  # required for trusted publishing
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

CHANGELOG.md

@@ -122,6 +122,26 @@ adds a flux-weighted **transition-matrix convergence** gate with
 - Made `autosampler.spaces` import lazily so lightweight modules (e.g. the CV
   registry) import without MDAnalysis / torch.
 
+### Production-readiness hardening (pre-release pass)
+- **Reproducibility:** threaded the configured seed into all learned-CV training
+  (SPIB no longer hardcodes `seed=42`; the torch RNG is reseeded before every
+  `fit`), and `SeedManager` now requests deterministic torch algorithms.
+- **Robustness:** the local backend tolerates a single walker's failure (and adds
+  an opt-in `execution.walker_timeout` hang guard) instead of aborting the batch;
+  delta-checkpoint resume reconstructs the full history (fixing a truncated
+  `autosampler-path`), writes atomically, and tolerates a corrupt delta; fixed a
+  target-mode spawn crash and a deep-TICA device mismatch.
+- **Packaging:** OpenMM is now an optional, lazily-imported backend so the base
+  `pip install` resolves; full PyPI metadata + single-sourced version;
+  `CITATION.cff`; tag-driven PyPI release workflow.
+- **CI/quality:** lint the whole tree (was a hand-picked subset), test on Python
+  3.10–3.12, build the docs in CI.
+- **Docs/examples/tests:** a self-contained CPU-only alanine-dipeptide
+  hello-world; example configs for SPIB / deep-TICA / WE / target / PBS; an
+  examples index; an API reference, references/citations page, and full CLI
+  reference; +23 tests (delta checkpoint, reproducibility, timeout, spawners,
+  paths). Suite 95 → 118.
+
 ### Fixed (Phase 2)
 - Renamed `AdaptiveSpaceModel.fited` → `fitted` (with a backwards-compatible
   loader for old checkpoints).

CITATION.cff

@@ -0,0 +1,36 @@
+cff-version: 1.2.0
+message: "If you use AutoSampler in your research, please cite it as below."
+title: "AutoSampler: an MSM-convergence-driven adaptive sampling framework for molecular dynamics"
+abstract: >-
+  AutoSampler is a modular framework for autonomous adaptive molecular-dynamics
+  sampling. It runs short MD walkers, projects frames into fixed or machine-learned
+  collective-variable spaces, restarts walkers from informative regions, builds a
+  Markov State Model each iteration, and stops when the MSM has converged.
+type: software
+# TODO(release): confirm the full author list, ORCIDs, and affiliations before
+# submitting the publication. The entries below are drawn from the repository and
+# must be verified.
+authors:
+  - family-names: Maity
+    given-names: Dibyendu
+    email: dibyendumaity1999@bose.res.in
+    # orcid: "https://orcid.org/0000-0000-0000-0000"
+    # affiliation: "S. N. Bose National Centre for Basic Sciences"
+  - family-names: Chakrabarty
+    given-names: Suman
+    email: chakrabarty.suman@gmail.com
+    # orcid: "https://orcid.org/0000-0000-0000-0000"
+    # affiliation: "S. N. Bose National Centre for Basic Sciences"
+version: 2.0.0
+license: MIT
+repository-code: "https://github.com/TeamSuman/AutoSampler"
+url: "https://github.com/TeamSuman/AutoSampler"
+keywords:
+  - molecular dynamics
+  - adaptive sampling
+  - Markov state model
+  - collective variables
+  - enhanced sampling
+# TODO(release): add the DOI once a Zenodo archive / preprint is minted, and add a
+# `preferred-citation:` block pointing to the journal article when published.
+# doi: "10.5281/zenodo.XXXXXXX"

MANIFEST.in

@@ -0,0 +1,3 @@
+include README.md CHANGELOG.md LICENSE CITATION.cff env.yml
+recursive-include docs *.md
+exclude AutoSampler_Changelog.pdf AutoSampler_Review.pdf

README.md

@@ -350,3 +350,9 @@ 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.
+
+## How to cite
+
+If you use AutoSampler in your research, please cite it. Citation metadata is in
+[`CITATION.cff`](CITATION.cff) (GitHub renders a "Cite this repository" button
+from it). A DOI and the accompanying publication will be added on release.

autosampler/binning/spatial.py

@@ -372,7 +372,7 @@ def _finite_voronoi_regions_2d(
         radius = float(np.ptp(vor.points, axis=0).max() * 2.0)
 
     all_ridges: dict[int, list[tuple[int, int, int]]] = {}
-    for (p1, p2), (v1, v2) in zip(vor.ridge_points, vor.ridge_vertices):
+    for (p1, p2), (v1, v2) in zip(vor.ridge_points, vor.ridge_vertices, strict=False):
         all_ridges.setdefault(p1, []).append((p2, v1, v2))
         all_ridges.setdefault(p2, []).append((p1, v1, v2))
 
@@ -405,7 +405,7 @@ def _finite_voronoi_regions_2d(
         vs = np.asarray([new_vertices[vertex] for vertex in new_region])
         centroid = vs.mean(axis=0)
         angles = np.arctan2(vs[:, 1] - centroid[1], vs[:, 0] - centroid[0])
-        new_region = [vertex for _, vertex in sorted(zip(angles, new_region))]
+        new_region = [vertex for _, vertex in sorted(zip(angles, new_region, strict=False))]
         new_regions.append(new_region)
 
     return new_regions, np.asarray(new_vertices, dtype=float)

autosampler/checkpoints/manager.py

@@ -1,8 +1,10 @@
 import logging
+import os
 import pickle
-import torch
 from pathlib import Path
-from typing import Any, Dict, Tuple
+from typing import Any
+
+import torch
 
 # On-disk checkpoint format version. Bump when the layout changes; ``load``
 # tolerates older checkpoints (a missing version file is treated as v1).
@@ -13,6 +15,50 @@
 _TORCH_ENCODER_MODES = ("tvae", "vampnet", "spib")
 
 
+def _atomic_pickle(obj: Any, path: Path) -> None:
+    """Pickle ``obj`` to ``path`` atomically (write tmp, then os.replace).
+
+    Prevents a crash mid-write (e.g. an HPC walltime kill) from leaving a
+    truncated file — important for delta history, where every checkpoint's
+    ``history.pkl`` is needed to reconstruct the full history on resume.
+    """
+    tmp = path.with_name(path.name + ".tmp")
+    with open(tmp, "wb") as handle:
+        pickle.dump(obj, handle)
+    os.replace(tmp, path)
+
+
+def reconstruct_history(checkpoint_root: Path, iteration: int) -> dict[Any, Any]:
+    """Merge the per-checkpoint delta ``history.pkl`` files (for all iterations
+    ``<= iteration``) into the full cumulative history.
+
+    Each key normally lives in exactly one delta; on overlap the newer checkpoint
+    wins. Unreadable deltas (truncated by a crash) are skipped with a warning
+    rather than aborting the whole restore.
+    """
+    iters = sorted(
+        int(path.name.removeprefix("iter_"))
+        for path in checkpoint_root.glob("iter_*")
+        if path.is_dir()
+        and path.name.removeprefix("iter_").isdigit()
+        and int(path.name.removeprefix("iter_")) <= iteration
+    )
+    full: dict[Any, Any] = {}
+    for it in iters:
+        hist_file = checkpoint_root / f"iter_{it}" / "history.pkl"
+        if not hist_file.exists():
+            continue
+        try:
+            with open(hist_file, "rb") as handle:
+                part = pickle.load(handle)
+        except Exception as exc:  # noqa: BLE001 - tolerate a corrupt delta
+            logging.warning("Skipping unreadable history delta %s: %s", hist_file, exc)
+            continue
+        if isinstance(part, dict):
+            full.update(part)
+    return full
+
+
 class CheckpointManager:
     """Handles serialization and reconstruction of the sampler state to allow exact deterministic restarts."""
 
@@ -25,56 +71,59 @@ def save(
         iteration: int,
         space_model: Any,
         scaler: Any,
-        bin_state: Dict[str, Any],
-        history: Dict[str, Any],
-        sampler_state: Dict[str, Any] | None = None,
+        bin_state: dict[str, Any],
+        history: dict[str, Any],
+        sampler_state: dict[str, Any] | None = None,
     ) -> None:
         """Save a complete state snapshot."""
         iter_dir = self.checkpoint_dir / f"iter_{iteration}"
         iter_dir.mkdir(exist_ok=True)
-        (iter_dir / "format_version").write_text(str(CHECKPOINT_FORMAT_VERSION))
 
         # 1. Save Space Model (TVAE or TICA)
         if space_model is not None:
-            with open(iter_dir / "space_model.pkl", "wb") as f:
-                pickle.dump(space_model, f)
+            _atomic_pickle(space_model, iter_dir / "space_model.pkl")
 
         if hasattr(space_model, "type"):
             if (
                 space_model.type in _TORCH_ENCODER_MODES
                 and getattr(space_model, "fitted", None) is not None
             ):
-                torch.save(space_model.fitted.state_dict(), iter_dir / "model.pt")
+                tmp = iter_dir / "model.pt.tmp"
+                torch.save(space_model.fitted.state_dict(), tmp)
+                os.replace(tmp, iter_dir / "model.pt")
             elif (
                 space_model.type == "tica"
                 and getattr(space_model, "model", None) is not None
             ):
-                with open(iter_dir / "model.pkl", "wb") as f:
-                    pickle.dump(space_model.model, f)
+                _atomic_pickle(space_model.model, iter_dir / "model.pkl")
 
         # 2. Save Feature Scaler
-        with open(iter_dir / "scaler.pkl", "wb") as f:
-            pickle.dump(scaler, f)
-
+        _atomic_pickle(scaler, iter_dir / "scaler.pkl")
+
         # 3. Save Bins & Spawn History
-        with open(iter_dir / "bin_state.pkl", "wb") as f:
-            pickle.dump(bin_state, f)
-
-        # Delta Checkpointing: Only save history since the last checkpoint
+        _atomic_pickle(bin_state, iter_dir / "bin_state.pkl")
+
+        # Delta checkpointing: each file stores only the history since the last
+        # checkpoint. load()/reconstruct_history() merge the deltas back into the
+        # full history. Writes are atomic so a crash can't truncate a delta and
+        # break the chain.
         last_ckpt = self._get_latest_checkpoint_before(iteration)
         delta_history = {
             k: v for k, v in history.items() if k > last_ckpt and k <= iteration
         }
-        with open(iter_dir / "history.pkl", "wb") as f:
-            pickle.dump(delta_history, f)
-
+        _atomic_pickle(delta_history, iter_dir / "history.pkl")
+
         if sampler_state is not None:
-            with open(iter_dir / "sampler_state.pkl", "wb") as f:
-                pickle.dump(sampler_state, f)
+            _atomic_pickle(sampler_state, iter_dir / "sampler_state.pkl")
+
+        # Write the format marker last: its presence signals a complete checkpoint.
+        tmp = iter_dir / "format_version.tmp"
+        tmp.write_text(str(CHECKPOINT_FORMAT_VERSION))
+        os.replace(tmp, iter_dir / "format_version")
 
     def load(
         self, iteration: int, space_model: Any = None
-    ) -> Tuple[Any, Any, Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
+    ) -> tuple[Any, Any, dict[str, Any], dict[str, Any], dict[str, Any]]:
         """Restore the state exactly as it was at the specified iteration."""
         iter_dir = self.checkpoint_dir / f"iter_{iteration}"
         if not iter_dir.exists():
@@ -106,35 +155,12 @@ def load(
         # 2. Load scaler
         with open(iter_dir / "scaler.pkl", "rb") as f:
             scaler = pickle.load(f)
-            
-        # 3. Load bins & history
+
+        # 3. Load bins & reconstruct the full (delta-checkpointed) history.
         with open(iter_dir / "bin_state.pkl", "rb") as f:
             bin_state = pickle.load(f)
-
-        with open(iter_dir / "history.pkl", "rb") as f:
-            history = pickle.load(f)
 
-        # Reconstruct full history for Delta Checkpointing
-        checkpoint_dirs = [
-            path
-            for path in self.checkpoint_dir.glob("iter_*")
-            if path.is_dir() and path.name.removeprefix("iter_").isdigit()
-        ]
-        previous_iters = sorted([
-            int(path.name.removeprefix("iter_"))
-            for path in checkpoint_dirs
-            if int(path.name.removeprefix("iter_")) < iteration
-        ], reverse=True)
-
-        for prev_iter in previous_iters:
-            prev_hist_file = self.checkpoint_dir / f"iter_{prev_iter}" / "history.pkl"
-            if prev_hist_file.exists():
-                with open(prev_hist_file, "rb") as f:
-                    part_hist = pickle.load(f)
-                    if isinstance(part_hist, dict):
-                        for k, v in part_hist.items():
-                            if k not in history:
-                                history[k] = v
+        history = reconstruct_history(self.checkpoint_dir, iteration)
 
         state_path = iter_dir / "sampler_state.pkl"
         if state_path.exists():

autosampler/cli.py

@@ -8,8 +8,9 @@
 import os
 import sys
 import tempfile
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Sequence
+from typing import Any
 
 os.environ.setdefault(
     "MPLCONFIGDIR", os.path.join(tempfile.gettempdir(), "autosampler-matplotlib")