STAMP v0.11.1 release

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 = [

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]

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])]