STAMP v0.13.0 release

docs/architecture.md

@@ -50,6 +50,22 @@ commands/<stage>.py   orchestration: selects adapter + backend, drives the stage
 
 <br>
 
+## Checkpointing
+
+`stamp.utils.checkpoint` provides shared checkpointing/resume utilities: `sync_checkpoint(work_dir, fingerprint)` keeps a work directory when its stored `checkpoint.json` matches the current fingerprint (normalised through JSON) and wipes it otherwise; `signatures` and `directory_signatures` fingerprint input files by size and modification time; `atomic_save_npz` and `atomic_write_json` write each finished unit so a partial file never looks complete.
+
+Command | Work directory | Unit checkpointed
+-- | -- | --
+`pick` | `raw/<picker>/` | tomogram picks (`<id>.json`)
+`decoy` | `raw/<method>/` | rejected-surface: surface scores (`scores/<id>.npz`); shifted: surfaces (`surfaces/<id>.npz`); synthetic-noise: volumes and picks
+`classify` | `raw/` | subvolumes, per-tomogram features, per-cluster alignment angles
+`identify` | `raw/{real,decoy}/` | one fit score per (class, candidate)
+`refine` | `<out>/<class>/` | finished half-maps
+
+Random sampling steps (decoy placement) are never checkpointed, only the deterministic work before them, so a resumed run is identical to an uninterrupted one for a given seed.
+
+<br>
+
 ## Adapters & Runners
 STAMP separates what a tool is from how to run it, as a `ToolAdapter` and `Runner` respectively:
 - `ToolAdapter`: builds a `ToolCommand` from `AdapterInputs` and parses a `RunResult` into `AdapterOutput`.

docs/cli.md

@@ -38,6 +38,16 @@ Option | Default | Meaning | Notes
 
 <br>
 
