From 81318564e55e4c6723cf6ec61b9f6e5248553553 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Mon, 28 Sep 2026 14:38:31 +0100
Subject: [PATCH 01/21] feat(log): add log utilities

- Modified utils/log.py to add two utility functions for logging startup messages and configuration values.
---
 src/stamp/utils/log.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/stamp/utils/log.py b/src/stamp/utils/log.py
index cbeb5d5..ab54715 100644
--- a/src/stamp/utils/log.py
+++ b/src/stamp/utils/log.py
@@ -3,9 +3,11 @@
 '''
 
 # Import external dependencies
-import sys
+import shlex, sys
+from dataclasses import asdict, is_dataclass
 from loguru import logger as logger
 from pathlib import Path
+from pydantic import BaseModel
 
 # Import internal STAMP objects
 import stamp.utils.io as io
@@ -146,4 +148,12 @@ def init_worker_logging(level_name: str, log_path: Path | None) -> None:
 # Create logger as instance of _StampLogger
 log = _StampLogger()
 
+# log_started: INFO message for a command starting (with terminal command used)
+def log_started(stage: str) -> None:
+    log.info(f'Started {stage} ({shlex.join(sys.argv)})')
 
+# log_config: DEBUG message for loaded/constructed config object's fields
+def log_config(name: str, model: BaseModel | object) -> None:
+    values = model.model_dump() if isinstance(model, BaseModel) else asdict(model) if is_dataclass(model) else vars(model)
+    fields = ', '.join(f'{k}={v!r}' for k, v in values.items())
+    log.debug(f'Configuration loaded for {name}: {fields}')

From 9615c9d36c4cbe33f08f30c6e21e0970804a17a8 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Mon, 28 Sep 2026 14:43:08 +0100
Subject: [PATCH 02/21] style(cli): add command start log messages

- Modified cli/pick.py, cli/decoy.py, cli/classify.py, cli/identify.py, cli/refine.py, commands/classify.py to add startup log message.
---
 src/stamp/cli/classify.py | 2 ++
 src/stamp/cli/decoy.py    | 2 ++
 src/stamp/cli/identify.py | 2 ++
 src/stamp/cli/pick.py     | 2 ++
 src/stamp/cli/refine.py   | 2 ++
 src/stamp/commands/run.py | 3 ++-
 6 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/stamp/cli/classify.py b/src/stamp/cli/classify.py
index cdc7c65..ae2df69 100644
--- a/src/stamp/cli/classify.py
+++ b/src/stamp/cli/classify.py
@@ -13,6 +13,7 @@
 from stamp.schemas.particles import ParticleSet
 from stamp.utils.errors import StampPipelineError
 from stamp.utils.io import resolve_output_dir
+from stamp.utils.log import log_started
 
 # Initialise Typer app
 classifyCli = typer.Typer(
@@ -115,6 +116,7 @@ def classify(
     ] = False,
 ) -> None:
     '''Cluster picked particles by structural similarity.'''
+    log_started('classification')
     is_decoy = is_decoy_particle_set(ParticleSet.model_validate(json.loads(particles.read_text())))
     classifyfuncs.run_classify(
         particles=particles,
diff --git a/src/stamp/cli/decoy.py b/src/stamp/cli/decoy.py
index 7b3531f..d493220 100644
--- a/src/stamp/cli/decoy.py
+++ b/src/stamp/cli/decoy.py
@@ -11,6 +11,7 @@
 import stamp.commands.decoy as decoyfuncs
 from stamp.utils.errors import StampPipelineError
 from stamp.utils.io import resolve_output_dir
+from stamp.utils.log import log_started
 
 # Initialise Typer app
 decoyCli = typer.Typer(
@@ -111,6 +112,7 @@ def decoy(
     ] = False,
 ) -> None:
     '''Generate a decoy dataset to run through STAMP alongside real data.'''
+    log_started('decoy generation')
     # Validate picker parameters
     try:
         parameters = json.loads(picker_params)
diff --git a/src/stamp/cli/identify.py b/src/stamp/cli/identify.py
index bb0e82d..54e3ab1 100644
--- a/src/stamp/cli/identify.py
+++ b/src/stamp/cli/identify.py
@@ -10,6 +10,7 @@
 # Import identify command functions
 import stamp.commands.identify as identifyfuncs
 from stamp.utils.io import resolve_output_dir
+from stamp.utils.log import log_started
 
 # Initialise Typer app
 identifyCli = typer.Typer(no_args_is_help=True, add_completion=False)
@@ -63,6 +64,7 @@ def identify(
     ] = 'tiff',
 ) -> None:
     '''Fit predicted structures to class averages and score candidates.'''
+    log_started('identification')
     identifyfuncs.run_identify(
         classes=classes,
         candidates=candidates,
diff --git a/src/stamp/cli/pick.py b/src/stamp/cli/pick.py
index 51a7617..9cd599d 100644
--- a/src/stamp/cli/pick.py
+++ b/src/stamp/cli/pick.py
@@ -10,6 +10,7 @@
 # Import pick command functions/variables
 import stamp.commands.pick as pickfuncs
 from stamp.utils.io import resolve_output_dir
+from stamp.utils.log import log_started
 
 # Initialise Typer app
 pickCli = typer.Typer(
@@ -102,6 +103,7 @@ def pick(
     ] = False,
 ) -> None:
     '''Run particle picking, reconcile across pickers, and assign half-sets.'''
+    log_started('picking')
     # Validate provided pickers
     picker_names = [name.strip() for name in pickers.split(',') if name.strip()]
     if not picker_names:
diff --git a/src/stamp/cli/refine.py b/src/stamp/cli/refine.py
index 14667d7..37bdfef 100644
--- a/src/stamp/cli/refine.py
+++ b/src/stamp/cli/refine.py
@@ -10,6 +10,7 @@
 # Import refine command functions/variables
 import stamp.commands.refine as refinefuncs
 from stamp.utils.io import resolve_output_dir
+from stamp.utils.log import log_started
 
 # Initialise Typer app
 refineCli = typer.Typer(no_args_is_help=True, add_completion=False)
@@ -79,6 +80,7 @@ def refine(
     ] = 'tiff',
 ) -> None:
     '''Refine identified classes with independent half-sets.'''
+    log_started('refinement')
     refinefuncs.run_refine(
         class_id=class_id, 
         identification=identification,
diff --git a/src/stamp/commands/run.py b/src/stamp/commands/run.py
index fe6accb..0e7f636 100644
--- a/src/stamp/commands/run.py
+++ b/src/stamp/commands/run.py
@@ -9,7 +9,7 @@
 from stamp.run.orchestrate import run_pipeline
 from stamp.run.report import write_report
 from stamp.schemas.config import load_run_config
-from stamp.utils.log import log
+from stamp.utils.log import log, log_started
 
 # run_full_pipeline: load the config, run every track, write the combined report
 def run_full_pipeline(
@@ -19,6 +19,7 @@ def run_full_pipeline(
     no_decoy: bool,
     backend: str | None
 ) -> None:
+    log_started('run')
     log.progress(f'Running STAMP pipeline from {config_path}')
     config = load_run_config(config_path)
     if no_decoy:

From 84cdd5df7c56e7cb5ea275859c7f9b36ec588ac3 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Mon, 28 Sep 2026 14:51:44 +0100
Subject: [PATCH 03/21] style(pick): add logging messages

- Modified commands/pick.py to add messages for vesicle summary, adapter/dispatch branching and resolved voxel size.
---
 src/stamp/commands/pick.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/stamp/commands/pick.py b/src/stamp/commands/pick.py
index d074880..22c90ee 100644
--- a/src/stamp/commands/pick.py
+++ b/src/stamp/commands/pick.py
@@ -46,6 +46,7 @@ def _execute(adapter: ToolAdapter, inputs: AdapterInputs, runner: Runner) -> lis
 # _select_adapter: return the adapter for a picker name, guarding it against the chosen backend
 def _select_adapter(picker_name: str, backend: str) -> ToolAdapter:
     if backend == 'mock':
+        log.debug(f'{picker_name}: using mock adapter (backend=mock)')
         return get_mock_adapter(picker_name)
 
     adapter = REAL_ADAPTERS.get(picker_name)
@@ -64,6 +65,7 @@ def _run_picker(
     backend: str,
 ) -> list[RawPick]:
     if backend != 'mock' and getattr(adapter, 'runs_in_process', False):
+        log.debug(f'{adapter.name}: dispatching in-process for {len(manifests)} tomogram(s)')
         inputs = AdapterInputs(
             input_paths=[m.segmentation_path for m in manifests],
             raw_tomogram_paths=[m.raw_tomogram_path for m in manifests],
@@ -74,6 +76,7 @@ def _run_picker(
         return adapter.run_in_process(inputs)  # type: ignore[attr-defined]
 
     if backend == 'mock' or adapter.batches_natively:
+        log.debug(f'{adapter.name}: dispatching as a single batched command for {len(manifests)} tomogram(s)')
         inputs = AdapterInputs(
             input_paths=[m.segmentation_path for m in manifests],
             raw_tomogram_paths=[m.raw_tomogram_path for m in manifests],
@@ -83,6 +86,7 @@ def _run_picker(
         )
         return _execute(adapter, inputs, runner)
 
+    log.debug(f'{adapter.name}: dispatching one command per tomogram for {len(manifests)} tomogram(s)')
     collected: list[RawPick] = []
     for manifest in manifests:
         inputs = AdapterInputs(
@@ -206,6 +210,7 @@ def run_pick(
         raise StampPipelineError(f'No .mrc files in {segmentation_dir} with a matching stem in {raw_tomogram_dir}')
     log.info(f'Matched {len(manifests)} tomograms/segmentations')
     resolved_voxel_size_angstrom = manifests[0].voxel_size_angstrom
+    log.debug(f'Resolved voxel size: {resolved_voxel_size_angstrom} Å')
 
     runner = select_runner(backend)
     output_dir.mkdir(parents=True, exist_ok=True)
@@ -265,6 +270,7 @@ def run_pick(
     particle_set_path.write_text(particle_set.model_dump_json(indent=2))
     log.info(f'Wrote {len(particle_set.particles)} consensus particles to {particle_set_path}')
 
+    log.progress('Computing vesicle summary')
     _write_vesicle_summary(all_reconciled, manifests, vesicle_labels_mrc_by_tomogram, output_dir, n_workers=n_workers)
     log.progress(f'Writing parameters...')
     write_sidecar(

From 2b3a8b455f9d4e8dc9a1f882ad47341e0ddcd49c Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Mon, 28 Sep 2026 14:55:39 +0100
Subject: [PATCH 04/21] style(classify): add logging messages

- Modified commands/classify.py to add messages for subvolume extraction, feature extraction and dispatch, and for configurations, resolved voxel size, and strict/combined branching.
---
 src/stamp/commands/classify.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py
index 3bd19b5..ee28837 100644
--- a/src/stamp/commands/classify.py
+++ b/src/stamp/commands/classify.py
@@ -27,7 +27,7 @@
 from stamp.schemas.subvolumes import Subvolumes
 from stamp.utils.errors import StampPipelineError
 from stamp.utils.io import archive_and_remove_directory, match_by_stem, resolve_directory_voxel_size_angstrom, write_sidecar
-from stamp.utils.log import log
+from stamp.utils.log import log, log_config
 from stamp.utils.plotting.core import PlotFormat, central_slice, finish, plot_path
 from stamp.utils.plotting.classify import class_average_grid, scatter_labels
 
@@ -61,6 +61,7 @@ def run_classify(
         voxel_size_angstrom = resolve_directory_voxel_size_angstrom(list(raw_tomogram_dir.glob('*.mrc')))
         if voxel_size_angstrom is None:
             raise StampPipelineError('Voxel size not given and not found in any MRC header in --raw-dir.')
+        log.debug(f'Resolved voxel size from MRC headers: {voxel_size_angstrom} Å')
     log.progress('Classifying particle set')
     particle_set = ParticleSet.model_validate(json.loads(particles.read_text()))
     is_decoy = is_decoy_particle_set(particle_set)
@@ -76,7 +77,10 @@ def run_classify(
     box_voxels = int(round(box_angstrom / voxel_size_angstrom))
     if box_voxels % 2 == 0:
         box_voxels += 1  # odd box keeps the particle exactly centred
+        log.debug(f'Using odd box voxel for centred particles: {box_voxels - 1} -> {box_voxels}')    
+    log.progress(f'Using box voxels: {box_voxels}')
 
+    log.progress('Starting subvolume extraction')
     subvols, kept, skipped = extract_particle_set(
         particle_set.particles, {k: str(v) for k, v in tomogram_paths.items()}, box_voxels,
         segmentation_paths=segmentation_paths, n_workers=n_workers, cache_dir=output_dir / 'raw' / 'subvolumes',
@@ -87,6 +91,7 @@ def run_classify(
         raise StampPipelineError('No particles could be extracted. Check --raw-dir and --box-length-a')
     log.info(f'Extracted {len(kept)} subvolumes at box {box_voxels}{" (membrane subtracted)" if segmentation_paths else ""}')
 
+    log.progress('Starting feature matrix construction')
     features = build_feature_matrix(
         subvols,
         n_radial_bins=n_radial_bins,
@@ -114,6 +119,7 @@ def run_classify(
         n_clusters=n_clusters,
         random_state=random_state,
     )
+    log_config('classify.clustering', config)
 
     output_dir.mkdir(parents=True, exist_ok=True)
 
@@ -129,8 +135,10 @@ def run_classify(
     if not strict_halfset_independence:
         log.warning('--no-strict-halfset-independence used, half-A and half-B particles are clustered together; FSC built on these classes will be inflated')
     if strict_halfset_independence:
+        log.debug('Dispatching to strict half-set-independent clustering')
         assignments, result_by_half, averages_by_half = _classify_strict(kept, subvols, features, config, output_dir, voxel_size_angstrom, align_settings)
     else:
+        log.debug('Dispatching to combined clustering (half-sets pooled)')
         assignments, result_by_half, averages_by_half = _classify_combined(kept, subvols, features, config, output_dir, voxel_size_angstrom, align_settings)
 
     assignments_path = output_dir / 'class_assignments.json'

From b0d3af6a3dc602598d11191072c8f0663cde3675 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Mon, 28 Sep 2026 15:11:34 +0100
Subject: [PATCH 05/21] style(identify): add logging messages

- Modified commands/identify.py to add messages for report writing and decoy control branching.
---
 src/stamp/commands/identify.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/stamp/commands/identify.py b/src/stamp/commands/identify.py
index 1957bdb..03169b3 100644
--- a/src/stamp/commands/identify.py
+++ b/src/stamp/commands/identify.py
@@ -128,7 +128,10 @@ def run_identify(
             log.info('Decoy control passed')
         else:
             log.warning(f'Decoy control failed: {decoy_control.reason}')
+    else:
+        log.debug('No --decoy-classes given, skipping decoy control')
 
+    log.progress('Writing identification report')
     _write_report(output_dir / 'identification_report.txt', results, real_scores, decoy_control, resolution, inplane_aligned)
 
     write_sidecar(

From 016e0012d28c4860535b47864d3507a6270a97c5 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Mon, 28 Sep 2026 15:14:27 +0100
Subject: [PATCH 06/21] style(refine): add logging messages

- Modified commands/refine.py to add messages for FSC, backend branching and resolved voxel size.
---
 src/stamp/commands/refine.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/stamp/commands/refine.py b/src/stamp/commands/refine.py
index 30a3d5e..f5ff93c 100644
--- a/src/stamp/commands/refine.py
+++ b/src/stamp/commands/refine.py
@@ -59,6 +59,7 @@ def _run_half(adapter, runner, particles, reference, workdir, parameters, backen
         raise StampPipelineError(f'Refine failed for {workdir.name}: {result.stderr}')
     final_map = workdir / f'final_{workdir.name[-1]}.mrc'
     if backend == 'mock':
+        log.debug(f'{workdir.name}: mock backend, synthesising final map from the seed reference')
         # mock backend runs nothing; synthesise a deterministic map from the seed
         with mrcfile.open(str(reference), permissive=True) as mrc:
             seed = np.asarray(mrc.data, dtype=np.float32)
@@ -91,6 +92,7 @@ def run_refine(
         voxel_size_angstrom = resolve_directory_voxel_size_angstrom(list(raw_tomogram_dir.glob('*.mrc')))
         if voxel_size_angstrom is None:
             raise StampValidationError('Voxel size not given and not found in any MRC header in --raw-dir.')
+        log.debug(f'Resolved voxel size from MRC headers: {voxel_size_angstrom} Å')
     particle_set = ParticleSet.model_validate(json.loads(particles.read_text()))
     assignments = [ClassAssignment.model_validate(row) for row in json.loads(class_assignments.read_text())]
     angle_by_particle = {row.particle_id: row.inplane_angle_degrees for row in assignments if row.inplane_angle_degrees is not None}
@@ -129,6 +131,7 @@ def run_refine(
             if mask is not None:
                 with mrcfile.open(str(mask), permissive=True) as mrc:
                     user_mask = np.asarray(mrc.data, dtype=np.float32)
+            log.progress(f'{target}: computing FSC')
             fsc = compute_fsc(map_a, map_b, voxel_size_angstrom, mask=user_mask if user_mask is not None else soft_sphere_mask(map_a.shape), threshold=threshold)
         except (StampValidationError, StampAdapterError) as exc:
             log.error(f'{target}: {exc}')

From 6b1825a5ca46c3844a1582a75a67d1a6bd1318ab Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Mon, 28 Sep 2026 15:25:12 +0100
Subject: [PATCH 07/21] style(run): add logging messages

- Modified run/orchestrate.py to add message debug for inferred synthetic shape.
- Modified schemas/config.py to add message for run config after loading.
---
 src/stamp/run/orchestrate.py | 4 +++-
 src/stamp/schemas/config.py  | 3 ++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/stamp/run/orchestrate.py b/src/stamp/run/orchestrate.py
index bfdfb39..944959a 100644
--- a/src/stamp/run/orchestrate.py
+++ b/src/stamp/run/orchestrate.py
@@ -56,7 +56,9 @@ def _synthetic_shape(config: RunConfig) -> tuple[int, int, int]:
         log.debug('No real tomograms found to infer synthetic shape from, defaulting to (200, 200, 200)')
         return (200, 200, 200)
     with mrcfile.open(str(first), permissive=True, header_only=True) as mrc:
-        return (int(mrc.header.nz), int(mrc.header.ny), int(mrc.header.nx))
+        shape = (int(mrc.header.nz), int(mrc.header.ny), int(mrc.header.nx))
+    log.debug(f'Inferred synthetic shape {shape} from {first}')
+    return shape
 
 # _decoy_pick: run decoy generation
 def _decoy_pick(config: RunConfig, output_dir: Path, real_particle_set: Path) -> Path:
diff --git a/src/stamp/schemas/config.py b/src/stamp/schemas/config.py
index e62b810..bb7b597 100644
--- a/src/stamp/schemas/config.py
+++ b/src/stamp/schemas/config.py
@@ -10,7 +10,7 @@
 
 # Import internal STAMP objects
 from stamp.utils.errors import StampPipelineError
-from stamp.utils.log import log
+from stamp.utils.log import log, log_config
 
 # _Strict: reject unknown keys, treat "" as unset
 class _Strict(BaseModel):
@@ -139,4 +139,5 @@ def load_run_config(path: Path) -> RunConfig:
     if config.stage.refine.backend is None:
         log.debug(f'stage.refine.backend not set, inheriting run.backend={config.run.backend!r}')
         config.stage.refine.backend = config.run.backend
+    log_config('run', config)
     return config

From 6f57770f60a4960c299b76cddc9117e0d8bb36ea Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Mon, 28 Sep 2026 15:44:26 +0100
Subject: [PATCH 08/21] fix(classify): stop unrequired mrc type conversion

- Modified classify/extract.py to only change MRC file data type to float32 if it isn't already float32, preventing double memory usage for no gain.
---
 src/stamp/classify/extract.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index a34394c..3f90f08 100644
--- a/src/stamp/classify/extract.py
+++ b/src/stamp/classify/extract.py
@@ -76,7 +76,9 @@ def box_fits_inside(
 # _extract_one_tomogram: worker entry point — loads one tomogram and extracts subvolumes for its particles
 def _extract_one_tomogram(job: _ExtractTomogramJob) -> tuple[list[Particle], list[Particle]]:
     with mrcfile.open(str(job.tomogram_path), permissive=True) as mrc:
-        tomogram = np.asarray(mrc.data).astype(np.float32)
+        tomogram = np.asarray(mrc.data)
+        if tomogram.dtype != np.float32:
+            tomogram=tomogram.astype(np.float32)
 
     if job.segmentation_path is not None:
         with mrcfile.open(str(job.segmentation_path), permissive=True) as mrc:

From 1858fc8a552ab6ac3e29c420eaaba10f2fbf7ea8 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 09:28:20 +0100
Subject: [PATCH 09/21] fix(classify): correct hdbscan string

---
 docs/cli.md               | 2 +-
 src/stamp/cli/classify.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index acaece1..0fb4097 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -123,7 +123,7 @@ Option | Default | Use
 `-o, --out-dir` | `.` | Output directory (outputs will be saved to `<dir>/stamp/classify` or `<dir/stamp/classify_decoy>`)
 `--box-length-a` | `300.0` | Extraction box edge length in Å (roughly 2x the largest expected particle)
 `--n-bins` | `12` | Radial bins in the rotational average
-`--method` | `hbdscan` | Clustering method to use (`hbdscan` or `kmeans`)
+`--method` | `hdbscan` | Clustering method to use (`hdbscan` or `kmeans`)
 `--min-cluster-size` | `20` | HDBSCAN minimum cluster size
 `--n-clusters` | `5` | Number of KMeans clusters to use
 `--n-components` | `20` | Number of PCA components to use
diff --git a/src/stamp/cli/classify.py b/src/stamp/cli/classify.py
index ae2df69..5f0c53e 100644
--- a/src/stamp/cli/classify.py
+++ b/src/stamp/cli/classify.py
@@ -53,9 +53,9 @@ def classify(
         typer.Option('--n-bins', help='Radial bins in the rotational average.', rich_help_panel = 'Subvolume extraction'),
     ] = 12,
     method: Annotated[
-        Literal['hbdscan', 'kmeans'],
+        Literal['hdbscan', 'kmeans'],
         typer.Option('--method', help='Clustering method to use.', rich_help_panel = 'Clustering'),
-    ] = 'hbdscan',
+    ] = 'hdbscan',
     min_cluster_size: Annotated[
         int,
         typer.Option('--min-cluster-size', help='Minimum cluster size (used by HDBSCAN).', rich_help_panel = 'Clustering'),

From c54dd82175059d46b116842f36ea246bb819ffb7 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 09:45:55 +0100
Subject: [PATCH 10/21] feat(extract): add subtraction job dataclass

- Modified classify/extract.py to add a dataclass for membrane subtraction jobs for splitting membrane subtraction from subtomogram extraction.
---
 src/stamp/classify/extract.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index 3f90f08..53428e1 100644
--- a/src/stamp/classify/extract.py
+++ b/src/stamp/classify/extract.py
@@ -21,10 +21,17 @@ class _ExtractTomogramJob:
     tomogram_id: str
     group: list[Particle]
     tomogram_path: str
-    segmentation_path: str | None
     box_voxels: int
     cache_dir: Path
 
+# _SubtractTomogramJob: one tomogram's membrane subtraction, written to disk for extraction to mmap
+@dataclass(frozen=True)
+class _SubtractTomogramJob:
+    tomogram_id: str
+    tomogram_path: str
+    segmentation_path: str
+    output_path: Path
+
 # quaternion_to_matrix: return a rotation matrix from a unit quaternion (w, x, y, z)
 def quaternion_to_matrix(quaternion: tuple[float, float, float, float]) -> np.ndarray:
     w, x, y, z = quaternion

From b91dc17ae29101d3d0ae44c57da782489631cab7 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 09:48:59 +0100
Subject: [PATCH 11/21] feat(extract): add membrane subtraction worker

- Modified classify/extract.py to add function for subtracting membranes as a process pool worker process.
---
 src/stamp/classify/extract.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index 53428e1..837d138 100644
--- a/src/stamp/classify/extract.py
+++ b/src/stamp/classify/extract.py
@@ -80,7 +80,22 @@ def box_fits_inside(
     position = np.array(position_xyz)
     return bool(np.all(position >= radius) and np.all(position <= extent_xyz - radius))
 
-# _extract_one_tomogram: worker entry point — loads one tomogram and extracts subvolumes for its particles
+# _subtract_one_tomogram: worker entry point to replace membrane voxels with the background median and write a float32 mrc
+def _subtract_one_tomogram(job: _SubtractTomogramJob) -> Path:
+    with mrcfile.open(job.tomogram_path, permissive=True) as mrc:
+        tomogram = np.array(mrc.data, dtype=np.float32)  # writable copy
+        voxel_size = mrc.voxel_size.copy()
+    with mrcfile.open(job.segmentation_path, permissive=True) as mrc:
+        membrane = np.asarray(mrc.data) > 0
+    if membrane.shape != tomogram.shape:
+        raise StampValidationError(f'segmentation shape {membrane.shape} does not match tomogram shape {tomogram.shape} for {job.tomogram_id!r}; they must be the same volume at the same binning')
+    # replace membrane voxels with the background median so the class average is not dominated by the membrane slab
+    tomogram[membrane] = np.median(tomogram[~membrane])
+    job.output_path.parent.mkdir(parents=True, exist_ok=True)
+    with mrcfile.new(str(job.output_path), overwrite=True) as out:
+        out.set_data(tomogram)
+        out.voxel_size = voxel_size
+    return job.output_path
 def _extract_one_tomogram(job: _ExtractTomogramJob) -> tuple[list[Particle], list[Particle]]:
     with mrcfile.open(str(job.tomogram_path), permissive=True) as mrc:
         tomogram = np.asarray(mrc.data)

From bc67452fdaa6d50fdee2a8296590e043b5d108f6 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 09:52:16 +0100
Subject: [PATCH 12/21] feat(extract): update subvolume extraction worker

- Modified classify/extract.py to update subvolume extraction worker to decouple membrane subtraction.
---
 src/stamp/classify/extract.py | 32 ++++++++++++--------------------
 1 file changed, 12 insertions(+), 20 deletions(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index 837d138..eb47684 100644
--- a/src/stamp/classify/extract.py
+++ b/src/stamp/classify/extract.py
@@ -96,35 +96,27 @@ def _subtract_one_tomogram(job: _SubtractTomogramJob) -> Path:
         out.set_data(tomogram)
         out.voxel_size = voxel_size
     return job.output_path
-def _extract_one_tomogram(job: _ExtractTomogramJob) -> tuple[list[Particle], list[Particle]]:
-    with mrcfile.open(str(job.tomogram_path), permissive=True) as mrc:
-        tomogram = np.asarray(mrc.data)
-        if tomogram.dtype != np.float32:
-            tomogram=tomogram.astype(np.float32)
-
-    if job.segmentation_path is not None:
-        with mrcfile.open(str(job.segmentation_path), permissive=True) as mrc:
-            segmentation = np.asarray(mrc.data)
-        if segmentation.shape != tomogram.shape:
-            raise StampValidationError(f'segmentation shape {segmentation.shape} does not match tomogram shape {tomogram.shape} for {job.tomogram_id!r}; they must be the same volume at the same binning')
-        # replace membrane voxels with the background median so the class average is not dominated by the membrane slab
-        tomogram[segmentation > 0] = np.median(tomogram[segmentation <= 0])
 
+# _extract_one_tomogram: worker entry point to mmaps one tomogram and extracts subvolumes for its particles
+def _extract_one_tomogram(job: _ExtractTomogramJob) -> tuple[list[Particle], list[Particle]]:
     kept: list[Particle] = []
     skipped: list[Particle] = []
     subvolumes: list[np.ndarray] = []
-    for particle in job.group:
-        if particle.orientation is None or not box_fits_inside(particle.position, tomogram.shape, job.box_voxels):
-            skipped.append(particle)
-            continue
-        subvolumes.append(extract_subvolume(tomogram, particle.position, particle.orientation, job.box_voxels))
-        kept.append(particle)
+    with mrcfile.mmap(job.tomogram_path, permissive=True) as mrc:
+        tomogram = mrc.data
+        if tomogram.dtype != np.float32:
+            tomogram = tomogram.astype(np.float32)
+        for particle in job.group:
+            if particle.orientation is None or not box_fits_inside(particle.position, tomogram.shape, job.box_voxels):
+                skipped.append(particle)
+                continue
+            subvolumes.append(extract_subvolume(tomogram, particle.position, particle.orientation, job.box_voxels))
+            kept.append(particle)
     if subvolumes:
         job.cache_dir.mkdir(parents=True, exist_ok=True)
         np.save(job.cache_dir / f'{job.tomogram_id}.npy', np.stack(subvolumes))
     return kept, skipped
 
-# extract_particle_set: extract subvolumes for a list of particles
 # extract_particle_set: extract subvolumes for a list of particles, cached per tomogram
 def extract_particle_set(
     particles: list[Particle],

From 35870d29bf5ab5bfb623412bc91d28e7046ad597 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 09:56:35 +0100
Subject: [PATCH 13/21] feat(extract): update extraction logic

- Modified classify/extract.py to update extraction logic to decouple membrane subtraction and subvolume extraction.
---
 src/stamp/classify/extract.py | 42 ++++++++++++++++++++++++++++-------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index eb47684..fac40b4 100644
--- a/src/stamp/classify/extract.py
+++ b/src/stamp/classify/extract.py
@@ -135,6 +135,9 @@ def extract_particle_set(
     kept: list[Particle] = []
     skipped: list[Particle] = []
 
+    # subtract membranes to cached mrc files
+    subtracted_dir = cache_dir / '.subtracted'
+    subtract_jobs: list[_SubtractTomogramJob] = []
     jobs: list[_ExtractTomogramJob] = []
     for tomogram_id, group in by_tomogram.items():
         path = tomogram_paths.get(tomogram_id)
@@ -142,24 +145,47 @@ def extract_particle_set(
             log.warning(f'{tomogram_id}: no matching raw tomogram path, skipping {len(group)} particle(s)')
             skipped.extend(group)
             continue
-        jobs.append(_ExtractTomogramJob(tomogram_id, group, path, segmentation_paths.get(tomogram_id), box_voxels, cache_dir))
+        segmentation_path = segmentation_paths.get(tomogram_id)
+        if segmentation_path is not None:
+            subtract_jobs.append(_SubtractTomogramJob(tomogram_id, path, segmentation_path, subtracted_dir / f'{tomogram_id}.mrc'))
+        else:
+            jobs.append(_ExtractTomogramJob(tomogram_id, group, path, box_voxels, cache_dir))
 
     result_by_tomogram: dict[str, tuple[list[Particle], list[Particle]]] = {}
 
-    def _on_success(job: _ExtractTomogramJob, result: tuple[list[Particle], list[Particle]]) -> None:
-        result_by_tomogram[job.tomogram_id] = result
-
-    def _on_error(job: _ExtractTomogramJob, exc: Exception) -> None:
+    def _on_error(job, exc: Exception) -> None:
         if isinstance(exc, StampValidationError):
             log.error(f'{job.tomogram_id}: {exc}')
-            skipped.extend(job.group)
+            skipped.extend(by_tomogram[job.tomogram_id])
         else:
             raise exc
 
-    run_parallel(jobs, _extract_one_tomogram, max_workers=n_workers, label='subvolume-extraction', on_success=_on_success, on_error=_on_error)
+    def _on_subtract_error(job: _SubtractTomogramJob, exc: Exception) -> None:
+        if not isinstance(exc, StampValidationError):
+            raise exc
+        # fall back to the raw tomogram so the particles are still classified, just without subtraction
+        log.warning(f'{job.tomogram_id}: {exc}; extracting from the raw tomogram without membrane subtraction')
+        jobs.append(_ExtractTomogramJob(job.tomogram_id, by_tomogram[job.tomogram_id], job.tomogram_path, box_voxels, cache_dir))
+
+    subtracted_ids: set[str] = set()
+
+    def _on_subtracted(job: _SubtractTomogramJob, output_path: Path) -> None:
+        subtracted_ids.add(job.tomogram_id)
+        jobs.append(_ExtractTomogramJob(job.tomogram_id, by_tomogram[job.tomogram_id], str(output_path), box_voxels, cache_dir))
+
+    def _on_success(job: _ExtractTomogramJob, result: tuple[list[Particle], list[Particle]]) -> None:
+        result_by_tomogram[job.tomogram_id] = result
+        if Path(job.tomogram_path).parent == subtracted_dir:
+            Path(job.tomogram_path).unlink(missing_ok=True)  # free disk as soon as a tomogram is done
+
+    try:
+        run_parallel(subtract_jobs, _subtract_one_tomogram, max_workers=n_workers, label='membrane-subtraction', on_success=_on_subtracted, on_error=_on_subtract_error)
+        run_parallel(jobs, _extract_one_tomogram, max_workers=n_workers, label='subvolume-extraction', on_success=_on_success, on_error=_on_error)
+    finally:
+        shutil.rmtree(subtracted_dir, ignore_errors=True)
 
     index: list[tuple[str, int]] = []
-    for job in jobs:
+    for job in sorted(jobs, key=lambda job: job.tomogram_id):
         result = result_by_tomogram.get(job.tomogram_id)
         if result is None:
             continue

From 6fca671ecec9f51a9d2844e2fa7f5f58bda94aff Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 10:00:04 +0100
Subject: [PATCH 14/21] feat(classify): finalise decoupling

- Modified classify/extract.py to add missing import and logging line for number of membrane-subtracted/raw density extractions.
- Modified commands/classify.py to remove log line for kept & skipped subvolumes (now reported directly).
- Modified tests/unit/test_classify.py to add unit test for fallback to raw tomogram if segmentation shape doesn't match.
---
 src/stamp/classify/extract.py  |  9 +++++----
 src/stamp/commands/classify.py |  1 -
 tests/unit/test_classify.py    | 18 ++++++++++++++++++
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index fac40b4..ac547c6 100644
--- a/src/stamp/classify/extract.py
+++ b/src/stamp/classify/extract.py
@@ -3,7 +3,7 @@
 '''
 
 # Import external dependencies
