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 `
/stamp/classify` or ``)
`--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/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/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/extract.py b/src/stamp/classify/extract.py
index a34394c..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
+import json, mrcfile, numpy as np, os, shutil
from dataclasses import dataclass
from pathlib import Path
from scipy.ndimage import map_coordinates
@@ -21,9 +21,17 @@ class _ExtractTomogramJob:
tomogram_id: str
group: list[Particle]
tomogram_path: str
- segmentation_path: str | None
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)
+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:
@@ -59,9 +67,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
@@ -73,34 +79,80 @@ 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
+
+# _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]]:
- with mrcfile.open(str(job.tomogram_path), permissive=True) as mrc:
- tomogram = np.asarray(mrc.data).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])
-
+ _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] = []
- 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)
+ 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
# extract_particle_set: extract subvolumes for a list of particles, cached per tomogram
def extract_particle_set(
particles: list[Particle],
@@ -119,6 +171,13 @@ def extract_particle_set(
kept: list[Particle] = []
skipped: list[Particle] = []
+ # 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():
path = tomogram_paths.get(tomogram_id)
@@ -126,31 +185,60 @@ 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))
-
- result_by_tomogram: dict[str, tuple[list[Particle], list[Particle]]] = {}
+ 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))
- def _on_success(job: _ExtractTomogramJob, result: tuple[list[Particle], list[Particle]]) -> None:
- result_by_tomogram[job.tomogram_id] = result
+ if n_resumed:
+ log.info(f'Resuming extraction: {n_resumed}/{len(by_tomogram)} tomogram(s) already extracted')
- 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))
+
+ def _on_subtracted(job: _SubtractTomogramJob, output_path: Path) -> None:
+ 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 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:
+ 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:
- 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)
+ 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/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/cli/classify.py b/src/stamp/cli/classify.py
index cdc7c65..5f0c53e 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(
@@ -52,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'),
@@ -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/classify.py b/src/stamp/commands/classify.py
index 3bd19b5..f6a01dc 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
@@ -27,7 +28,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 +62,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 +78,34 @@ 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}')
+ # 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,
segmentation_paths=segmentation_paths, n_workers=n_workers, cache_dir=output_dir / 'raw' / 'subvolumes',
@@ -85,8 +114,8 @@ 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(
subvols,
n_radial_bins=n_radial_bins,
@@ -94,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})')
@@ -114,6 +144,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 +160,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'
@@ -232,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, {}
@@ -241,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
@@ -261,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):
@@ -317,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
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(
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(
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}')
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:
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
diff --git a/src/stamp/utils/checkpoint.py b/src/stamp/utils/checkpoint.py
new file mode 100644
index 0000000..e4e1c45
--- /dev/null
+++ b/src/stamp/utils/checkpoint.py
@@ -0,0 +1,51 @@
+'''
+STAMP: fingerprint-gated working directories for resumable commands
+'''
+
+# Import external dependencies
+import hashlib, json, numpy as np, os, 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}'
+
+# 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
+ 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/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}')
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 6919f39..c333aae 100644
--- a/tests/unit/test_classify.py
+++ b/tests/unit/test_classify.py
@@ -149,6 +149,51 @@ 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
+
+ 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:
@@ -275,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:
@@ -368,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))
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()
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" },