+## Resuming interrupted commands
+
+STAMP's `pick`, `decoy`, `classify`, `identify` and `refine` commands periodically checkpoint finished work throughout processing. Re-running a command with the same inputs/parameters will be detected, and the command will resume from the previous checkpoint. Partial output lives in `raw/` (for `refine`, in the class directory) alongside a `checkpoint.json` fingerprint of everything that affects the results.
+
+If a command's inputs or parameters differ, any old partial output is discarded with a warning. Parameters which don't affect results (e.g. `--n-processes`, plots, `--keep-raw`) and steps applied after checkpointed work (for `pick`: consensus rule, distance threshold, seed and beam angle filter) can change between runs without invalidating a given checkpoint.
+
+`stamp run` re-runs will also reuse any valid partial output within a stage.
+
+<br>
+
 ## `stamp pick`
 
 ```
@@ -64,6 +74,8 @@ Argument/Option | Default | Use
 `--max-beam-angle-deviation` | `None` | Drop particles whose normal deviates from the beam-orthogonal plane by more than this many degrees
 `--vesicle-labels-mrc` | `None` | Path to an MRC file produced by EValuator's `label` command, or a directory to be matched by filename stem
 `--normalise-per-vesicle` | `False` | Normalise `stamp-native` scores per vesicle instead of per tomogram
+`-n, --n-processes` | `1` | Processes to use for the `stamp-native` picker (1 = sequential)
+`--keep-raw` | `False` | Keep the raw intermediate output instead of archiving it to `raw.tar.gz`
 
 <br>
 
@@ -82,20 +94,22 @@ Option | Default | Use
 `-s, --seg-dir` | *n/a (required unless `--method synthetic-noise`)* | Directory containing segmented MRC files
 `-r, --raw-dir` | *n/a (required unless `--method synthetic-noise`)* | Directory of raw tomogram MRCs
 `--picker-params` | `'{}'` | JSON of native-picker parameters (must match the values used for `stamp pick` so decoys are drawn from the same candidate pool)
-`--n-decoys-per-tomogram` | `50` | Number of decoy positions to pick per tomogram
-`--min-distance-from-real-a` | `100.0` | Minimum distance from a real pick in Å for `--method rejected-surface`
+`--n-decoys-per-tomogram` | `None` *(adaptive: matches the real pick count per tomogram)* | Number of decoy positions to pick per tomogram (required for `--method synthetic-noise`)
+`--min-distance-from-real-a` | `100.0` | Minimum distance in Å from a real pick for `--method rejected-surface`, or from the membrane surface for `--method shifted`
 `--min-pick-distance-a` | `60.0` | Minimum distance from a real pick in Å for `--method shifted`
 `--min-shift-a` | `200.0` | Minimum displacement in Å for `--method shifted`
 `--max-shift-a` | `600.0` | Maximum displacement in Å for `--method shifted`
 `--n-synthetic-tomograms` | `3` | Number of volumes to produce for `--method synthetic-noise`
 `--synthetic-shape` | `'200,200,200'` | Volume shape (in `z,y,x`) to use for `--method synthetic-noise`
 `--seed` | `0` | Seed to use for sampling and half-set assignment
-`-backend` | `local` | Backend to use for picking (`local` or `mock`)
+`--backend` | `local` | Backend to use for picking (`local` or `mock`)
 `--plots/--no-plots` | `--plots` | Write decoy position plots
 `--pick-plot-style` | `segmented` | Style of decoy position plots (`segmented` or `none`)
 `--plot-format` | `tiff` | Decoy plot file format (`png` or `jpg` or `tiff` or `svg`)
 `--pick-zstack-movie/--no-pick-zstack-movie` | `--pick-zstack-movie` | Write a per-tomogram movie through Z with picks highlighted
 `--pick-plot-3d/--no-pick-plot-3d` | `--no-pick-plot-3d` | Write a static 3D pick visualisation per tomogram
+`-n, --n-processes` | `1` | Processes to use for per-tomogram decoy generation (1 = sequential)
+`--keep-raw` | `False` | Keep the raw intermediate output instead of archiving it to `raw.tar.gz`
 
 ### `stamp decoy` methods
 `stamp decoy` provides three methods for generating decoy picks:
@@ -137,6 +151,8 @@ Option | Default | Use
 `--n-azimuthal-samples` | `64` | Number of azimuthal sampling points (must be >= `2*(modes+1)`)
 `--plots/--no-plots` | `--plots` | Write embedding and class-average plots
 `--plot-format` | `tiff` | Embedding and class average plot file format (`png` or `jpg` or `tiff` or `svg`)
+`-n, --n-processes` | `1` | Processes to use for per-cluster in-plane alignment, extraction and feature building (1 = sequential)
+`--keep-raw` | `False` | Keep the raw intermediate output instead of archiving it to `raw.tar.gz`
 
 
 <br>
@@ -159,6 +175,7 @@ Option | Default | Use
 `--fetch-missing` | `False` | Fetch model from AlphaFold when a candidate has no local `structure_path`
 `--plots/--no-plots` | `--plots` | Write fit-score heatmap and decoy-control histogram plots
 `--plot-format` | `tiff` | Fit-score heatmap and decoy-control histogram plot file format (`png` or `jpg` or `tiff` or `svg`)
+`-n, --n-processes` | `1` | Processes to use for candidate fitting (1 = sequential)
 
 <br>
 
@@ -181,7 +198,7 @@ Option | Default | Use
 `--mask` | `None` | Path to MRC file to use as mask
 `--iterations` | `5` | Number of refinement iterations to use
 `--backend` | `local` | Backend to use for refinement (`local` or `mock`)
-`--combined-halfset` | `False` | Use a single combined refinement step with an internal split, instead of two independent trees
+`--threshold` | `0.143` | FSC threshold used to read the resolution
 `--plots/--no-plots` | `--plots` | Write FSC curve plots
 `--plot-format` | `tiff` | FSC curve plot file format  (`png` or `jpg` or `tiff` or `svg`)
 

pyproject.toml

@@ -1,13 +1,14 @@
 [project]
 name = "stamp"
-version = "0.12.0"
+version = "0.13.0"
 description = "A workflow for Sub-Tomogram Averaging Membrane Proteins"
 readme = "README.md"
 authors = [
     { name = "Sam Holland", email = "hol_sam@icloud.com" }
 ]
 requires-python = ">=3.14"
 dependencies = [
+    "fast-hdbscan>=0.3.2",
     "loguru>=0.7.3",
     "matplotlib>=3.11.1",
     "mrcfile>=1.5.4",

src/stamp/adapters/native.py

@@ -90,5 +90,9 @@ def _on_error(job: _PickJob, exc: Exception) -> None:
                 log.error(f'{job.tomogram_id}: {exc}')
             else:
                 raise exc
-        run_parallel(jobs, _pick_one, max_workers=n_workers, label='stamp-native', on_success=_on_success, on_error=_on_error)
+        # a tomogram's cache file is written atomically on completion, so its presence means it is done
+        pending = [job for job in jobs if not (inputs.output_directory / f'{job.tomogram_id}.json').is_file()]
+        if len(pending) < len(jobs):
+            log.info(f'Resuming stamp-native: {len(jobs) - len(pending)} tomogram(s) loaded from checkpoint')
+        run_parallel(pending, _pick_one, max_workers=n_workers, label='stamp-native', on_success=_on_success, on_error=_on_error)
         return load_cached_picks(inputs.output_directory, inputs.tomogram_ids)

src/stamp/classify/cluster.py

@@ -5,11 +5,11 @@
 # Import external dependencies
 import numpy as np
 from dataclasses import dataclass
-from sklearn.cluster import HDBSCAN, KMeans
+from sklearn.cluster import KMeans
 from sklearn.decomposition import PCA
 
 # Import internal STAMP objects
-from stamp.utils.log import log
+from stamp.utils.log import log, log_config
 
 # NOISE_CLUSTER_ID: id for noise cluster
 NOISE_CLUSTER_ID = 'noise'
@@ -40,48 +40,51 @@ def _cluster_embedding(embedding: np.ndarray, config: ClusteringConfig) -> tuple
     if config.method == 'hdbscan':
         if embedding.shape[0] < config.min_cluster_size:
             raise ValueError(f'only {embedding.shape[0]} particles but min_cluster_size is {config.min_cluster_size}')
-        labels = HDBSCAN(min_cluster_size=config.min_cluster_size).fit_predict(embedding)
+        from fast_hdbscan import HDBSCAN
+        labels = HDBSCAN(min_cluster_size=config.min_cluster_size, min_samples=max(config.min_cluster_size - 1, 1)).fit_predict(np.asarray(embedding, dtype=np.float64))
     else:
         n_clusters = min(config.n_clusters, embedding.shape[0])
         labels = KMeans(n_clusters=n_clusters, random_state=config.random_state, n_init=10).fit_predict(embedding)
     centroids = {int(label): embedding[labels == label].mean(axis=0) for label in np.unique(labels) if label != -1}
     return labels, centroids
 
-# reduce_and_cluster: PCA followed by HDBSCAN or KMeans
-def reduce_and_cluster(features: np.ndarray, config: ClusteringConfig) -> ClusteringResult:
+# reduce_features: PCA of a feature matrix, return (embedding, explained variance ratio)
+def reduce_features(features: np.ndarray, config: ClusteringConfig) -> tuple[np.ndarray, np.ndarray]:
     if features.shape[0] == 0:
         raise ValueError('cannot cluster an empty feature matrix')
     n_components = min(config.n_components, features.shape[0] - 1, features.shape[1])
     if n_components < 2:
         raise ValueError(f'Too few particles ({features.shape[0]}) to run PCA (at least 2 needed)')
-
+    log.progress(f'Starting PCA: {features.shape[0]} particles, {features.shape[1]} features, {n_components} components')
     pca = PCA(n_components=n_components, random_state=config.random_state)
     embedding = pca.fit_transform(features)
-    log.debug(f'PCA: {n_components} components, {pca.explained_variance_ratio_[:n_components].sum():.1%} variance explained')
+    log.progress(f'PCA complete: {pca.explained_variance_ratio_.sum():.1%} variance explained')
+    return embedding, pca.explained_variance_ratio_
+
+
+# reduce_and_cluster: PCA followed by HDBSCAN or KMeans
+def reduce_and_cluster(features: np.ndarray, config: ClusteringConfig) -> ClusteringResult:
+    embedding, explained_variance_ratio = reduce_features(features, config)
+    log.progress(f'Starting {config.method} clustering')
+    log_config(config.method, config)
     labels, centroids = _cluster_embedding(embedding, config)
-    log.debug(f'{config.method}: {len(set(labels))} cluster(s) found')
-    return ClusteringResult(
-        labels=labels,
-        embedding=embedding,
-        explained_variance_ratio=pca.explained_variance_ratio_,
-        centroids=centroids,
-    )
+    log.progress(f'Clustering complete: {len(set(labels) - {-1})} cluster(s), {int((labels == -1).sum())} noise')
+    return ClusteringResult(labels=labels, embedding=embedding, explained_variance_ratio=explained_variance_ratio, centroids=centroids)
 
 # reduce_and_cluster_shared: one PCA basis fitted on every particle, each index group clustered in that basis
 def reduce_and_cluster_shared(features: np.ndarray, groups: dict[str, np.ndarray], config: ClusteringConfig) -> dict[str, ClusteringResult]:
-    n_components = min(config.n_components, features.shape[0] - 1, features.shape[1])
-    if n_components < 1:
-        raise ValueError(f'too few particles ({features.shape[0]}) to run PCA (at least 2 needed)')
-    pca = PCA(n_components=n_components, random_state=config.random_state)
-    embedding_all = pca.fit_transform(features)
+    embedding_all, explained_variance_ratio = reduce_features(features, config)
     results: dict[str, ClusteringResult] = {}
     for name, indices in groups.items():
+        log.progress(f'Starting {config.method} clustering for half {name}')
+        log_config(config.method, config)
         embedding = embedding_all[indices]
         labels, centroids = _cluster_embedding(embedding, config)
+        log.progress(f'Half {name} clustering complete: {len(set(labels) - {-1})} cluster(s), {int((labels == -1).sum())} noise')
         results[name] = ClusteringResult(
             labels=labels,
             embedding=embedding,
-            explained_variance_ratio=pca.explained_variance_ratio_,
+            explained_variance_ratio=explained_variance_ratio,
             centroids=centroids,
         )
     return results

src/stamp/classify/features.py

@@ -151,6 +151,7 @@ def build_feature_matrix(
         by_tomogram.setdefault(tomogram_id, []).append(local_index)
     with pool_context as pool:
         completed = 0
+        n_resumed = 0
         for tomogram_id, local_indices in by_tomogram.items():
             rows = [subvols.index[i][1] for i in local_indices]
             tomogram_magnitudes = _load_cached_magnitudes(cache_dir, tomogram_id, rows)
@@ -170,10 +171,13 @@ def build_feature_matrix(
                 if cache_dir is not None:
                     atomic_save_npz(cache_dir / f'{tomogram_id}.npz', rows=np.array(rows), magnitudes=tomogram_magnitudes)
             else:
+                n_resumed += 1
                 log.debug(f'{tomogram_id}: features loaded from checkpoint')
             completed += len(local_indices)
             for local_index, result in zip(local_indices, tomogram_magnitudes):
                 magnitudes[local_index] = result
+    if n_resumed:
+        log.info(f'Resuming feature extraction: {n_resumed}/{len(by_tomogram)} tomogram(s) loaded from checkpoint')
     stacked = np.stack(magnitudes)  # (n_particles, n_modes, n_radial_bins, box_voxels)
 
     n_particles = stacked.shape[0]

src/stamp/commands/classify.py

@@ -3,7 +3,7 @@
 '''
 
 # Import external dependencies
-import json, matplotlib.pyplot as plt, numpy as np
+import json, os, matplotlib.pyplot as plt, numpy as np
 from pathlib import Path
 
 # Import internal STAMP objects
@@ -17,7 +17,7 @@
     reduce_and_cluster,
     reduce_and_cluster_shared,
 )
-from stamp.utils.checkpoint import file_signature, sync_checkpoint
+from stamp.utils.checkpoint import file_signature, signatures, sync_checkpoint
 from stamp.classify.extract import extract_particle_set
 from stamp.classify.features import build_feature_matrix
 from stamp.backends.base import ToolCommand
@@ -84,8 +84,8 @@ def run_classify(
     # everything that changes results; n_workers, plots and keep_raw are left out so a run can resume with different resources
     fingerprint = {
         'particles': file_signature(particles, content=True),
-        'raw': {stem: file_signature(Path(path)) for stem, path in sorted(tomogram_paths.items())},
-        'segmentation': {stem: file_signature(Path(path)) for stem, path in sorted(segmentation_paths.items())},
+        'raw': signatures(tomogram_paths),
+        'segmentation': signatures(segmentation_paths),
         'voxel_size_angstrom': voxel_size_angstrom,
         'box_voxels': box_voxels,
         'n_radial_bins': n_radial_bins,
@@ -125,7 +125,7 @@ def run_classify(
         n_workers=n_workers,
         cache_dir=output_dir / 'raw' / 'features',
     )
-    log.debug(f'Feature vector: {features.shape[1]} dimensions (modes 0-{azimuthal_modes})')
+    log.progress(f'Feature matrix complete: {features.shape[0]} particles x {features.shape[1]} features (modes 0-{azimuthal_modes})')
 
     degenerate = ~features.any(axis=1)
     if degenerate.any():
@@ -145,6 +145,9 @@ def run_classify(
         random_state=random_state,
     )
     log_config('classify.clustering', config)
+
+    # set environment variable NUMBA_NUM_THREADS so fast_hdbscan doesn't try to use all CPUs on compute clusters
+    os.environ['NUMBA_NUM_THREADS'] = str(n_workers)
 
     output_dir.mkdir(parents=True, exist_ok=True)
 

src/stamp/commands/decoy.py

@@ -4,6 +4,7 @@
 
 # Import external dependencies
 import json
+from dataclasses import asdict
 from pathlib import Path
 
 # Import internal STAMP objects
@@ -20,6 +21,7 @@
 from stamp.picking.native import NativePickerConfig
 from stamp.run.state import stage_dir
 from stamp.schemas.particles import ParticleSet
+from stamp.utils.checkpoint import directory_signatures, file_signature, sync_checkpoint
 from stamp.utils.errors import StampPipelineError, StampValidationError
 from stamp.utils.io import archive_and_remove_directory, load_tomogram_manifests, resolve_directory_voxel_size_angstrom, write_sidecar
 from stamp.utils.log import log
@@ -67,6 +69,26 @@ def run_decoy(
 
     output_dir.mkdir(parents=True, exist_ok=True)
 
+    # everything that changes decoy positions; workers, plots and keep_raw are left out so a run can resume with different resources
+    uses_real_data = method != METHOD_SYNTHETIC_NOISE
+    fingerprint = {
+        'method': method,
+        'seed': seed,
+        'picker_config': asdict(config),
+        'n_decoys_per_tomogram': n_decoys_per_tomogram,
+        'min_distance_from_real_angstrom': min_distance_from_real_angstrom,
+        'min_shift_angstrom': min_shift_angstrom,
+        'max_shift_angstrom': max_shift_angstrom,
+        'min_distance_from_picks_angstrom': min_distance_from_picks_angstrom,
+        'n_synthetic_tomograms': n_synthetic_tomograms,
+        'synthetic_shape': synthetic_shape,
+        'real_particle_set': file_signature(real_particle_set, content=True) if uses_real_data and real_particle_set else None,
+        'segmentation': directory_signatures(segmentation_dir) if uses_real_data else {},
+        'raw': directory_signatures(raw_tomogram_dir) if uses_real_data else {},
+    }
+    if sync_checkpoint(output_dir / 'raw', fingerprint):
+        log.info('Found matching partial output, resuming')
+
     if method == METHOD_SYNTHETIC_NOISE:
         shape = tuple(int(value) for value in synthetic_shape.split(','))
         if len(shape) != 3: