STAMP v0.12.0 release

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

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 = [

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,28 +69,55 @@ 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],
     *,
     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 {}
     box_voxels = subvols.box_voxels
     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

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,38 +171,74 @@ 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)
         if path is None:
             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