STAMP v0.11.0 release

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stamp"
-version = "0.10.0"
+version = "0.11.0"
 description = "A workflow for Sub-Tomogram Averaging Membrane Proteins"
 readme = "README.md"
 authors = [

src/stamp/classify/align.py

@@ -8,6 +8,7 @@
 from scipy.ndimage import rotate
 
 # Import internal STAMP objects
+from stamp.schemas.subvolumes import Subvolumes
 from stamp.utils.parallel import run_parallel
 
 # _ROLL_AXES: box axes spanning the plane perpendicular to the membrane normal
@@ -63,27 +64,28 @@ 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_inplace: per-cluster iterative azimuthal alignment; returns {subvolume_row_index: angle_degrees}
+# align_inplane: per-cluster iterative azimuthal alignment; returns {particle_index: angle_degrees}
 def align_inplane(
-    subvolumes: np.ndarray,
+    subvols: Subvolumes,
     cluster_ids: list[str],
     *,
     angular_step_degrees: float = 10.0,
     iterations: int = 3,
     n_workers: int = 1,
 ) -> dict[int, float]:
-    if subvolumes.shape[0] == 0:
+    if len(subvols) == 0:
         return {}
-    box_voxels = subvolumes.shape[-1]
+    box_voxels = subvols.box_voxels
     mask = _annulus_mask(box_voxels)
     angles = np.arange(0.0, 360.0, angular_step_degrees)
     jobs = []
     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]
-        jobs.append(_AlignClusterJob(members, subvolumes[members], mask, angles, iterations))
+        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'):
         resolved.update(cluster_result)
-    return resolved
+    return resolved

src/stamp/classify/average.py

@@ -7,19 +7,28 @@
 from pathlib import Path
 
 # Import internal STAMP objects
+from stamp.schemas.subvolumes import Subvolumes
 from stamp.utils.log import log
 
-# compute_class_averages: mean subvolume per cluster
+# compute_class_averages: mean subvolume per cluster, streamed in batches
 def compute_class_averages(
-    subvolumes: np.ndarray, cluster_ids: list[str]
+    subvols: Subvolumes,
+    cluster_ids: list[str],
+    *,
+    batch_size: int = 512,
 ) -> dict[str, tuple[np.ndarray, int]]:
     averages: dict[str, tuple[np.ndarray, int]] = {}
     for cluster_id in sorted(set(cluster_ids)):
         if cluster_id == 'noise':
             continue
-        mask = np.array([identifier == cluster_id for identifier in cluster_ids])
-        members = subvolumes[mask]
-        averages[cluster_id] = (members.mean(axis=0), int(mask.sum()))
+        members = [index for index, cid in enumerate(cluster_ids) if cid == cluster_id]
+        member_subvols = subvols.take(members)
+        total = np.zeros((subvols.box_voxels,) * 3, dtype=np.float64)
+        count = 0
+        for _indices, batch in member_subvols.iter_batches(batch_size):
+            total += batch.sum(axis=0, dtype=np.float64)
+            count += batch.shape[0]
+        averages[cluster_id] = ((total / count).astype(np.float32), count)
     return averages
 
 # write_class_averages: write each class average as an MRC with its voxel size set

src/stamp/classify/extract.py

@@ -10,10 +10,21 @@
 
 # Import STAMP schema
 from stamp.schemas.particles import Particle
+from stamp.schemas.subvolumes import Subvolumes
 from stamp.utils.errors import StampValidationError
 from stamp.utils.log import log
 from stamp.utils.parallel import run_parallel
 
+# _ExtractTomogramJob: one tomogram's particle group, picklable for ProcessPoolExecutor
+@dataclass(frozen=True)
+class _ExtractTomogramJob:
+    tomogram_id: str
+    group: list[Particle]
+    tomogram_path: str
+    segmentation_path: str | None
+    box_voxels: int
+    cache_dir: 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
@@ -62,18 +73,8 @@ def box_fits_inside(
     position = np.array(position_xyz)
     return bool(np.all(position >= radius) and np.all(position <= extent_xyz - radius))
 
-# _ExtractTomogramJob: one tomogram's particle group, picklable for ProcessPoolExecutor
-@dataclass(frozen=True)
-class _ExtractTomogramJob:
-    tomogram_id: str
-    group: list[Particle]
-    tomogram_path: str
-    segmentation_path: str | None
-    box_voxels: int
-    cache_dir: Path | None = None
-
 # _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], list[np.ndarray]]:
+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)
 
@@ -94,21 +95,22 @@ 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)
-    # cache to disk once tomogram finishes
-    if subvolumes and job.cache_dir is not None:
+    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, 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],
     tomogram_paths: dict[str, str],
     box_voxels: int,
     segmentation_paths: dict[str, str] | None = None,
     n_workers: int = 1,
-    cache_dir: Path | None = None,
-) -> tuple[np.ndarray, list[Particle], list[Particle]]:
+    *,
+    cache_dir: Path,
+) -> tuple[Subvolumes, list[Particle], list[Particle]]:
     by_tomogram: dict[str, list[Particle]] = {}
     for particle in particles:
         by_tomogram.setdefault(particle.tomogram_id, []).append(particle)
@@ -126,11 +128,10 @@ def extract_particle_set(
             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], list[np.ndarray]]] = {}
+    result_by_tomogram: dict[str, tuple[list[Particle], list[Particle]]] = {}
 
-    def _on_success(job: _ExtractTomogramJob, result: tuple[list[Particle], list[Particle], list[np.ndarray]]) -> None:
-        tomogram_kept, tomogram_skipped, subvolumes = result
-        result_by_tomogram[job.tomogram_id] = (tomogram_kept, tomogram_skipped, subvolumes)
+    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:
         if isinstance(exc, StampValidationError):
@@ -141,21 +142,15 @@ def _on_error(job: _ExtractTomogramJob, exc: Exception) -> None:
 
     run_parallel(jobs, _extract_one_tomogram, max_workers=n_workers, label='subvolume-extraction', on_success=_on_success, on_error=_on_error)
 
-    parts: list[np.ndarray] = []
+    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, subvolumes = result
+        tomogram_kept, tomogram_skipped = result
         kept.extend(tomogram_kept)
         skipped.extend(tomogram_skipped)
-        cache_path = job.cache_dir / f'{job.tomogram_id}.npy' if job.cache_dir is not None else None
-        if cache_path is not None and cache_path.exists():
-            parts.append(np.load(cache_path, mmap_mode='r'))
-        elif subvolumes:
-            parts.append(np.stack(subvolumes))
+        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')
-    if not parts:
-        return np.empty((0, box_voxels, box_voxels, box_voxels)), kept, skipped
-    return np.concatenate(parts, axis=0), kept, skipped
+    return Subvolumes(cache_dir, index, box_voxels), kept, skipped

src/stamp/classify/features.py

@@ -8,6 +8,7 @@
 from scipy.ndimage import map_coordinates
 
 # Import internal STAMP objects
+from stamp.schemas.subvolumes import Subvolumes
 from stamp.utils.log import log
 from stamp.utils.parallel import run_parallel_ordered
 
@@ -102,23 +103,29 @@ 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 stack of canonical subvolumes
+# build_feature_matrix: feature matrix for a Subvolumes instance, read in per-tomogram batches
 def build_feature_matrix(
-    subvolumes: np.ndarray,
+    subvols: Subvolumes,
     n_radial_bins: int = 12,
     max_azimuthal_mode: int = 4,
     n_azimuthal_samples: int = 64,
     min_radius_fraction: float = 0.25,
     n_workers: int = 1,
+    *,
+    batch_size: int = 512,
 ) -> np.ndarray:
