From dcedb6266727e96bed925b9c28bf9f9b96241fbd Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 14:43:54 +0100
Subject: [PATCH 01/10] perf(classify): drop subvolumes from memory

- Modified classify/extract.py to drop subvolumes from memory after extraction.
---
 src/stamp/classify/extract.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index efe5137..87821e9 100644
--- a/src/stamp/classify/extract.py
+++ b/src/stamp/classify/extract.py
@@ -94,10 +94,11 @@ 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
+    # cache to disk once tomogram finishes, then drop from memory rather than shipping the full array back over the pool
     if subvolumes and job.cache_dir is not None:
         job.cache_dir.mkdir(parents=True, exist_ok=True)
         np.save(job.cache_dir / f'{job.tomogram_id}.npy', np.stack(subvolumes))
+        subvolumes = []
     return kept, skipped, subvolumes
 
 # extract_particle_set: extract subvolumes for a list of particles

From 74bd10be0bc0069224767ca7b5fb3743189858a9 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 14:52:20 +0100
Subject: [PATCH 02/10] feat(schemas): add schema for subvolumes

- Added schemas/subvolumes.py to define a schema for Subvolumes to allow lazy particle access instead of reading all into memory.
- Modified tests/unit/test_schemas.py to add tests for subvolume schema.
---
 src/stamp/schemas/subvolumes.py | 66 +++++++++++++++++++++++++++++++++
 tests/unit/test_schemas.py      | 56 +++++++++++++++++++++++++++-
 2 files changed, 120 insertions(+), 2 deletions(-)
 create mode 100644 src/stamp/schemas/subvolumes.py

diff --git a/src/stamp/schemas/subvolumes.py b/src/stamp/schemas/subvolumes.py
new file mode 100644
index 0000000..cf16343
--- /dev/null
+++ b/src/stamp/schemas/subvolumes.py
@@ -0,0 +1,66 @@
+'''
+STAMP: lazy index over per-tomogram cached subvolume stacks
+'''
+
+# Import external dependencies
+import numpy as np
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+
+# Subvolumes: particle -> (tomogram_id, row) index over mmap'd per-tomogram .npy caches
+@dataclass(frozen=True)
+class Subvolumes:
+    cache_dir: Path
+    index: list[tuple[str, int]]  # index[i] = (tomogram_id, row) for particle i, in extraction order
+    box_voxels: int
+    _mmaps: dict[str, np.ndarray] = field(default_factory=dict, repr=False, compare=False)
+
+    def __len__(self) -> int:
+        return len(self.index)
+
+    # _mmap_for: open (or reuse) the mmap for one tomogram's cache file
+    def _mmap_for(self, tomogram_id: str) -> np.ndarray:
+        cached = self._mmaps.get(tomogram_id)
+        if cached is not None:
+            return cached
+        array = np.load(self.cache_dir / f'{tomogram_id}.npy', mmap_mode='r')
+        self._mmaps[tomogram_id] = array  # mutating the dict's contents, not the frozen field itself
+        return array
+
+    def __getitem__(self, i: int) -> np.ndarray:
+        tomogram_id, row = self.index[i]
+        return np.asarray(self._mmap_for(tomogram_id)[row])
+
+    # take: view over a subset of particles by index, no data copy
+    def take(self, indices: Sequence[int]) -> 'Subvolumes':
+        return Subvolumes(self.cache_dir, [self.index[i] for i in indices], self.box_voxels)
+
+    # iter_batches: yield (global_indices, dense_batch) grouped by tomogram, for disk-local reads
+    def iter_batches(self, batch_size: int = 512) -> Iterator[tuple[list[int], np.ndarray]]:
+        by_tomogram: dict[str, list[int]] = {}
+        for local_index, (tomogram_id, _row) in enumerate(self.index):
+            by_tomogram.setdefault(tomogram_id, []).append(local_index)
+        for tomogram_id, local_indices in by_tomogram.items():
+            mmap = self._mmap_for(tomogram_id)
+            for start in range(0, len(local_indices), batch_size):
+                chunk = local_indices[start : start + batch_size]
+                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
+    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]] = {}
+        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]
diff --git a/tests/unit/test_schemas.py b/tests/unit/test_schemas.py
index 1864923..e44c873 100644
--- a/tests/unit/test_schemas.py
+++ b/tests/unit/test_schemas.py
@@ -3,7 +3,7 @@
 '''
 
 # Import external dependencies
-import pytest
+import numpy as np, pytest
 from datetime import datetime, timezone
 from pathlib import Path
 from pydantic import ValidationError
@@ -18,8 +18,15 @@
     ParticleSet,
 )
 from stamp.schemas.provenance import ProvenanceSidecar
+from stamp.schemas.subvolumes import Subvolumes
 from stamp.utils.errors import StampPipelineError
 
+# _write_cache: helper writing a fake per-tomogram cache of n subvolumes
+def _write_cache(cache_dir, tomogram_id: str, n: int, box_voxels: int, offset: int = 0) -> np.ndarray:
+    data = np.arange(n * box_voxels**3, dtype=np.float32).reshape(n, box_voxels, box_voxels, box_voxels) + offset
+    np.save(cache_dir / f'{tomogram_id}.npy', data)
+    return data
+
 # TestManifestSchema: class containing unit tests for schemas/manifest.py
 class TestManifestSchema:
     def test_tomogram_manifest_defaults_not_decoy(self):
@@ -134,4 +141,49 @@ def test_provenance_sidecar_construction(self):
             timestamp=datetime.now(timezone.utc),
             input_checksums={'tomo001_seg.mrc': 'abc123'},
         )
-        assert sidecar.stage == 'pick'
\ No newline at end of file
+        assert sidecar.stage == 'pick'
+
+class TestSubvolumes:
+    def test_getitem_matches_dense_reference(self, tmp_path):
+        box_voxels = 3
+        a = _write_cache(tmp_path, 'tomoA', 2, box_voxels)
+        b = _write_cache(tmp_path, 'tomoB', 3, box_voxels, offset=1000)
+        index = [('tomoA', 0), ('tomoA', 1), ('tomoB', 0), ('tomoB', 1), ('tomoB', 2)]
+        store = Subvolumes(tmp_path, index, box_voxels)
+        dense = np.concatenate([a, b], axis=0)
+        for i in range(len(store)):
+            np.testing.assert_array_equal(store[i], dense[i])
+
+    def test_take_returns_view_without_copying_cache_files(self, tmp_path):
+        box_voxels = 2
+        _write_cache(tmp_path, 'tomoA', 4, box_voxels)
+        index = [('tomoA', i) for i in range(4)]
+        store = Subvolumes(tmp_path, index, box_voxels)
+        subset = store.take([0, 2])
+        assert len(subset) == 2
+        np.testing.assert_array_equal(subset[0], store[0])
+        np.testing.assert_array_equal(subset[1], store[2])
+
+    def test_iter_batches_covers_every_index_grouped_by_tomogram(self, tmp_path):
+        box_voxels = 2
+        _write_cache(tmp_path, 'tomoA', 5, box_voxels)
+        _write_cache(tmp_path, 'tomoB', 3, box_voxels, offset=500)
+        index = [('tomoA', i) for i in range(5)] + [('tomoB', i) for i in range(3)]
+        store = Subvolumes(tmp_path, index, box_voxels)
+        seen: dict[int, np.ndarray] = {}
+        for indices, batch in store.iter_batches(batch_size=2):
+            for local_index, subvolume in zip(indices, batch):
+                seen[local_index] = subvolume
+        assert set(seen) == set(range(len(store)))
+        for i in seen:
+            np.testing.assert_array_equal(seen[i], store[i])
+
+    def test_write_rolled_produces_a_readable_store(self, tmp_path):
+        box_voxels = 2
+        _write_cache(tmp_path, 'tomoA', 3, box_voxels)
+        index = [('tomoA', i) for i in range(3)]
+        store = Subvolumes(tmp_path, index, box_voxels)
+        rolled = ((i, store[i] * 2) for i in range(len(store)))
+        new_store = store.write_rolled(rolled, tmp_path / 'rolled')
+        for i in range(len(store)):
+            np.testing.assert_array_equal(new_store[i], store[i] * 2)

From fc33f9eecd1515d2b7a6464720b6ae7426a84314 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 15:19:50 +0100
Subject: [PATCH 03/10] refactor(classify): return subvolumes from extract

- Modified classify/extract.py to return Subvolumes instance rather than actual subvolumes.
- Modified tests/unit/test_classify.py to fix test suite following refactor.
---
 src/stamp/classify/extract.py | 56 ++++++++++++++++-------------------
 tests/unit/test_classify.py   |  8 ++---
 2 files changed, 29 insertions(+), 35 deletions(-)

diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py
index 87821e9..a34394c 100644
--- a/src/stamp/classify/extract.py
+++ b/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,22 +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, then drop from memory rather than shipping the full array back over the pool
-    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))
-        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)
@@ -127,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):
@@ -142,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
diff --git a/tests/unit/test_classify.py b/tests/unit/test_classify.py
index e83259a..ef6b528 100644
--- a/tests/unit/test_classify.py
+++ b/tests/unit/test_classify.py
@@ -123,10 +123,10 @@ def test_membrane_voxels_replaced_with_background(self, tmp_path) -> None:
             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, _ = extract_particle_set([particle], {'t': str(tmp_path / 't.mrc')}, 11, segmentation_paths={'t': str(tmp_path / 's.mrc')})
+        subvolumes, kept, _ = 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
         # the box centre sampled the membrane plane; it must now sit near background
-        assert abs(float(subvolumes[0, 5, 5, 5])) < 5.0
+        assert abs(float(subvolumes[0][5, 5, 5])) < 5.0
 
     def test_no_segmentation_leaves_tomogram_untouched(self, tmp_path) -> None:
         shape = (40, 40, 40)
@@ -136,8 +136,8 @@ def test_no_segmentation_leaves_tomogram_untouched(self, tmp_path) -> None:
             mrc.set_data(tomogram)
             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, _ = extract_particle_set([particle], {'t': str(tmp_path / 't.mrc')}, 11)
-        assert float(subvolumes[0, 5, 5, 5]) == 50.0
+        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
 
 # TestFeatures: class containing unit tests for src/stamp/classify/features.py
 class TestFeatures:

From 966e13ef864143fe7253062c8cd1ea70f27e1336 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 15:35:21 +0100
Subject: [PATCH 04/10] refactor(classify): stream feature extraction

- Modified classify/features.py to stream feature extraction from Subvolumes instance.
---
 src/stamp/classify/features.py | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/stamp/classify/features.py b/src/stamp/classify/features.py
index 50c49a8..356bddc 100644
--- a/src/stamp/classify/features.py
+++ b/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,
+    store: 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(store)} subvolumes, max_mode={max_azimuthal_mode}')
+    if len(store) == 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(store)
+    for indices, batch in store.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)
\ No newline at end of file

From 2ebb337e3daa05c3f3e867b716a35561c1776984 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 16:20:01 +0100
Subject: [PATCH 05/10] refactor(classify): batch cluster alignment

- Modified classify/align.py to batch cluster alignment using Subvolumes instance.
---
 src/stamp/classify/align.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/stamp/classify/align.py b/src/stamp/classify/align.py
index 4c4b1e0..011564d 100644
--- a/src/stamp/classify/align.py
+++ b/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,18 +64,18 @@ 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,
+    store: 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(store) == 0:
         return {}
-    box_voxels = subvolumes.shape[-1]
+    box_voxels = store.box_voxels
     mask = _annulus_mask(box_voxels)
     angles = np.arange(0.0, 360.0, angular_step_degrees)
     jobs = []
@@ -82,8 +83,9 @@ def align_inplane(
         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([store[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
\ No newline at end of file

From 2d3e0c2e203b1a3ae84f8b11b6d560cbbaa5b4d4 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 16:22:29 +0100
Subject: [PATCH 06/10] refactor(classify): batch class averaging

- Modified classify/average.py to batch class averaging using a Subvolumes instance.
---
 src/stamp/classify/average.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/stamp/classify/average.py b/src/stamp/classify/average.py
index 76f17c3..ce90b70 100644
--- a/src/stamp/classify/average.py
+++ b/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]
+    store: 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_store = store.take(members)
+        total = np.zeros((store.box_voxels,) * 3, dtype=np.float64)
+        count = 0
+        for _indices, batch in member_store.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

From 4f0b2ecceaf1ca950896aa52e5732492aa7e0d94 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 16:44:26 +0100
Subject: [PATCH 07/10] refactor(classify): batch classification

- Modified commands/classify.py to use Subvolumes-based batching throughout classification.
---
 src/stamp/commands/classify.py | 51 +++++++++++++++++++---------------
 1 file changed, 29 insertions(+), 22 deletions(-)

diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py
index a539f0e..3bd19b5 100644
--- a/src/stamp/commands/classify.py
+++ b/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

From b60e2dbf9343fe761ea7312888ebe2f79b2753ba Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 16:52:58 +0100
Subject: [PATCH 08/10] refactor(classify): standardise Subvolumes names

---
 src/stamp/classify/align.py    |  8 ++++----
 src/stamp/classify/average.py  |  8 ++++----
 src/stamp/classify/features.py | 10 +++++-----
 tests/unit/test_schemas.py     | 34 +++++++++++++++++-----------------
 4 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/src/stamp/classify/align.py b/src/stamp/classify/align.py
index 011564d..c2539a0 100644
--- a/src/stamp/classify/align.py
+++ b/src/stamp/classify/align.py
@@ -66,16 +66,16 @@ def _align_one_cluster(job: _AlignClusterJob) -> dict[int, float]:
 
 # align_inplane: per-cluster iterative azimuthal alignment; returns {particle_index: angle_degrees}
 def align_inplane(
-    store: Subvolumes,
+    subvols: Subvolumes,
     cluster_ids: list[str],
     *,
     angular_step_degrees: float = 10.0,
     iterations: int = 3,
     n_workers: int = 1,
 ) -> dict[int, float]:
-    if len(store) == 0:
+    if len(subvols) == 0:
         return {}
-    box_voxels = store.box_voxels
+    box_voxels = subvols.box_voxels
     mask = _annulus_mask(box_voxels)
     angles = np.arange(0.0, 360.0, angular_step_degrees)
     jobs = []
@@ -83,7 +83,7 @@ def align_inplane(
         if cluster_id == 'noise':
             continue
         members = [index for index, cid in enumerate(cluster_ids) if cid == cluster_id]
-        member_subvolumes = np.stack([store[i] for i in members])  # bounded to one cluster's members, not the whole dataset
+        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'):
diff --git a/src/stamp/classify/average.py b/src/stamp/classify/average.py
index ce90b70..431596d 100644
--- a/src/stamp/classify/average.py
+++ b/src/stamp/classify/average.py
@@ -12,7 +12,7 @@
 
 # compute_class_averages: mean subvolume per cluster, streamed in batches
 def compute_class_averages(
-    store: Subvolumes,
+    subvols: Subvolumes,
     cluster_ids: list[str],
     *,
     batch_size: int = 512,
@@ -22,10 +22,10 @@ def compute_class_averages(
         if cluster_id == 'noise':
             continue
         members = [index for index, cid in enumerate(cluster_ids) if cid == cluster_id]
-        member_store = store.take(members)
-        total = np.zeros((store.box_voxels,) * 3, dtype=np.float64)
+        member_subvols = subvols.take(members)
+        total = np.zeros((subvols.box_voxels,) * 3, dtype=np.float64)
         count = 0
-        for _indices, batch in member_store.iter_batches(batch_size):
+        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)
diff --git a/src/stamp/classify/features.py b/src/stamp/classify/features.py
index 356bddc..3762370 100644
--- a/src/stamp/classify/features.py
+++ b/src/stamp/classify/features.py
@@ -105,7 +105,7 @@ def _standardise_columns(matrix: np.ndarray) -> np.ndarray:
 
 # build_feature_matrix: feature matrix for a Subvolumes instance, read in per-tomogram batches
 def build_feature_matrix(
-    store: Subvolumes,
+    subvols: Subvolumes,
     n_radial_bins: int = 12,
     max_azimuthal_mode: int = 4,
     n_azimuthal_samples: int = 64,
@@ -114,15 +114,15 @@ def build_feature_matrix(
     *,
     batch_size: int = 512,
 ) -> np.ndarray:
-    log.debug(f'Building features for {len(store)} subvolumes, max_mode={max_azimuthal_mode}')
-    if len(store) == 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: list[np.ndarray | None] = [None] * len(store)
-    for indices, batch in store.iter_batches(batch_size):
+    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
diff --git a/tests/unit/test_schemas.py b/tests/unit/test_schemas.py
index e44c873..bc5fcc8 100644
--- a/tests/unit/test_schemas.py
+++ b/tests/unit/test_schemas.py
@@ -149,41 +149,41 @@ def test_getitem_matches_dense_reference(self, tmp_path):
         a = _write_cache(tmp_path, 'tomoA', 2, box_voxels)
         b = _write_cache(tmp_path, 'tomoB', 3, box_voxels, offset=1000)
         index = [('tomoA', 0), ('tomoA', 1), ('tomoB', 0), ('tomoB', 1), ('tomoB', 2)]
-        store = Subvolumes(tmp_path, index, box_voxels)
+        subvols = Subvolumes(tmp_path, index, box_voxels)
         dense = np.concatenate([a, b], axis=0)
-        for i in range(len(store)):
-            np.testing.assert_array_equal(store[i], dense[i])
+        for i in range(len(subvols)):
+            np.testing.assert_array_equal(subvols[i], dense[i])
 
     def test_take_returns_view_without_copying_cache_files(self, tmp_path):
         box_voxels = 2
         _write_cache(tmp_path, 'tomoA', 4, box_voxels)
         index = [('tomoA', i) for i in range(4)]
-        store = Subvolumes(tmp_path, index, box_voxels)
-        subset = store.take([0, 2])
+        subvols = Subvolumes(tmp_path, index, box_voxels)
+        subset = subvols.take([0, 2])
         assert len(subset) == 2
-        np.testing.assert_array_equal(subset[0], store[0])
-        np.testing.assert_array_equal(subset[1], store[2])
+        np.testing.assert_array_equal(subset[0], subvols[0])
+        np.testing.assert_array_equal(subset[1], subvols[2])
 
     def test_iter_batches_covers_every_index_grouped_by_tomogram(self, tmp_path):
         box_voxels = 2
         _write_cache(tmp_path, 'tomoA', 5, box_voxels)
         _write_cache(tmp_path, 'tomoB', 3, box_voxels, offset=500)
         index = [('tomoA', i) for i in range(5)] + [('tomoB', i) for i in range(3)]
-        store = Subvolumes(tmp_path, index, box_voxels)
+        subvols = Subvolumes(tmp_path, index, box_voxels)
         seen: dict[int, np.ndarray] = {}
-        for indices, batch in store.iter_batches(batch_size=2):
+        for indices, batch in subvols.iter_batches(batch_size=2):
             for local_index, subvolume in zip(indices, batch):
                 seen[local_index] = subvolume
-        assert set(seen) == set(range(len(store)))
+        assert set(seen) == set(range(len(subvols)))
         for i in seen:
-            np.testing.assert_array_equal(seen[i], store[i])
+            np.testing.assert_array_equal(seen[i], subvols[i])
 
-    def test_write_rolled_produces_a_readable_store(self, tmp_path):
+    def test_write_rolled_produces_a_readable_subvols(self, tmp_path):
         box_voxels = 2
         _write_cache(tmp_path, 'tomoA', 3, box_voxels)
         index = [('tomoA', i) for i in range(3)]
-        store = Subvolumes(tmp_path, index, box_voxels)
-        rolled = ((i, store[i] * 2) for i in range(len(store)))
-        new_store = store.write_rolled(rolled, tmp_path / 'rolled')
-        for i in range(len(store)):
-            np.testing.assert_array_equal(new_store[i], store[i] * 2)
+        subvols = Subvolumes(tmp_path, index, box_voxels)
+        rolled = ((i, subvols[i] * 2) for i in range(len(subvols)))
+        new_subvols = subvols.write_rolled(rolled, tmp_path / 'rolled')
+        for i in range(len(subvols)):
+            np.testing.assert_array_equal(new_subvols[i], subvols[i] * 2)

From 125145065db07e94dd9b3a229aeeb5bfa8d9a8cc Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Fri, 25 Sep 2026 16:57:22 +0100
Subject: [PATCH 09/10] test(classify): update classification tests

- Modified tests/unit/test_classify.py to fix unit tests broken due to refactoring.
- Modified tests/integration/test_classify_command.py to add tests for chunking.
---
 tests/integration/test_classify_command.py | 60 ++++++++++++++++++++++
 tests/unit/test_classify.py                | 58 +++++++++++++--------
 2 files changed, 98 insertions(+), 20 deletions(-)

diff --git a/tests/integration/test_classify_command.py b/tests/integration/test_classify_command.py
index c4b9d7c..ae0e579 100644
--- a/tests/integration/test_classify_command.py
+++ b/tests/integration/test_classify_command.py
@@ -8,7 +8,12 @@
 from typer.testing import CliRunner
 
 # Import main CLI
+from stamp.classify.extract import extract_particle_set
+from stamp.classify.features import build_feature_matrix
+from stamp.classify.align import align_inplane
+from stamp.classify.average import compute_class_averages
 from stamp.cli.cli import stamp_app
+from stamp.schemas.particles import HalfSet, Particle
 
 # Initialise runner
 runner = CliRunner()
@@ -113,3 +118,58 @@ def test_classify_strict_halfset_mode(self, tmp_path: Path) -> None:
         )
         assert result.exit_code == 0, result.output
         assert 'Cross-half cluster matching' in result.output
+
+class TestClassifyChunking:
+    def _guard_no_full_concatenate(self, monkeypatch, max_rows: int) -> None:
+        '''Fail the test if anything tries to stack more than max_rows at once.'''
+        real_stack = np.stack
+        def guarded_stack(arrays, *args, **kwargs):
+            if len(arrays) > max_rows:
+                raise AssertionError(f'np.stack called with {len(arrays)} arrays, exceeding the per-batch bound of {max_rows}')
+            return real_stack(arrays, *args, **kwargs)
+        monkeypatch.setattr(np, 'stack', guarded_stack)
+
+    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
+        n_particles_per_tomogram = 20
+        box_voxels = 9
+        rng = np.random.default_rng(0)
+        raw_dir = tmp_path / 'raw'
+        raw_dir.mkdir()
+        tomogram_paths: dict[str, str] = {}
+        particles: list[Particle] = []
+        grid = [(10 + (i % 5) * 5, 10 + (i // 5) * 5, 20) for i in range(n_particles_per_tomogram)]
+        for t in range(n_tomograms):
+            tomogram_id = f'tomo{t:03d}'
+            volume = rng.normal(0.0, 0.2, size=(40, 40, 40)).astype(np.float32)
+            for x, y, z in grid:
+                volume[z - 2 : z + 2, y - 2 : y + 2, x - 2 : x + 2] -= 3.0
+            path = raw_dir / f'{tomogram_id}.mrc'
+            with mrcfile.new(path, overwrite=True) as mrc:
+                mrc.set_data(volume)
+            tomogram_paths[tomogram_id] = str(path)
+            for index, (x, y, z) in enumerate(grid):
+                particles.append(
+                    Particle(
+                        particle_id=f'{tomogram_id}_p{index:03d}',
+                        tomogram_id=tomogram_id,
+                        position=(float(x), float(y), float(z)),
+                        orientation=(1.0, 0.0, 0.0, 0.0),
+                        source_picker='stamp-native',
+                        half_set=HalfSet.A if index % 2 == 0 else HalfSet.B,
+                    )
+                )
+        cache_dir = tmp_path / 'subvolumes'
+        self._guard_no_full_concatenate(monkeypatch, max_rows=n_particles_per_tomogram)
+        subvols, kept, skipped = extract_particle_set(particles, tomogram_paths, box_voxels, cache_dir=cache_dir)
+        assert not skipped
+        assert len(subvols) == n_tomograms * n_particles_per_tomogram
+        monkeypatch.undo()
+        features = build_feature_matrix(subvols, n_workers=1, batch_size=n_particles_per_tomogram)
+        assert features.shape[0] == len(subvols)
+        cluster_ids = ['c00'] * len(subvols)
+        angles = align_inplane(subvols, cluster_ids, n_workers=1)
+        assert len(angles) == len(subvols)
+        averages = compute_class_averages(subvols, cluster_ids)
+        assert averages['c00'][1] == len(subvols)
\ No newline at end of file
diff --git a/tests/unit/test_classify.py b/tests/unit/test_classify.py
index ef6b528..6919f39 100644
--- a/tests/unit/test_classify.py
+++ b/tests/unit/test_classify.py
@@ -4,6 +4,7 @@
 
 # Import external dependencies
 import mrcfile, numpy as np, pytest
+from pathlib import Path
 from scipy.ndimage import gaussian_filter, rotate
 
 # Import internal STAMP objects
@@ -26,6 +27,15 @@
 from stamp.commands.classify import _classify_combined
 from stamp.picking.geometry import quaternion_from_reference_to
 from stamp.schemas.particles import HalfSet, Particle
+from stamp.schemas.subvolumes import Subvolumes
+
+# _subvols_from_array: write a dense subvolume stack to a single-tomogram cache and wrap it as a Subvolumes subvols
+def _subvols_from_array(tmp_path: Path, array: np.ndarray) -> Subvolumes:
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    box_voxels = array.shape[-1]
+    np.save(tmp_path / 'tomo000.npy', np.asarray(array, dtype=np.float32))
+    index = [('tomo000', i) for i in range(array.shape[0])]
+    return Subvolumes(tmp_path, index, box_voxels)
 
 # _three_blobs: three well-separated Gaussian blobs in 2D
 def _three_blobs(n_per_group: int = 40, seed: int = 0) -> np.ndarray:
@@ -147,25 +157,27 @@ def test_rotational_average_shape(self) -> None:
         averaged = rotational_average(np.ones((11, 11, 11)), bin_index, valid, 4)
         assert averaged.shape == (4, 11)
 
-    def test_features_are_invariant_to_in_plane_rotation(self) -> None:
+    def test_features_are_invariant_to_in_plane_rotation(self, tmp_path: Path) -> None:
         '''A subvolume rotated about z should give the same feature vector.'''
         rng = np.random.default_rng(0)
         subvolume = rng.normal(size=(21, 21, 21))
         # Blur slightly so interpolation during rotation doesn't dominate.
         subvolume = np.cumsum(np.cumsum(subvolume, axis=0), axis=1) / 100.0
         rotated = rotate(subvolume, angle=37.0, axes=(0, 1), reshape=False, order=1)
-        features = build_feature_matrix(np.stack([subvolume, rotated]), n_radial_bins=8, max_azimuthal_mode=0)
+        subvols = _subvols_from_array(tmp_path, np.stack([subvolume, rotated]))
+        features = build_feature_matrix(subvols, n_radial_bins=8, max_azimuthal_mode=0)
         correlation = np.corrcoef(features[0], features[1])[0, 1]
         assert correlation > 0.9, f'in-plane rotation changed the features (r={correlation:.3f})'
 
-    def test_features_distinguish_different_structures(self) -> None:
+    def test_features_distinguish_different_structures(self, tmp_path: Path) -> None:
         '''A blob and an empty box give different features.'''
         box = 21
         centre = box // 2
         with_blob = np.zeros((box, box, box))
         with_blob[centre - 1 : centre + 2, centre - 1 : centre + 2, centre + 4 : centre + 7] = 1.0
         without_blob = np.zeros((box, box, box))
-        features = build_feature_matrix(np.stack([with_blob, without_blob]), n_radial_bins=8)
+        subvols = _subvols_from_array(tmp_path, np.stack([with_blob, without_blob]))
+        features = build_feature_matrix(subvols, n_radial_bins=8)
         assert not np.allclose(features[0], features[1])
 
     def test_empty_input_returns_empty_matrix(self) -> None:
@@ -173,13 +185,14 @@ def test_empty_input_returns_empty_matrix(self) -> None:
         features = build_feature_matrix(np.empty((0, 11, 11, 11)), n_radial_bins=4)
         assert features.shape[0] == 0
 
-    def test_mode_zero_reproduces_rotational_average(self) -> None:
+    def test_mode_zero_reproduces_rotational_average(self, tmp_path: Path) -> None:
         '''max_azimuthal_mode=0 must give features equivalent to rotational average.'''
         rng = np.random.default_rng(0)
         raw = rng.normal(size=(4, 21, 21, 21))
         subvolumes = np.cumsum(np.cumsum(raw, axis=1), axis=2) / 100.0
 
-        azimuthal = build_feature_matrix(subvolumes, n_radial_bins=10, max_azimuthal_mode=0)
+        subvols = _subvols_from_array(tmp_path, subvolumes)
+        azimuthal = build_feature_matrix(subvols, n_radial_bins=10, max_azimuthal_mode=0)
         bin_index, valid = cylindrical_bins(21, n_radial_bins=10)
         classic = np.stack([rotational_average(v, bin_index, valid, 10).ravel() for v in subvolumes])
         classic = (classic - classic.mean(axis=1, keepdims=True)) / classic.std(axis=1, keepdims=True)
@@ -209,7 +222,7 @@ def test_c3_particle_has_power_at_mode_three(self) -> None:
         assert per_mode[3] > per_mode[2]
         assert per_mode[3] > per_mode[4]
 
-    def test_c3_and_c4_separate_with_azimuthal_features_but_not_without(self) -> None:
+    def test_c3_and_c4_separate_with_azimuthal_features_but_not_without(self, tmp_path: Path) -> None:
         '''Two particles differing only in symmetry are indistinguishable to a rotational average and distinguishable once azimuthal modes are included.'''
         rng = np.random.default_rng(0)
         subvolumes = []
@@ -232,8 +245,9 @@ def _between_group_separation(features: np.ndarray) -> float:
             within = 0.5 * (group_three.std(0).mean() + group_four.std(0).mean())
             return between / max(within, 1e-9)
 
-        rotational_only = build_feature_matrix(subvolumes, n_radial_bins=12, max_azimuthal_mode=0)
-        with_azimuthal = build_feature_matrix(subvolumes, n_radial_bins=12, max_azimuthal_mode=4)
+        subvols = _subvols_from_array(tmp_path, subvolumes)
+        rotational_only = build_feature_matrix(subvols, n_radial_bins=12, max_azimuthal_mode=0)
+        with_azimuthal = build_feature_matrix(subvols, n_radial_bins=12, max_azimuthal_mode=4)
 
         assert _between_group_separation(with_azimuthal) > (1.5 * _between_group_separation(rotational_only)), 'azimuthal features did not improve C3-vs-C4 separation'
 
@@ -242,21 +256,23 @@ def test_azimuthal_sampling_must_resolve_requested_mode(self) -> None:
         with pytest.raises(ValueError, match='aliasing'):
             azimuthal_magnitudes(_c_n_particle(), n_radial_bins=12, n_azimuthal_samples=6, max_mode=6)
 
-    def test_inner_bins_excluded_from_azimuthal_block(self) -> None:
+    def test_inner_bins_excluded_from_azimuthal_block(self, tmp_path: Path) -> None:
         '''Raising min_radius_fraction should shrink the feature vector.'''
         subvolumes = np.stack([_c_n_particle() for _ in range(3)])
+        subvols = _subvols_from_array(tmp_path, subvolumes)
         wide = build_feature_matrix(
-            subvolumes, n_radial_bins=12, max_azimuthal_mode=4, min_radius_fraction=0.0
+            subvols, n_radial_bins=12, max_azimuthal_mode=4, min_radius_fraction=0.0
         )
         narrow = build_feature_matrix(
-            subvolumes, n_radial_bins=12, max_azimuthal_mode=4, min_radius_fraction=0.5
+            subvols, n_radial_bins=12, max_azimuthal_mode=4, min_radius_fraction=0.5
         )
         assert narrow.shape[1] < wide.shape[1]
 
-    def test_degenerate_rows_are_all_zero(self):
+    def test_degenerate_rows_are_all_zero(self, tmp_path: Path):
         '''build_feature_matrix z-scores a flat subvolume to zeros.'''
         stack = np.stack([np.zeros((11, 11, 11)), np.random.default_rng(0).random((11, 11, 11))])
-        features = build_feature_matrix(stack, n_radial_bins=4, max_azimuthal_mode=0)
+        subvols = _subvols_from_array(tmp_path, stack)
+        features = build_feature_matrix(subvols, n_radial_bins=4, max_azimuthal_mode=0)
         assert not features[0].any() and features[1].any()
 
 # TestCluster: class containing unit tests for src/stamp/classify/cluster.py
@@ -338,23 +354,25 @@ def test_shared_basis_matches_recover_planted_split(self):
 
 # TestAlign: class containing unit tests for src/stamp/classify/align.py
 class TestAlign:
-    def test_recovers_planted_rolls(self):
+    def test_recovers_planted_rolls(self, tmp_path: Path):
         '''Rolled copies of one motif align back to a common frame.'''
         box = 21
         base = _asymmetric_blob(box)
         planted = [0.0, 40.0, 120.0, 250.0, 300.0]
         stack = np.stack([roll_about_normal(base, angle) for angle in planted])
-        angles = align_inplane(stack, ['c00'] * len(planted), angular_step_degrees=10.0, iterations=4)
+        subvols = _subvols_from_array(tmp_path, stack)
+        angles = align_inplane(subvols, ['c00'] * len(planted), angular_step_degrees=10.0, iterations=4)
 
         aligned = np.stack([roll_about_normal(stack[i], angles[i]) for i in range(len(planted))])
         reference = aligned.mean(axis=0)
         for frame in aligned:
             assert np.corrcoef(frame.ravel(), reference.ravel())[0, 1] > 0.98
 
-    def test_noise_cluster_skipped(self):
+    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))
-        angles = align_inplane(stack, ['noise', 'noise', 'c00', 'c00'])
+        subvols = _subvols_from_array(tmp_path, stack)
+        angles = align_inplane(subvols, ['noise', 'noise', 'c00', 'c00'])
         assert set(angles) == {2, 3}
 
 # TestVesicle: class containing unit tests for vesicle_id carry-through in classification
@@ -366,7 +384,7 @@ def test_class_assignment_carries_vesicle_id(self, tmp_path: Path) -> None:
         centres = np.array([[0.0, 0.0], [20.0, 0.0]])
         features = np.vstack([rng.normal(centre, 1.0, size=(n_per_group, 2)) for centre in centres])
         n_particles = features.shape[0]
-        subvolumes = np.zeros((n_particles, box, box, box), dtype=np.float32)
+        subvolumes = _subvols_from_array(tmp_path / 'subvolumes', np.zeros((n_particles, box, box, box), dtype=np.float32))
 
         particles = [
             Particle(
@@ -397,7 +415,7 @@ def test_class_assignment_vesicle_id_none_when_unset(self, tmp_path: Path) -> No
         centres = np.array([[0.0, 0.0], [20.0, 0.0]])
         features = np.vstack([rng.normal(centre, 1.0, size=(n_per_group, 2)) for centre in centres])
         n_particles = features.shape[0]
-        subvolumes = np.zeros((n_particles, box, box, box), dtype=np.float32)
+        subvolumes = _subvols_from_array(tmp_path / 'subvolumes', np.zeros((n_particles, box, box, box), dtype=np.float32))
 
         particles = [
             Particle(

From 0273828f419f2283093b1d458334d97124fd0106 Mon Sep 17 00:00:00 2001
From: Sam Holland <hol_sam@icloud.com>
Date: Sat, 26 Sep 2026 11:51:47 +0100
Subject: [PATCH 10/10] chore: update stamp version to v0.11.0

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

diff --git a/pyproject.toml b/pyproject.toml
index 3fc42cd..3bfeaa1 100644
--- a/pyproject.toml
+++ b/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 = [
diff --git a/uv.lock b/uv.lock
index 5303fe1..971cce4 100644
--- a/uv.lock
+++ b/uv.lock
@@ -786,7 +786,7 @@ wheels = [
 
 [[package]]
 name = "stamp"
-version = "0.10.0"
+version = "0.11.0"
 source = { editable = "." }
 dependencies = [
     { name = "loguru" },