From 0eee169caee2416766d7cbc71199dad6322bc2e1 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Sun, 27 Sep 2026 14:10:05 +0100 Subject: [PATCH 1/3] feat(utils): allow parallel callers to define pool - Modified utils/parallel.py to update function signatures to allow callers to pass pre-configured process pools. --- src/stamp/utils/parallel.py | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/src/stamp/utils/parallel.py b/src/stamp/utils/parallel.py index dae2b88..98190fc 100644 --- a/src/stamp/utils/parallel.py +++ b/src/stamp/utils/parallel.py @@ -23,11 +23,16 @@ def run_parallel( label: str, on_success: Callable[[T, R], None] | None = None, on_error: Callable[[T, Exception], None] | None = None, + pool: ProcessPoolExecutor | None = None, + progress_total: int | None = None, + progress_offset: int = 0, ) -> list[R]: total = len(items) if total == 0: return [] - report_every = max(1, total // 10) + # progress_total lets callers define an overall total to report against (eg if using batches) + report_total = progress_total if progress_total is not None else total + report_every = max(1, report_total // 10) results: list[R] = [] # _handle: per-item bookkeeping @@ -42,20 +47,29 @@ def _handle(item: T, index: int, get_result: Callable[[], R]) -> None: if on_success is not None: on_success(item, result) results.append(result) - if index % report_every == 0 or index == total: - log.progress(f'{label}: {index}/{total} complete') + reported_index = progress_offset + index + if reported_index % report_every == 0 or reported_index == report_total: + log.progress(f'{label}: {reported_index}/{report_total} complete') if max_workers <= 1 or total == 1: for index, item in enumerate(items, start=1): _handle(item, index, lambda item=item: worker_fn(item)) return results - level_name, log_path = get_worker_log_config() - with ProcessPoolExecutor(max_workers=max_workers, initializer=init_worker_logging, initargs=(level_name, log_path)) as pool: + # use the provided pool when given + if pool is not None: future_to_item = {pool.submit(worker_fn, item): item for item in items} for index, future in enumerate(as_completed(future_to_item), start=1): item = future_to_item[future] _handle(item, index, future.result) + return results + + level_name, log_path = get_worker_log_config() + with ProcessPoolExecutor(max_workers=max_workers, initializer=init_worker_logging, initargs=(level_name, log_path)) as owned_pool: + future_to_item = {owned_pool.submit(worker_fn, item): item for item in items} + for index, future in enumerate(as_completed(future_to_item), start=1): + item = future_to_item[future] + _handle(item, index, future.result) return results # _call_indexed: wrapper pairing a worker's result with its original index @@ -70,11 +84,17 @@ def run_parallel_ordered( *, max_workers: int = 1, label: str, + pool: ProcessPoolExecutor | None = None, + progress_total: int | None = None, + progress_offset: int = 0, ) -> list[R]: indexed_results = run_parallel( list(enumerate(items)), partial(_call_indexed, worker_fn), max_workers=max_workers, label=label, + pool=pool, + progress_total=progress_total, + progress_offset=progress_offset, ) return [result for _, result in sorted(indexed_results, key=lambda pair: pair[0])] From 4c5f275a6ba02492c03d3836b935448275de4ef9 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Sun, 27 Sep 2026 14:15:46 +0100 Subject: [PATCH 2/3] fix(classify): use single feature extraction pool - Modified classify/features.py to create process pool for feature extractions, so a new pool is not created for each batch. --- src/stamp/classify/features.py | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/src/stamp/classify/features.py b/src/stamp/classify/features.py index 3762370..5a7b802 100644 --- a/src/stamp/classify/features.py +++ b/src/stamp/classify/features.py @@ -4,12 +4,14 @@ # Import external dependencies import numpy as np +from concurrent.futures import ProcessPoolExecutor +from contextlib import nullcontext from functools import partial 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.log import get_worker_log_config, init_worker_logging, log from stamp.utils.parallel import run_parallel_ordered # cylindrical_bins: precompute per-voxel radial bin index and validity mask @@ -122,10 +124,28 @@ def build_feature_matrix( 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(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 + total_particles = len(subvols) + # create process pool for feature extraction to be reused across batches + pool_context = ( + ProcessPoolExecutor(max_workers=n_workers, initializer=init_worker_logging, initargs=get_worker_log_config()) + if n_workers > 1 + else nullcontext() + ) + with pool_context as pool: + completed = 0 + for indices, batch in subvols.iter_batches(batch_size): + batch_results = run_parallel_ordered( + list(batch), + worker, + max_workers=n_workers, + label='feature-extraction', + pool=pool, + progress_total=total_particles, + progress_offset=completed, + ) + completed += len(batch_results) + 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] From ba0c54ed3358a6ba442c1f7063c28b7dbe268da3 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Sun, 27 Sep 2026 14:16:45 +0100 Subject: [PATCH 3/3] chore: update stamp version to v0.11.1 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3bfeaa1..1efdbc6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "stamp" -version = "0.11.0" +version = "0.11.1" description = "A workflow for Sub-Tomogram Averaging Membrane Proteins" readme = "README.md" authors = [ diff --git a/uv.lock b/uv.lock index 971cce4..2b4d892 100644 --- a/uv.lock +++ b/uv.lock @@ -786,7 +786,7 @@ wheels = [ [[package]] name = "stamp" -version = "0.11.0" +version = "0.11.1" source = { editable = "." } dependencies = [ { name = "loguru" },