diff --git a/pyproject.toml b/pyproject.toml index a467af2..7782c31 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/src/stamp/classify/align.py b/src/stamp/classify/align.py index 3a1b1fa..8b31c72 100644 --- a/src/stamp/classify/align.py +++ b/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') diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py index 04a50f7..8d57c82 100644 --- a/src/stamp/commands/classify.py +++ b/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( diff --git a/src/stamp/schemas/subvolumes.py b/src/stamp/schemas/subvolumes.py index 2d4f9ce..ddfd68a 100644 --- a/src/stamp/schemas/subvolumes.py +++ b/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) diff --git a/uv.lock b/uv.lock index 047d694..89f579d 100644 --- a/uv.lock +++ b/uv.lock @@ -850,7 +850,7 @@ wheels = [ [[package]] name = "stamp" -version = "0.13.1" +version = "0.13.2" source = { editable = "." } dependencies = [ { name = "fast-hdbscan" },