-import mrcfile, numpy as np
+import mrcfile, numpy as np, shutil
 from dataclasses import dataclass
 from pathlib import Path
 from scipy.ndimage import map_coordinates
@@ -66,9 +66,7 @@ def extract_subvolume(
     sample_xyz = grid_xyz + np.array(position_xyz)[:, None]
     sample_zyx = sample_xyz[::-1]  # single xyz -> zyx conversion
 
-    values = map_coordinates(
-        tomogram.astype(np.float32), sample_zyx, order=1, mode='constant', cval=0.0
-    )
+    values = map_coordinates(np.asarray(tomogram, dtype=np.float32), sample_zyx, order=1, mode='constant', cval=0.0)
     return values.reshape(box_voxels, box_voxels, box_voxels)
 
 # box_fits_inside: whether a rotated box at this position stays inside the volume
@@ -195,4 +193,7 @@ def _on_success(job: _ExtractTomogramJob, result: tuple[list[Particle], list[Par
         index.extend((job.tomogram_id, row) for row in range(len(tomogram_kept)))
 
     log.debug(f'extract_particle_set: {len(kept)} kept, {len(skipped)} skipped')
+    n_subtracted = sum(particle.tomogram_id in subtracted_ids for particle in kept)
+    parts = [f'{n} ({label})' for n, label in ((n_subtracted, 'membrane-subtracted'), (len(kept) - n_subtracted, 'raw density')) if n]
+    log.info(f'Extracted {" + ".join(parts)} subvolumes at box {box_voxels}')
     return Subvolumes(cache_dir, index, box_voxels), kept, skipped
diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py
index ee28837..935b6cf 100644
--- a/src/stamp/commands/classify.py
+++ b/src/stamp/commands/classify.py
@@ -89,7 +89,6 @@ def run_classify(
         log.warning(f'Skipped {len(skipped)} particles whose {box_voxels}-voxel box fell outside the volume or had no matching tomogram or had no orientation')
     if not kept:
         raise StampPipelineError('No particles could be extracted. Check --raw-dir and --box-length-a')
-    log.info(f'Extracted {len(kept)} subvolumes at box {box_voxels}{" (membrane subtracted)" if segmentation_paths else ""}')
 
     log.progress('Starting feature matrix construction')
     features = build_feature_matrix(
diff --git a/tests/unit/test_classify.py b/tests/unit/test_classify.py
index 6919f39..3817ab9 100644
--- a/tests/unit/test_classify.py
+++ b/tests/unit/test_classify.py
@@ -149,6 +149,24 @@ def test_no_segmentation_leaves_tomogram_untouched(self, tmp_path) -> None:
         subvolumes, kept, _ = extract_particle_set([particle], {'t': str(tmp_path / 't.mrc')}, 11, cache_dir=tmp_path)
         assert float(subvolumes[0][5, 5, 5]) == 50.0
 
+    def test_shape_mismatch_falls_back_to_raw_tomogram(self, tmp_path, caplog) -> None:
+        tomogram = np.zeros((40, 40, 40), dtype=np.float32)
+        tomogram[:, :, 20] = 50.0
+        segmentation = np.zeros((30, 30, 30), dtype=np.float32)  # wrong shape
+        with mrcfile.new(tmp_path / 't.mrc', overwrite=True) as mrc:
+            mrc.set_data(tomogram)
+            mrc.voxel_size = 1.0
+        with mrcfile.new(tmp_path / 's.mrc', overwrite=True) as mrc:
+            mrc.set_data(segmentation)
+            mrc.voxel_size = 1.0
+        particle = Particle(particle_id='p0', tomogram_id='t', position=(20.0, 20.0, 20.0), orientation=(1.0, 0.0, 0.0, 0.0), source_picker='test', confidence=1.0, half_set=HalfSet.A)
+        subvolumes, kept, skipped = extract_particle_set([particle], {'t': str(tmp_path / 't.mrc')}, 11, segmentation_paths={'t': str(tmp_path / 's.mrc')}, cache_dir=tmp_path)
+        assert len(kept) == 1 and not skipped
+        assert float(subvolumes[0][5, 5, 5]) == 50.0  # raw density, membrane untouched
+        assert not (tmp_path / '.subtracted').exists()
+        assert 'without membrane subtraction' in caplog.text
+        assert '1 (raw density)' in caplog.text
+
 # TestFeatures: class containing unit tests for src/stamp/classify/features.py
 class TestFeatures:
     def test_rotational_average_shape(self) -> None:

From 8b774533e2f056c9ae86940f19dde7d7b647cc1a Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 10:14:25 +0100
Subject: [PATCH 15/21] feat(utils): add checkpointing utilities

- Added utils/checkpoint.py to define utility functions for getting file signatures and syncing checkpoints.
- Modified tests/unit/test_utils.py to add unit tests covering checkpoint utilities.
---
 src/stamp/utils/checkpoint.py | 36 +++++++++++++++++++++++++++++++++++
 tests/unit/test_utils.py      | 12 ++++++++++++
 2 files changed, 48 insertions(+)
 create mode 100644 src/stamp/utils/checkpoint.py

diff --git a/src/stamp/utils/checkpoint.py b/src/stamp/utils/checkpoint.py
new file mode 100644
index 0000000..08f2f1f
--- /dev/null
+++ b/src/stamp/utils/checkpoint.py
@@ -0,0 +1,36 @@
+'''
+STAMP: fingerprint-gated working directories for resumable commands
+'''
+
+# Import external dependencies
+import hashlib, json, shutil
+from pathlib import Path
+
+# Import internal STAMP objects
+from stamp.utils.log import log
+
+# _CHECKPOINT_NAME: fingerprint file stored inside the working directory
+_CHECKPOINT_NAME = 'checkpoint.json'
+
+# file_signature: identity of an input file, content hash when content is True else size and mtime
+def file_signature(path: Path, *, content: bool = False) -> str:
+    if content:
+        return hashlib.sha256(path.read_bytes()).hexdigest()
+    stat = path.stat()
+    return f'{stat.st_size}:{stat.st_mtime_ns}'
+
+# sync_checkpoint: keep work_dir when its stored fingerprint matches exactly, else wipe it; returns True when resuming
+def sync_checkpoint(work_dir: Path, fingerprint: dict) -> bool:
+    path = work_dir / _CHECKPOINT_NAME
+    try:
+        stored = json.loads(path.read_text())
+    except (FileNotFoundError, json.JSONDecodeError):
+        stored = None
+    if stored == fingerprint:
+        return True
+    if work_dir.exists() and any(work_dir.iterdir()):
+        log.warning(f'{work_dir} was produced by a different command or inputs, starting fresh')
+    shutil.rmtree(work_dir, ignore_errors=True)
+    work_dir.mkdir(parents=True)
+    path.write_text(json.dumps(fingerprint, indent=2, sort_keys=True))
+    return False
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 544d209..1f8e65b 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -336,3 +336,15 @@ def test_run_parallel_ordered_matches_sequential(self):
         sequential = run_parallel_ordered(items, _square, max_workers=1, label='sq')
         batch = run_parallel_ordered(items, _square, max_workers=4, label='sq')
         assert sequential == batch
+
+# TestCheckpoint: class containing unit tests for src/stamp/utils/checkpoint.py
+class TestCheckpoint:
+    def test_matching_fingerprint_resumes_and_mismatch_wipes(self, tmp_path: Path) -> None:
+        from stamp.utils.checkpoint import sync_checkpoint
+        raw = tmp_path / 'raw'
+        assert sync_checkpoint(raw, {'box': 11}) is False
+        (raw / 'kept.npy').write_text('x')
+        assert sync_checkpoint(raw, {'box': 11}) is True
+        assert (raw / 'kept.npy').exists()
+        assert sync_checkpoint(raw, {'box': 13}) is False
+        assert not (raw / 'kept.npy').exists()

From 2dead152d5f0f74e221b0ca019b9435c0ab4d5de Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 10:18:40 +0100
Subject: [PATCH 16/21] feat(classify): check for classify checkpoint

- Modified commands/classify.py to create a command fingerprint (containing all analysis options) and check for checkpoint with matching signature.
---
 src/stamp/commands/classify.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py
index 935b6cf..74d4885 100644
--- a/src/stamp/commands/classify.py
+++ b/src/stamp/commands/classify.py
@@ -17,6 +17,7 @@
     reduce_and_cluster,
     reduce_and_cluster_shared,
 )
+from stamp.utils.checkpoint import file_signature, sync_checkpoint
 from stamp.classify.extract import extract_particle_set
 from stamp.classify.features import build_feature_matrix
 from stamp.backends.base import ToolCommand
@@ -80,6 +81,30 @@ def run_classify(
         log.debug(f'Using odd box voxel for centred particles: {box_voxels - 1} -> {box_voxels}')    
     log.progress(f'Using box voxels: {box_voxels}')
 
+    # 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())},
+        'voxel_size_angstrom': voxel_size_angstrom,
+        'box_voxels': box_voxels,
+        'n_radial_bins': n_radial_bins,
+        'azimuthal_modes': azimuthal_modes,
+        'n_azimuthal_samples': n_azimuthal_samples,
+        'min_radius_fraction': min_radius_fraction,
+        'method': method,
+        'min_cluster_size': min_cluster_size,
+        'n_clusters': n_clusters,
+        'n_components': n_components,
+        'strict_halfset_independence': strict_halfset_independence,
+        'inplane_alignment': inplane_alignment,
+        'inplane_angular_step_degrees': inplane_angular_step_degrees,
+        'inplane_iterations': inplane_iterations,
+        'random_state': random_state,
+    }
+    if sync_checkpoint(output_dir / 'raw', fingerprint):
+        log.info('Found matching partial output, resuming')
+
     log.progress('Starting subvolume extraction')
     subvols, kept, skipped = extract_particle_set(
         particle_set.particles, {k: str(v) for k, v in tomogram_paths.items()}, box_voxels,

From 9038885ae9f15e216fe3e242955043e2bf9eb39c Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 10:45:27 +0100
Subject: [PATCH 17/21] feat(classify): resume extraction

- Modified classify/extract.py to resume extraction from disk when files present.
---
 src/stamp/classify/extract.py | 77 +++++++++++++++++++++++++++--------
 1 file changed, 61 insertions(+), 16 deletions(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index ac547c6..7f6fb08 100644
--- a/src/stamp/classify/extract.py
+++ b/src/stamp/classify/extract.py
@@ -3,7 +3,7 @@
 '''
 
 # Import external dependencies
-import mrcfile, numpy as np, shutil
+import json, mrcfile, numpy as np, os, shutil
 from dataclasses import dataclass
 from pathlib import Path
 from scipy.ndimage import map_coordinates
@@ -23,6 +23,7 @@ class _ExtractTomogramJob:
     tomogram_path: str
     box_voxels: int
     cache_dir: Path
+    subtracted: bool = False
 
 # _SubtractTomogramJob: one tomogram's membrane subtraction, written to disk for extraction to mmap
 @dataclass(frozen=True)
@@ -95,8 +96,40 @@ def _subtract_one_tomogram(job: _SubtractTomogramJob) -> Path:
         out.voxel_size = voxel_size
     return job.output_path
 
+# _marker_path: per-tomogram completion marker, written only after its subvolume cache is in place
+def _marker_path(cache_dir: Path, tomogram_id: str) -> Path:
+    return cache_dir / f'{tomogram_id}.json'
+
+# _write_marker: record which particles were kept or skipped so an aborted run can resume without recomputing
+def _write_marker(job: _ExtractTomogramJob, kept: list[Particle], skipped: list[Particle]) -> None:
+    marker = {
+        'box_voxels': job.box_voxels,
+        'subtracted': job.subtracted,
+        'kept': [particle.particle_id for particle in kept],
+        'skipped': [particle.particle_id for particle in skipped],
+    }
+    path = _marker_path(job.cache_dir, job.tomogram_id)
+    temporary = path.with_name(path.name + '.tmp')
+    temporary.write_text(json.dumps(marker))
+    os.replace(temporary, path)  # atomic, a partial marker never looks complete
+
+# _read_marker: (kept, skipped, subtracted) from a valid marker for this group, else None
+def _read_marker(cache_dir: Path, tomogram_id: str, group: list[Particle], box_voxels: int) -> tuple[list[Particle], list[Particle], bool] | None:
+    try:
+        marker = json.loads(_marker_path(cache_dir, tomogram_id).read_text())
+    except (FileNotFoundError, json.JSONDecodeError):
+        return None
+    by_id = {particle.particle_id: particle for particle in group}
+    ids = marker.get('kept', []) + marker.get('skipped', [])
+    if marker.get('box_voxels') != box_voxels or len(ids) != len(by_id) or set(ids) != set(by_id):
+        return None
+    if marker['kept'] and not (cache_dir / f'{tomogram_id}.npy').is_file():
+        return None
+    return [by_id[i] for i in marker['kept']], [by_id[i] for i in marker['skipped']], bool(marker.get('subtracted'))
+
 # _extract_one_tomogram: worker entry point to mmaps one tomogram and extracts subvolumes for its particles
 def _extract_one_tomogram(job: _ExtractTomogramJob) -> tuple[list[Particle], list[Particle]]:
+    _marker_path(job.cache_dir, job.tomogram_id).unlink(missing_ok=True)  # invalidate any stale marker before rewriting
     kept: list[Particle] = []
     skipped: list[Particle] = []
     subvolumes: list[np.ndarray] = []
@@ -110,9 +143,14 @@ def _extract_one_tomogram(job: _ExtractTomogramJob) -> tuple[list[Particle], lis
                 continue
             subvolumes.append(extract_subvolume(tomogram, particle.position, particle.orientation, job.box_voxels))
             kept.append(particle)
+    job.cache_dir.mkdir(parents=True, exist_ok=True)
     if subvolumes:
-        job.cache_dir.mkdir(parents=True, exist_ok=True)
-        np.save(job.cache_dir / f'{job.tomogram_id}.npy', np.stack(subvolumes))
+        final = job.cache_dir / f'{job.tomogram_id}.npy'
+        temporary = final.with_name(final.name + '.tmp')
+        with open(temporary, 'wb') as handle:
+            np.save(handle, np.stack(subvolumes))
+        os.replace(temporary, final)  # atomic, a truncated cache never looks complete
+    _write_marker(job, kept, skipped)
     return kept, skipped
 
 # extract_particle_set: extract subvolumes for a list of particles, cached per tomogram
@@ -133,8 +171,12 @@ def extract_particle_set(
     kept: list[Particle] = []
     skipped: list[Particle] = []
 
-    # subtract membranes to cached mrc files
+    # subtract membranes to cached mrc files, any leftovers from an aborted run are recomputed
     subtracted_dir = cache_dir / '.subtracted'
+    shutil.rmtree(subtracted_dir, ignore_errors=True)
+    result_by_tomogram: dict[str, tuple[list[Particle], list[Particle]]] = {}
+    subtracted_ids: set[str] = set()
+    n_resumed = 0
     subtract_jobs: list[_SubtractTomogramJob] = []
     jobs: list[_ExtractTomogramJob] = []
     for tomogram_id, group in by_tomogram.items():
@@ -143,13 +185,21 @@ def extract_particle_set(
             log.warning(f'{tomogram_id}: no matching raw tomogram path, skipping {len(group)} particle(s)')
             skipped.extend(group)
             continue
+        resumed = _read_marker(cache_dir, tomogram_id, group, box_voxels)
+        if resumed is not None:
+            result_by_tomogram[tomogram_id] = (resumed[0], resumed[1])
+            if resumed[2]:
+                subtracted_ids.add(tomogram_id)
+            n_resumed += 1
+            continue
         segmentation_path = segmentation_paths.get(tomogram_id)
         if segmentation_path is not None:
             subtract_jobs.append(_SubtractTomogramJob(tomogram_id, path, segmentation_path, subtracted_dir / f'{tomogram_id}.mrc'))
         else:
             jobs.append(_ExtractTomogramJob(tomogram_id, group, path, box_voxels, cache_dir))
 
-    result_by_tomogram: dict[str, tuple[list[Particle], list[Particle]]] = {}
+    if n_resumed:
+        log.info(f'Resuming extraction: {n_resumed}/{len(by_tomogram)} tomogram(s) already extracted')
 
     def _on_error(job, exc: Exception) -> None:
         if isinstance(exc, StampValidationError):
@@ -165,15 +215,13 @@ def _on_subtract_error(job: _SubtractTomogramJob, exc: Exception) -> None:
         log.warning(f'{job.tomogram_id}: {exc}; extracting from the raw tomogram without membrane subtraction')
         jobs.append(_ExtractTomogramJob(job.tomogram_id, by_tomogram[job.tomogram_id], job.tomogram_path, box_voxels, cache_dir))
 
-    subtracted_ids: set[str] = set()
-
     def _on_subtracted(job: _SubtractTomogramJob, output_path: Path) -> None:
-        subtracted_ids.add(job.tomogram_id)
-        jobs.append(_ExtractTomogramJob(job.tomogram_id, by_tomogram[job.tomogram_id], str(output_path), box_voxels, cache_dir))
+        jobs.append(_ExtractTomogramJob(job.tomogram_id, by_tomogram[job.tomogram_id], str(output_path), box_voxels, cache_dir, subtracted=True))
 
     def _on_success(job: _ExtractTomogramJob, result: tuple[list[Particle], list[Particle]]) -> None:
         result_by_tomogram[job.tomogram_id] = result
-        if Path(job.tomogram_path).parent == subtracted_dir:
+        if job.subtracted:
+            subtracted_ids.add(job.tomogram_id)
             Path(job.tomogram_path).unlink(missing_ok=True)  # free disk as soon as a tomogram is done
 
     try:
@@ -183,14 +231,11 @@ def _on_success(job: _ExtractTomogramJob, result: tuple[list[Particle], list[Par
         shutil.rmtree(subtracted_dir, ignore_errors=True)
 
     index: list[tuple[str, int]] = []
-    for job in sorted(jobs, key=lambda job: job.tomogram_id):
-        result = result_by_tomogram.get(job.tomogram_id)
-        if result is None:
-            continue
-        tomogram_kept, tomogram_skipped = result
+    for tomogram_id in sorted(result_by_tomogram):
+        tomogram_kept, tomogram_skipped = result_by_tomogram[tomogram_id]
         kept.extend(tomogram_kept)
         skipped.extend(tomogram_skipped)
-        index.extend((job.tomogram_id, row) for row in range(len(tomogram_kept)))
+        index.extend((tomogram_id, row) for row in range(len(tomogram_kept)))
 
     log.debug(f'extract_particle_set: {len(kept)} kept, {len(skipped)} skipped')
     n_subtracted = sum(particle.tomogram_id in subtracted_ids for particle in kept)

From a79c09fd08d4dd8b2ad98380d9d79b2f93bfb463 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 10:59:44 +0100
Subject: [PATCH 18/21] feat(utils): add checkpointing utilities

- Modified utils/checkpoint.py to add additional checkpointing utilities.
---
 src/stamp/utils/checkpoint.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/stamp/utils/checkpoint.py b/src/stamp/utils/checkpoint.py
index 08f2f1f..e4e1c45 100644
--- a/src/stamp/utils/checkpoint.py
+++ b/src/stamp/utils/checkpoint.py
@@ -3,7 +3,7 @@
 '''
 
 # Import external dependencies
-import hashlib, json, shutil
+import hashlib, json, numpy as np, os, shutil
 from pathlib import Path
 
 # Import internal STAMP objects
@@ -19,6 +19,21 @@ def file_signature(path: Path, *, content: bool = False) -> str:
     stat = path.stat()
     return f'{stat.st_size}:{stat.st_mtime_ns}'
 
+# atomic_save_npz: write arrays via a temporary file so a partial file never looks complete
+def atomic_save_npz(path: Path, **arrays: np.ndarray) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(path.name + '.tmp')
+    with open(temporary, 'wb') as handle:
+        np.savez(handle, **arrays)
+    os.replace(temporary, path)
+
+# atomic_write_json: write a json file via a temporary file so a partial file never looks complete
+def atomic_write_json(path: Path, payload: object) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(path.name + '.tmp')
+    temporary.write_text(json.dumps(payload))
+    os.replace(temporary, path)
+
 # sync_checkpoint: keep work_dir when its stored fingerprint matches exactly, else wipe it; returns True when resuming
 def sync_checkpoint(work_dir: Path, fingerprint: dict) -> bool:
     path = work_dir / _CHECKPOINT_NAME

From 1e27cb0545be3570612b8d981bf18ee191591945 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 11:10:57 +0100
Subject: [PATCH 19/21] feat(classify): resume feature/alignment processes

- Modified classify/align.py to resume from checkpoint where possible.
- Modified classify/features.py to resume from checkpoint where possible.
- Modified commands/classify.py to pass cache/checkpoint directories through to feature extraction/alignment functions.
---
 src/stamp/classify/align.py    | 40 ++++++++++++++++++++++---
 src/stamp/classify/features.py | 54 ++++++++++++++++++++++++++--------
 src/stamp/commands/classify.py |  7 +++--
 3 files changed, 82 insertions(+), 19 deletions(-)

diff --git a/src/stamp/classify/align.py b/src/stamp/classify/align.py
index c2539a0..3a1b1fa 100644
--- a/src/stamp/classify/align.py
+++ b/src/stamp/classify/align.py
@@ -4,11 +4,15 @@
 
 # Import external dependencies
 import numpy as np
+import json
 from dataclasses import dataclass
+from pathlib import Path
 from scipy.ndimage import rotate
 
 # Import internal STAMP objects
 from stamp.schemas.subvolumes import Subvolumes
+from stamp.utils.checkpoint import atomic_write_json
+from stamp.utils.log import log
 from stamp.utils.parallel import run_parallel
 
 # _ROLL_AXES: box axes spanning the plane perpendicular to the membrane normal
@@ -48,6 +52,7 @@ def _best_angle(
 # _AlignClusterJob: one cluster's subvolumes
 @dataclass(frozen=True)
 class _AlignClusterJob:
+    cluster_id: str
     member_indices: list[int]
     member_subvolumes: np.ndarray
     mask: np.ndarray
@@ -64,7 +69,19 @@ def _align_one_cluster(job: _AlignClusterJob) -> dict[int, float]:
         current = {local: _best_angle(job.member_subvolumes[local], reference_masked, job.mask, job.angles) for local in current}
     return {job.member_indices[local]: angle for local, angle in current.items()}
 
-# align_inplane: per-cluster iterative azimuthal alignment; returns {particle_index: angle_degrees}
+# _load_cluster_checkpoint: saved angles for a cluster, or None when absent or computed for different members
+def _load_cluster_checkpoint(checkpoint_dir: Path | None, cluster_id: str, members: list[int]) -> dict[int, float] | None:
+    if checkpoint_dir is None:
+        return None
+    try:
+        saved = json.loads((checkpoint_dir / f'{cluster_id}.json').read_text())
+    except (OSError, json.JSONDecodeError):
+        return None
+    if saved.get('members') != members:
+        return None
+    return dict(zip(members, saved['angles']))
+
+# align_inplane: per-cluster iterative azimuthal alignment, each finished cluster checkpointed in checkpoint_dir when given; returns {particle_index: angle_degrees}
 def align_inplane(
     subvols: Subvolumes,
     cluster_ids: list[str],
@@ -72,6 +89,7 @@ def align_inplane(
     angular_step_degrees: float = 10.0,
     iterations: int = 3,
     n_workers: int = 1,
+    checkpoint_dir: Path | None = None,
 ) -> dict[int, float]:
     if len(subvols) == 0:
         return {}
@@ -79,13 +97,27 @@ def align_inplane(
     mask = _annulus_mask(box_voxels)
     angles = np.arange(0.0, 360.0, angular_step_degrees)
     jobs = []
+    resolved: dict[int, float] = {}
+    n_resumed = 0
     for cluster_id in sorted(set(cluster_ids)):
         if cluster_id == 'noise':
             continue
         members = [index for index, cid in enumerate(cluster_ids) if cid == cluster_id]
+        saved = _load_cluster_checkpoint(checkpoint_dir, cluster_id, members)
+        if saved is not None:
+            resolved.update(saved)
+            n_resumed += 1
+            continue
         member_subvolumes = np.stack([subvols[i] for i in members])  # bounded to one cluster's members, not the whole dataset
-        jobs.append(_AlignClusterJob(members, member_subvolumes, mask, angles, iterations))
-    resolved: dict[int, float] = {}
-    for cluster_result in run_parallel(jobs, _align_one_cluster, max_workers=n_workers, label='inplane-align'):
+        jobs.append(_AlignClusterJob(cluster_id, members, member_subvolumes, mask, angles, iterations))
+    if n_resumed:
+        log.info(f'Resuming in-plane alignment: {n_resumed} cluster(s) loaded from checkpoint')
+
+    # on_success: checkpoint each cluster as it finishes so an abort only loses in-flight clusters
+    def _on_success(job: _AlignClusterJob, cluster_result: dict[int, float]) -> None:
         resolved.update(cluster_result)
+        if checkpoint_dir is not None:
+            atomic_write_json(checkpoint_dir / f'{job.cluster_id}.json', {'members': job.member_indices, 'angles': [cluster_result[i] for i in job.member_indices]})
+
+    run_parallel(jobs, _align_one_cluster, max_workers=n_workers, label='inplane-align', on_success=_on_success)
     return resolved
\ No newline at end of file
diff --git a/src/stamp/classify/features.py b/src/stamp/classify/features.py
index 5a7b802..75f7e01 100644
--- a/src/stamp/classify/features.py
+++ b/src/stamp/classify/features.py
@@ -7,10 +7,12 @@
 from concurrent.futures import ProcessPoolExecutor
 from contextlib import nullcontext
 from functools import partial
+from pathlib import Path
 from scipy.ndimage import map_coordinates
 
 # Import internal STAMP objects
 from stamp.schemas.subvolumes import Subvolumes
+from stamp.utils.checkpoint import atomic_save_npz
 from stamp.utils.log import get_worker_log_config, init_worker_logging, log
 from stamp.utils.parallel import run_parallel_ordered
 
@@ -105,7 +107,19 @@ def _standardise_columns(matrix: np.ndarray) -> np.ndarray:
     stds[stds == 0.0] = 1.0
     return (matrix - means) / stds
 
-# build_feature_matrix: feature matrix for a Subvolumes instance, read in per-tomogram batches
+# _load_cached_magnitudes: one tomogram's cached magnitudes, or None when absent or built for different rows
+def _load_cached_magnitudes(cache_dir: Path | None, tomogram_id: str, rows: list[int]) -> np.ndarray | None:
+    if cache_dir is None:
+        return None
+    try:
+        with np.load(cache_dir / f'{tomogram_id}.npz') as cached:
+            if np.array_equal(cached['rows'], rows):
+                return cached['magnitudes']
+    except (OSError, ValueError, KeyError):
+        pass
+    return None
+
+# build_feature_matrix: feature matrix for a Subvolumes instance, read in per-tomogram batches, per-tomogram magnitudes cached in cache_dir when given
 def build_feature_matrix(
     subvols: Subvolumes,
     n_radial_bins: int = 12,
@@ -115,6 +129,7 @@ def build_feature_matrix(
     n_workers: int = 1,
     *,
     batch_size: int = 512,
+    cache_dir: Path | None = None,
 ) -> np.ndarray:
     log.debug(f'Building features for {len(subvols)} subvolumes, max_mode={max_azimuthal_mode}')
     if len(subvols) == 0:
@@ -131,20 +146,33 @@ def build_feature_matrix(
         if n_workers > 1
         else nullcontext()
     )
+    by_tomogram: dict[str, list[int]] = {}
+    for local_index, (tomogram_id, _row) in enumerate(subvols.index):
+        by_tomogram.setdefault(tomogram_id, []).append(local_index)
     with pool_context as pool:
         completed = 0
-        for indices, batch in subvols.iter_batches(batch_size):
-            batch_results = run_parallel_ordered(
-                list(batch),
-                worker,
-                max_workers=n_workers,
-                label='feature-extraction',
-                pool=pool,
-                progress_total=total_particles,
-                progress_offset=completed,
-            )
-            completed += len(batch_results)
-            for local_index, result in zip(indices, batch_results):
+        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)
+            if tomogram_magnitudes is None:
+                results: list[np.ndarray] = []
+                for _indices, batch in subvols.take(local_indices).iter_batches(batch_size):
+                    results.extend(run_parallel_ordered(
+                        list(batch),
+                        worker,
+                        max_workers=n_workers,
+                        label='feature-extraction',
+                        pool=pool,
+                        progress_total=total_particles,
+                        progress_offset=completed + len(results),
+                    ))
+                tomogram_magnitudes = np.stack(results)
+                if cache_dir is not None:
+                    atomic_save_npz(cache_dir / f'{tomogram_id}.npz', rows=np.array(rows), magnitudes=tomogram_magnitudes)
+            else:
+                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
     stacked = np.stack(magnitudes)  # (n_particles, n_modes, n_radial_bins, box_voxels)
 
diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py
index 74d4885..f6a01dc 100644
--- a/src/stamp/commands/classify.py
+++ b/src/stamp/commands/classify.py
@@ -123,6 +123,7 @@ def run_classify(
         n_azimuthal_samples=n_azimuthal_samples,
         min_radius_fraction=min_radius_fraction,
         n_workers=n_workers,
+        cache_dir=output_dir / 'raw' / 'features',
     )
     log.debug(f'Feature vector: {features.shape[1]} dimensions (modes 0-{azimuthal_modes})')
 
@@ -264,6 +265,7 @@ def _resolve_and_apply_inplane(
     cluster_ids: list[str],
     align_settings: dict,
     rolled_cache_dir: Path,
+    checkpoint_dir: Path,
 ) -> tuple[Subvolumes, dict[int, float]]:
     if not align_settings['enabled']:
         return subvols, {}
@@ -273,6 +275,7 @@ def _resolve_and_apply_inplane(
         angular_step_degrees=align_settings['angular_step_degrees'],
         iterations=align_settings['iterations'],
         n_workers=align_settings['n_workers'],
+        checkpoint_dir=checkpoint_dir,
     )
     rolled = ((index, roll_about_normal(subvols[index], angles.get(index, 0.0))) for index in range(len(subvols)))
     return subvols.write_rolled(rolled, rolled_cache_dir), angles
@@ -293,7 +296,7 @@ def _classify_combined(
 
     _report_clusters(cluster_ids, result.explained_variance_ratio)
 
-    aligned, angles = _resolve_and_apply_inplane(subvols, cluster_ids, align_settings, output_dir / 'raw' / 'subvolumes_rolled')
+    aligned, angles = _resolve_and_apply_inplane(subvols, cluster_ids, align_settings, output_dir / 'raw' / 'subvolumes_rolled', output_dir / 'raw' / 'align' / 'combined')
 
     averages_by_half: dict[str, dict[str, tuple[np.ndarray, int]]] = {}
     for half_set in (HalfSet.A, HalfSet.B):
@@ -349,7 +352,7 @@ def _classify_strict(
         ('B', half_b, indices_b, result_b, matches),
     ):
         cluster_ids = [label_to_cluster_id(int(label)) for label in result.labels]
-        aligned, local_angles = _resolve_and_apply_inplane(subvols.take(indices), cluster_ids, align_settings, output_dir / 'raw' / f'subvolumes_rolled_{half_label}')
+        aligned, local_angles = _resolve_and_apply_inplane(subvols.take(indices), cluster_ids, align_settings, output_dir / 'raw' / f'subvolumes_rolled_{half_label}', output_dir / 'raw' / 'align' / half_label)
         averages = compute_class_averages(aligned, cluster_ids)
         write_class_averages(averages, output_dir / 'class_averages', voxel_size_angstrom, half_label)
         result_by_half[half_label] = result

From ab6d26b2abb619e5de76fdc60a642953b73af187 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 12:03:57 +0100
Subject: [PATCH 20/21] test(classify): update classification tests

- Modified tests/integration/test_classify_command.py to add tests for resuming from checkpointed run.
- Modified tests/unit/test_classify.py to add unit tests for resuming different modules from checkpoints.
---
 tests/integration/test_classify_command.py | 26 +++++++++
 tests/unit/test_classify.py                | 65 ++++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/tests/integration/test_classify_command.py b/tests/integration/test_classify_command.py
index ae0e579..221847d 100644
--- a/tests/integration/test_classify_command.py
+++ b/tests/integration/test_classify_command.py
@@ -129,6 +129,32 @@ def guarded_stack(arrays, *args, **kwargs):
             return real_stack(arrays, *args, **kwargs)
         monkeypatch.setattr(np, 'stack', guarded_stack)
 
+    def test_classify_resumes_from_matching_partial_output(self, tmp_path: Path) -> None:
+        '''A second identical run reuses extraction, features and alignment; a changed command starts fresh.'''
+        particle_set_path, raw_dir = _build_dataset(tmp_path)
+        argv = [
+            '-vv', '-d', str(tmp_path), 'classify',
+            '--particles', str(particle_set_path), '--raw-dir', str(raw_dir), '--out-dir', str(tmp_path),
+            '--voxel-size-a', '10.0', '--box-length-a', '200.0', '--method', 'kmeans',
+            '--n-clusters', '2', '--n-components', '10', '--keep-raw',
+        ]
+        assert runner.invoke(stamp_app, argv).exit_code == 0
+        output_dir = tmp_path / 'stamp' / 'classify'
+        first = (output_dir / 'class_assignments.json').read_text()
+
+        result = runner.invoke(stamp_app, argv)
+        assert result.exit_code == 0, result.output
+        assert 'Resuming extraction' in result.output
+        assert 'features loaded from checkpoint' in result.output
+        assert 'Resuming in-plane alignment' in result.output
+        assert (output_dir / 'class_assignments.json').read_text() == first
+
+        changed = argv[:argv.index('--n-clusters')] + ['--n-clusters', '3'] + argv[argv.index('--n-clusters') + 2:]
+        result = runner.invoke(stamp_app, changed)
+        assert result.exit_code == 0, result.output
+        assert 'starting fresh' in result.output
+        assert 'Resuming extraction' not in result.output
+
     def test_classify_pipeline_never_stacks_the_whole_dataset(self, tmp_path, monkeypatch):
         '''extract/features/align/average batch per-tomogram, instead of stacking the whole dataset.'''
         n_tomograms = 4
diff --git a/tests/unit/test_classify.py b/tests/unit/test_classify.py
index 3817ab9..c333aae 100644
--- a/tests/unit/test_classify.py
+++ b/tests/unit/test_classify.py
@@ -167,6 +167,33 @@ def test_shape_mismatch_falls_back_to_raw_tomogram(self, tmp_path, caplog) -> No
         assert 'without membrane subtraction' in caplog.text
         assert '1 (raw density)' in caplog.text
 
+    def test_aborted_extraction_resumes_without_recomputing(self, tmp_path, monkeypatch) -> None:
+        from stamp.classify import extract as extract_module
+        particles, paths = [], {}
+        for name in ('a', 'b'):
+            with mrcfile.new(tmp_path / f'{name}.mrc', overwrite=True) as mrc:
+                mrc.set_data(np.random.default_rng(0).normal(size=(40, 40, 40)).astype(np.float32))
+            paths[name] = str(tmp_path / f'{name}.mrc')
+            particles.append(Particle(particle_id=f'p_{name}', tomogram_id=name, position=(20.0, 20.0, 20.0), orientation=(1.0, 0.0, 0.0, 0.0), source_picker='test', confidence=1.0, half_set=HalfSet.A))
+        cache = tmp_path / 'cache'
+        full, kept_full, _ = extract_particle_set(particles, paths, 11, cache_dir=cache)
+
+        # simulate an abort during tomogram b: its marker never got written
+        (cache / 'b.json').unlink()
+        (cache / 'b.npy').unlink()
+        calls: list[str] = []
+        original = extract_module._extract_one_tomogram
+        monkeypatch.setattr(extract_module, '_extract_one_tomogram', lambda job: (calls.append(job.tomogram_id), original(job))[1])
+        resumed, kept_resumed, _ = extract_particle_set(particles, paths, 11, cache_dir=cache)
+        assert calls == ['b']
+        assert [p.particle_id for p in kept_resumed] == [p.particle_id for p in kept_full]
+        assert np.array_equal(resumed[1], full[1])
+
+        # a different box size invalidates the markers
+        calls.clear()
+        extract_particle_set(particles, paths, 13, cache_dir=cache)
+        assert sorted(calls) == ['a', 'b']
+
 # TestFeatures: class containing unit tests for src/stamp/classify/features.py
 class TestFeatures:
     def test_rotational_average_shape(self) -> None:
@@ -293,6 +320,22 @@ def test_degenerate_rows_are_all_zero(self, tmp_path: Path):
         features = build_feature_matrix(subvols, n_radial_bins=4, max_azimuthal_mode=0)
         assert not features[0].any() and features[1].any()
 
+    def test_feature_checkpoint_skips_recomputation(self, tmp_path: Path, monkeypatch) -> None:
+        from stamp.classify import features as features_module
+        subvols = _subvols_from_array(tmp_path / 'sv', np.random.default_rng(0).random((6, 11, 11, 11)))
+        cache = tmp_path / 'features'
+        first = build_feature_matrix(subvols, n_radial_bins=4, cache_dir=cache)
+
+        def _fail(*args, **kwargs):
+            raise AssertionError('features recomputed despite checkpoint')
+
+        monkeypatch.setattr(features_module, 'azimuthal_magnitudes', _fail)
+        assert np.array_equal(build_feature_matrix(subvols, n_radial_bins=4, cache_dir=cache), first)
+
+        # a checkpoint built for different rows is ignored
+        with pytest.raises(AssertionError):
+            build_feature_matrix(subvols.take([0, 1, 2]), n_radial_bins=4, cache_dir=cache)
+
 # TestCluster: class containing unit tests for src/stamp/classify/cluster.py
 class TestCluster:
     def test_hdbscan_recovers_three_groups(self) -> None:
@@ -386,6 +429,28 @@ def test_recovers_planted_rolls(self, tmp_path: Path):
         for frame in aligned:
             assert np.corrcoef(frame.ravel(), reference.ravel())[0, 1] > 0.98
 
+    def test_alignment_checkpoint_resumes_per_cluster(self, tmp_path: Path, monkeypatch) -> None:
+        from stamp.classify import align as align_module
+        stack = np.random.default_rng(0).random((6, 15, 15, 15))
+        subvols = _subvols_from_array(tmp_path / 'sv', stack)
+        cluster_ids = ['c00', 'c00', 'c00', 'c01', 'c01', 'c01']
+        checkpoint = tmp_path / 'align'
+        first = align_inplane(subvols, cluster_ids, checkpoint_dir=checkpoint)
+        assert sorted(p.name for p in checkpoint.glob('*.json')) == ['c00.json', 'c01.json']
+
+        # simulate an abort before c01 finished
+        (checkpoint / 'c01.json').unlink()
+        aligned_clusters: list[list[int]] = []
+        original = align_module._align_one_cluster
+        monkeypatch.setattr(align_module, '_align_one_cluster', lambda job: (aligned_clusters.append(job.member_indices), original(job))[1])
+        assert align_inplane(subvols, cluster_ids, checkpoint_dir=checkpoint) == first
+        assert aligned_clusters == [[3, 4, 5]]
+
+        # a cluster whose membership changed is recomputed
+        aligned_clusters.clear()
+        align_inplane(subvols, ['c00', 'c00', 'c01', 'c01', 'c01', 'c01'], checkpoint_dir=checkpoint)
+        assert sorted(aligned_clusters) == [[0, 1], [2, 3, 4, 5]]
+
     def test_noise_cluster_skipped(self, tmp_path: Path):
         '''Noise rows get no angle.'''
         stack = np.random.default_rng(0).random((4, 15, 15, 15))

From 2eb5191532255ec447ca7fa9a37701770fd639ca Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Tue, 29 Sep 2026 12:05:54 +0100
Subject: [PATCH 21/21] chore: update stamp version to v0.12.0

---
 pyproject.toml | 2 +-
 uv.lock        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index f5dc65c..00a0f39 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "stamp"
-version = "0.11.2"
+version = "0.12.0"
 description = "A workflow for Sub-Tomogram Averaging Membrane Proteins"
 readme = "README.md"
 authors = [
diff --git a/uv.lock b/uv.lock
index 9dd97bd..25b3090 100644
--- a/uv.lock
+++ b/uv.lock
@@ -786,7 +786,7 @@ wheels = [
 
 [[package]]
 name = "stamp"
-version = "0.11.2"
+version = "0.12.0"
 source = { editable = "." }
 dependencies = [
     { name = "loguru" },