STAMP v0.13.2 release

pyproject.toml

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

src/stamp/classify/align.py

@@ -3,8 +3,7 @@
 '''
 
 # Import external dependencies
-import numpy as np
-import json
+import json, numpy as np
 from dataclasses import dataclass
 from pathlib import Path
 from scipy.ndimage import rotate
@@ -15,6 +14,9 @@
 from stamp.utils.log import log
 from stamp.utils.parallel import run_parallel
 
+# _BATCH_SIZE: subvolumes read per chunk in a cluster worker
+_BATCH_SIZE = 64
+
 # _ROLL_AXES: box axes spanning the plane perpendicular to the membrane normal
 _ROLL_AXES = (0, 1)
 
@@ -49,24 +51,32 @@ def _best_angle(
             best_angle, best_score = float(angle), score
     return best_angle
 
-# _AlignClusterJob: one cluster's subvolumes
+# _AlignClusterJob: one cluster's members (lazy view so no subvolume data is held or pickled)
 @dataclass(frozen=True)
 class _AlignClusterJob:
     cluster_id: str
     member_indices: list[int]
-    member_subvolumes: np.ndarray
+    member_subvolumes: Subvolumes  # local index i <-> member_indices[i]
     mask: np.ndarray
     angles: np.ndarray
     iterations: int
 
-# _align_one_cluster: multiprocessing worker entry point
+# _align_one_cluster: multiprocessing worker entry point using streamed batches
 def _align_one_cluster(job: _AlignClusterJob) -> dict[int, float]:
-    current = {local: 0.0 for local in range(len(job.member_indices))}
+    n_members = len(job.member_indices)
+    current = {local: 0.0 for local in range(n_members)}
     for _ in range(max(job.iterations, 1)):
-        reference = np.mean([roll_about_normal(job.member_subvolumes[local], angle) for local, angle in current.items()], axis=0)
-        reference_masked = reference[job.mask]
+        total = np.zeros((job.member_subvolumes.box_voxels,) * 3, dtype=np.float64)
+        for locals_, batch in job.member_subvolumes.iter_batches(_BATCH_SIZE):
+            for local, subvolume in zip(locals_, batch):
+                total += roll_about_normal(subvolume, current[local])
+        reference_masked = (total / n_members)[job.mask]
         reference_masked = reference_masked - reference_masked.mean()
-        current = {local: _best_angle(job.member_subvolumes[local], reference_masked, job.mask, job.angles) for local in current}
+        updated: dict[int, float] = {}
+        for locals_, batch in job.member_subvolumes.iter_batches(_BATCH_SIZE):
+            for local, subvolume in zip(locals_, batch):
+                updated[local] = _best_angle(subvolume, reference_masked, job.mask, job.angles)
+        current = updated
     return {job.member_indices[local]: angle for local, angle in current.items()}
 
 # _load_cluster_checkpoint: saved angles for a cluster, or None when absent or computed for different members
@@ -108,8 +118,7 @@ def align_inplane(
             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(cluster_id, members, member_subvolumes, mask, angles, iterations))
+        jobs.append(_AlignClusterJob(cluster_id, members, subvols.take(members), mask, angles, iterations))
     if n_resumed:
         log.info(f'Resuming in-plane alignment: {n_resumed} cluster(s) loaded from checkpoint')
 

src/stamp/commands/classify.py

@@ -280,9 +280,8 @@ def _resolve_and_apply_inplane(
         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)))
+    rolled = ((index, roll_about_normal(subvolume, angles.get(index, 0.0))) for indices, batch in subvols.iter_batches() for index, subvolume in zip(indices, batch))
     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(

src/stamp/schemas/subvolumes.py

@@ -55,19 +55,26 @@ def iter_batches(self, batch_size: int = 512) -> Iterator[tuple[list[int], np.nd
                 rows = [self.index[i][1] for i in chunk]
                 yield chunk, np.asarray(mmap[rows])
 
-    # write_rolled: persist (global_index, subvolume) pairs to a new per-tomogram cache, return the resulting store
+    # write_rolled: persist (global_index, subvolume) pairs to a new per-tomogram cache via preallocated memmaps, return the resulting store
     def write_rolled(self, rolled: Iterator[tuple[int, np.ndarray]], new_cache_dir: Path) -> 'Subvolumes':
         new_cache_dir.mkdir(parents=True, exist_ok=True)
-        buffers: dict[str, list[np.ndarray]] = {}
-        order: dict[str, list[int]] = {}
+        row_counts: dict[str, int] = {}
+        new_index: list[tuple[str, int]] = []
+        for tomogram_id, _row in self.index:
+            new_index.append((tomogram_id, row_counts.get(tomogram_id, 0)))
+            row_counts[tomogram_id] = row_counts.get(tomogram_id, 0) + 1
+        shape = (self.box_voxels,) * 3
+        outputs: dict[str, np.ndarray] = {}
+        written = 0
         for global_index, subvolume in rolled:
-            tomogram_id, _row = self.index[global_index]
-            buffers.setdefault(tomogram_id, []).append(subvolume)
-            order.setdefault(tomogram_id, []).append(global_index)
-        new_index: list[tuple[str, int] | None] = [None] * len(self.index)
-        for tomogram_id, subvolumes in buffers.items():
-            np.save(new_cache_dir / f'{tomogram_id}.npy', np.stack(subvolumes))
-            for row, global_index in enumerate(order[tomogram_id]):
-                new_index[global_index] = (tomogram_id, row)
-        assert all(entry is not None for entry in new_index)  # every particle must be rolled exactly once
-        return Subvolumes(new_cache_dir, new_index, self.box_voxels)  # type: ignore[arg-type]
+            tomogram_id, row = new_index[global_index]
+            if tomogram_id not in outputs:
+                if len(outputs) >= _MAX_OPEN_MMAPS:  # flush and close the oldest to bound open handles
+                    outputs.pop(next(iter(outputs))).flush()
+                outputs[tomogram_id] = np.lib.format.open_memmap(new_cache_dir / f'{tomogram_id}.npy', mode='r+' if (new_cache_dir / f'{tomogram_id}.npy').exists() else 'w+', dtype=np.float32, shape=(row_counts[tomogram_id], *shape))
+            outputs[tomogram_id][row] = subvolume
+            written += 1
+        for output in outputs.values():
+            output.flush()
+        assert written == len(self.index)  # every particle must be rolled exactly once
+        return Subvolumes(new_cache_dir, new_index, self.box_voxels)