-    log.debug(f'Building features for {subvolumes.shape[0]} subvolumes, max_mode={max_azimuthal_mode}')
-    if subvolumes.shape[0] == 0:
+    log.debug(f'Building features for {len(subvols)} subvolumes, max_mode={max_azimuthal_mode}')
+    if len(subvols) == 0:
         return np.empty((0, 0))
     if not 0.0 <= min_radius_fraction < 1.0:
         raise ValueError('min_radius_fraction must be between [0, 1)')
 
     worker = partial(azimuthal_magnitudes, n_radial_bins=n_radial_bins, n_azimuthal_samples=n_azimuthal_samples, max_mode=max_azimuthal_mode)
-    magnitudes = run_parallel_ordered(list(subvolumes), worker, max_workers=n_workers, label='feature-extraction')
+    magnitudes: list[np.ndarray | None] = [None] * len(subvols)
+    for indices, batch in subvols.iter_batches(batch_size):
+        batch_results = run_parallel_ordered(list(batch), worker, max_workers=n_workers, label='feature-extraction')
+        for local_index, result in zip(indices, batch_results):
+            magnitudes[local_index] = result
     stacked = np.stack(magnitudes)  # (n_particles, n_modes, n_radial_bins, box_voxels)
 
     n_particles = stacked.shape[0]
@@ -132,4 +139,4 @@ def build_feature_matrix(
         block = stacked[:, mode, first_outer_bin:, :].reshape(n_particles, -1)
         blocks.append(_standardise_columns(block))
 
-    return np.hstack(blocks)
+    return np.hstack(blocks)

src/stamp/commands/classify.py

@@ -24,6 +24,7 @@
 from stamp.run.state import stage_dir
 from stamp.utils.halfset import split_by_half_set
 from stamp.schemas.particles import ClassAssignment, HalfSet, ParticleSet
+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
@@ -76,7 +77,7 @@ def run_classify(
     if box_voxels % 2 == 0:
         box_voxels += 1  # odd box keeps the particle exactly centred
 
-    subvolumes, kept, skipped = extract_particle_set(
+    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,7 +88,7 @@ def run_classify(
     log.info(f'Extracted {len(kept)} subvolumes at box {box_voxels}{" (membrane subtracted)" if segmentation_paths else ""}')
 
     features = build_feature_matrix(
-        subvolumes,
+        subvols,
         n_radial_bins=n_radial_bins,
         max_azimuthal_mode=azimuthal_modes,
         n_azimuthal_samples=n_azimuthal_samples,
@@ -99,8 +100,9 @@ def run_classify(
     degenerate = ~features.any(axis=1)
     if degenerate.any():
         log.warning(f'Dropped {int(degenerate.sum())} particles with a constant/empty subvolume (edge fill)')
+        keep_indices = [i for i, bad in enumerate(degenerate) if not bad]
         features = features[~degenerate]
-        subvolumes = subvolumes[~degenerate]
+        subvols = subvols.take(keep_indices)
         kept = [particle for particle, bad in zip(kept, degenerate) if not bad]
     if not kept:
         raise StampPipelineError('No particles left after dropping degenerate subvolumes')
@@ -127,9 +129,9 @@ 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:
-        assignments, result_by_half, averages_by_half = _classify_strict(kept, subvolumes, features, config, output_dir, voxel_size_angstrom, align_settings)
+        assignments, result_by_half, averages_by_half = _classify_strict(kept, subvols, features, config, output_dir, voxel_size_angstrom, align_settings)
     else:
-        assignments, result_by_half, averages_by_half = _classify_combined(kept, subvolumes, features, config, output_dir, voxel_size_angstrom, align_settings)
+        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'
     assignments_path.write_text(json.dumps([a.model_dump() for a in assignments], indent=2))
@@ -224,26 +226,30 @@ def build_classify_commands(config, output_dir: Path, track: str = 'real') -> li
     argv += ['--plot-format', config.plots.format]
     return [ToolCommand(tool='classify', argv=argv, working_directory=target, output_paths=[target / 'class_averages'])]
 
-# _resolve_and_apply_inplane: estimate per-row azimuth, return rolled subvolumes and {row: angle}
+# _resolve_and_apply_inplane: estimate per-row azimuth, return a rolled subvols and {row: angle}
 def _resolve_and_apply_inplane(
-    subvolumes: np.ndarray, cluster_ids: list[str], align_settings: dict
-) -> tuple[np.ndarray, dict[int, float]]:
+    subvols: Subvolumes,
+    cluster_ids: list[str],
+    align_settings: dict,
+    rolled_cache_dir: Path,
+) -> tuple[Subvolumes, dict[int, float]]:
     if not align_settings['enabled']:
-        return subvolumes, {}
+        return subvols, {}
     angles = align_inplane(
-        subvolumes,
+        subvols,
         cluster_ids,
         angular_step_degrees=align_settings['angular_step_degrees'],
         iterations=align_settings['iterations'],
         n_workers=align_settings['n_workers'],
     )
-    rolled = np.stack([roll_about_normal(volume, angles.get(index, 0.0)) for index, volume in enumerate(subvolumes)])
+    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
     return rolled, angles
 
 # _classify_combined: cluster all particles together, then average each half separately
 def _classify_combined(
     particles,
-    subvolumes,
+    subvols: Subvolumes,
     features,
     config,
     output_dir: Path,
@@ -255,14 +261,15 @@ def _classify_combined(
 
     _report_clusters(cluster_ids, result.explained_variance_ratio)
 
-    aligned, angles = _resolve_and_apply_inplane(subvolumes, cluster_ids, align_settings)
+    aligned, angles = _resolve_and_apply_inplane(subvols, cluster_ids, align_settings, output_dir / 'raw' / 'subvolumes_rolled')
 
     averages_by_half: dict[str, dict[str, tuple[np.ndarray, int]]] = {}
     for half_set in (HalfSet.A, HalfSet.B):
-        mask = np.array([p.half_set == half_set for p in particles])
-        if not mask.any():
+        half_indices = [i for i, p in enumerate(particles) if p.half_set == half_set]
+        if not half_indices:
             continue
-        averages = compute_class_averages(aligned[mask], [cid for cid, keep in zip(cluster_ids, mask) if keep])
+        half_cluster_ids = [cluster_ids[i] for i in half_indices]
+        averages = compute_class_averages(aligned.take(half_indices), half_cluster_ids)
         write_class_averages(averages, output_dir / 'class_averages', voxel_size_angstrom, half_set.value)
         averages_by_half[half_set.value] = averages
 
@@ -281,7 +288,7 @@ def _classify_combined(
 # _classify_strict: cluster each half independently, then match clusters by centroid
 def _classify_strict(
     particles,
-    subvolumes,
+    subvols: Subvolumes,
     features,
     config,
     output_dir: Path,
@@ -290,12 +297,12 @@ def _classify_strict(
 ) -> tuple[list[ClassAssignment], dict[str, ClusteringResult], dict[str, dict[str, tuple[np.ndarray, int]]]]:
     half_a, half_b = split_by_half_set(list(particles))
     index_of = {particle.particle_id: i for i, particle in enumerate(particles)}
-    indices_a = np.array([index_of[p.particle_id] for p in half_a])
-    indices_b = np.array([index_of[p.particle_id] for p in half_b])
-    if indices_a.size == 0 or indices_b.size == 0:
+    indices_a = [index_of[p.particle_id] for p in half_a]
+    indices_b = [index_of[p.particle_id] for p in half_b]
+    if not indices_a or not indices_b:
         raise StampPipelineError('Strict mode needs particles in both half-sets')
 
-    results = reduce_and_cluster_shared(features, {'A': indices_a, 'B': indices_b}, config)
+    results = reduce_and_cluster_shared(features, {'A': np.array(indices_a), 'B': np.array(indices_b)}, config)
     result_a, result_b = results['A'], results['B']
     matches = match_clusters_across_halves(result_a.centroids, result_b.centroids)
     log.debug('Cross-half cluster matching (B -> A, shared PCA space):')
@@ -310,7 +317,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(subvolumes[indices], cluster_ids, align_settings)
+        aligned, local_angles = _resolve_and_apply_inplane(subvols.take(indices), cluster_ids, align_settings, output_dir / 'raw' / f'subvolumes_rolled_{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