STAMP v0.10.0 release

docs/configuration.md

@@ -18,7 +18,7 @@ stop_after = "refine"                           # pick | classify | identify | r
 [decoy]
 enabled = true
 method = "rejected-surface"                     # rejected-surface | shifted | synthetic-noise
-n_decoys_per_tomogram = 50
+# n_decoys_per_tomogram = 50                    # omit for adaptive n
 min_distance_from_real_angstrom = 100.0
 min_distance_from_picks_angstrom = 60.0
 min_shift_angstrom = 200.0
@@ -96,7 +96,7 @@ Key | Type | Default
 -- | -- | --
 `enabled` | bool | `True`
 `method` | `rejected-surface`\|`shifted`\|`synthetic-noise` | `rejected-surface`
-`n_decoys_per_tomogram` | int | `50`
+`n_decoys_per_tomogram` | int\|`None` | `None`
 `min_distance_from_real_angstrom` | float | `100.0`
 `min_distance_from_picks_angstrom` | float | `60.0`
 `min_shift_angstrom` | float | `200.0`

pyproject.toml

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

src/stamp/adapters/native.py

@@ -3,7 +3,6 @@
 '''
 
 # Import external dependencies
-import json
 from dataclasses import dataclass, replace
 from pathlib import Path
 
@@ -13,6 +12,7 @@
 from stamp.picking.native import PICKER_NAME, NativePickerConfig, pick_tomogram
 from stamp.schemas.picks import RawPick
 from stamp.utils.errors import StampPipelineError, StampValidationError
+from stamp.utils.io import cache_tomogram_picks, load_cached_picks
 from stamp.utils.log import log
 from stamp.utils.parallel import run_parallel
 
@@ -81,17 +81,14 @@ def run_in_process(self, inputs: AdapterInputs) -> list[RawPick]:
                 tomogram_config = replace(config, normalise_per_vesicle=False, vesicle_labels_mrc=None)
             jobs.append(_PickJob(segmentation_path, tomogram_path, tomogram_id, tomogram_config))
 
-        picks: list[RawPick] = []
         def _on_success(job: _PickJob, result: tuple[str, list[RawPick]]) -> None:
             tomogram_id, tomogram_picks = result
             log.debug(f'{tomogram_id}: {len(tomogram_picks)} raw picks')
-            cache_path = inputs.output_directory / f'{tomogram_id}.json'
-            cache_path.write_text(json.dumps([pick.model_dump() for pick in tomogram_picks], indent=2))
-            picks.extend(tomogram_picks)
+            cache_tomogram_picks(inputs.output_directory, tomogram_id, tomogram_picks)
         def _on_error(job: _PickJob, exc: Exception) -> None:
             if isinstance(exc, StampValidationError):
                 log.error(f'{job.tomogram_id}: {exc}')
             else:
                 raise exc
         run_parallel(jobs, _pick_one, max_workers=n_workers, label='stamp-native', on_success=_on_success, on_error=_on_error)
-        return picks
+        return load_cached_picks(inputs.output_directory, inputs.tomogram_ids)

src/stamp/classify/extract.py

@@ -5,6 +5,7 @@
 # Import external dependencies
 import mrcfile, numpy as np
 from dataclasses import dataclass
+from pathlib import Path
 from scipy.ndimage import map_coordinates
 
 # Import STAMP schema
@@ -69,6 +70,7 @@ class _ExtractTomogramJob:
     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]]:
@@ -92,6 +94,10 @@ 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
+    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))
     return kept, skipped, subvolumes
 
 # extract_particle_set: extract subvolumes for a list of particles
@@ -101,13 +107,13 @@ def extract_particle_set(
     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]]:
     by_tomogram: dict[str, list[Particle]] = {}
     for particle in particles:
         by_tomogram.setdefault(particle.tomogram_id, []).append(particle)
 
     segmentation_paths = segmentation_paths or {}
-    subvolumes: list[np.ndarray] = []
     kept: list[Particle] = []
     skipped: list[Particle] = []
 
@@ -118,13 +124,13 @@ def extract_particle_set(
             log.warning(f'{tomogram_id}: no matching raw tomogram path, skipping {len(group)} particle(s)')
             skipped.extend(group)
             continue
-        jobs.append(_ExtractTomogramJob(tomogram_id, group, path, segmentation_paths.get(tomogram_id), box_voxels))
+        jobs.append(_ExtractTomogramJob(tomogram_id, group, path, segmentation_paths.get(tomogram_id), box_voxels, cache_dir))
 
-    # keep results ordered by input tomogram order so kept/subvolumes stay aligned deterministically
     result_by_tomogram: dict[str, tuple[list[Particle], list[Particle], list[np.ndarray]]] = {}
 
     def _on_success(job: _ExtractTomogramJob, result: tuple[list[Particle], list[Particle], list[np.ndarray]]) -> None:
-        result_by_tomogram[job.tomogram_id] = result
+        tomogram_kept, tomogram_skipped, subvolumes = result
+        result_by_tomogram[job.tomogram_id] = (tomogram_kept, tomogram_skipped, subvolumes)
 
     def _on_error(job: _ExtractTomogramJob, exc: Exception) -> None:
         if isinstance(exc, StampValidationError):
@@ -135,16 +141,21 @@ 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] = []
     for job in jobs:
         result = result_by_tomogram.get(job.tomogram_id)
         if result is None:
             continue
-        tomogram_kept, tomogram_skipped, tomogram_subvolumes = result
+        tomogram_kept, tomogram_skipped, subvolumes = result
         kept.extend(tomogram_kept)
         skipped.extend(tomogram_skipped)
-        subvolumes.extend(tomogram_subvolumes)
+        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))
 
     log.debug(f'extract_particle_set: {len(kept)} kept, {len(skipped)} skipped')
-    if not subvolumes:
+    if not parts:
         return np.empty((0, box_voxels, box_voxels, box_voxels)), kept, skipped
-    return np.stack(subvolumes), kept, skipped
+    return np.concatenate(parts, axis=0), kept, skipped

src/stamp/cli/classify.py

@@ -109,6 +109,10 @@ def classify(
         Literal['png', 'jpg', 'tiff', 'svg'],
         typer.Option('--plot-format', help='Image format for static plots.', rich_help_panel = 'Plotting'),
     ] = 'tiff',
+    keep_raw: Annotated[
+        bool,
+        typer.Option('--keep-raw', help='Keep the raw per-tomogram subvolume cache instead of archiving it to raw.tar.gz.'),
+    ] = False,
 ) -> None:
     '''Cluster picked particles by structural similarity.'''
     is_decoy = is_decoy_particle_set(ParticleSet.model_validate(json.loads(particles.read_text())))
@@ -135,4 +139,5 @@ def classify(
         n_azimuthal_samples=n_azimuthal_samples,
         make_plots=make_plots,
         plot_format=plot_format,
+        keep_raw=keep_raw,
     )

src/stamp/cli/cli.py

@@ -8,7 +8,7 @@
 from typing import Annotated, Literal
 
 # Import internal commands
-from stamp.cli import classify, config, decoy, filter, identify, pick, plot, refine, run, tools
+from stamp.cli import classify, config, decoy, filter, identify, internal, pick, plot, refine, run, tools
 from stamp.utils.errors import StampError
 from stamp.utils.log import configure_logging, log
 
@@ -47,6 +47,12 @@
     help='Misc STAMP tools.',
     rich_help_panel='Utilities'
 )
+stamp_app.add_typer(
+    internal.internalCli,
+    name='internal',
+    help='STAMP internal commands.',
+    hidden=True,
+)
 
 # logging_callback: provide logging options and configure logging
 @stamp_app.callback()

src/stamp/cli/decoy.py

@@ -9,6 +9,7 @@
 
 # Import decoy command functions/variables
 import stamp.commands.decoy as decoyfuncs
+from stamp.utils.errors import StampPipelineError
 from stamp.utils.io import resolve_output_dir
 
 # Initialise Typer app
@@ -49,9 +50,9 @@ def decoy(
         typer.Option('--picker-params', help='JSON of native-picker parameters. Must match the values used for `stamp pick`, or decoys won\'t be drawn from the same candidate pool.', rich_help_panel = 'Inputs'),
     ] = '{}',
     n_decoys_per_tomogram: Annotated[
-        int,
-        typer.Option('--n-decoys-per-tomogram', help='Decoy positions per tomogram.', rich_help_panel = 'Decoy placement'),
-    ] = 50,
+        int | None,
+        typer.Option('--n-decoys-per-tomogram', help='Decoy positions per tomogram. Omit for adaptive targeting (matches real pick count per tomogram).', rich_help_panel = 'Decoy placement'),
+    ] = None,
     min_distance_from_real_angstrom: Annotated[
         float,
         typer.Option('--min-distance-from-real-a', help='Minimum separation from any real pick.', rich_help_panel = 'Rejected surface decoy'),
@@ -84,6 +85,10 @@ def decoy(
         int,
         typer.Option('-n', '--n-processes', help='Number of processes to use for per-tomogram decoy generation (1 = sequential).', rich_help_panel = 'Decoy placement'),
     ] = 1,
+    keep_raw: Annotated[
+        bool,
+        typer.Option('--keep-raw', help='Keep the raw per-picker output directory instead of archiving it to raw.tar.gz.'),
+    ] = False,
     make_plots: Annotated[
         bool,
         typer.Option('--plots', help='Write decoy position plots.', rich_help_panel = 'Plotting'),
@@ -138,6 +143,7 @@ def decoy(
         synthetic_shape=synthetic_shape,
         seed=seed,
         n_workers=n_workers,
+        keep_raw=keep_raw,
         make_plots=make_plots,
         pick_plot_style=pick_plot_style,
         plot_format=plot_format,

src/stamp/cli/refine.py

@@ -45,6 +45,10 @@ def refine(
         float | None,
         typer.Option('--voxel-size-a', help='Voxel size in Å. Read from MRC headers if omitted.'),
     ] = None,
+    threshold: Annotated[
+        float,
+        typer.Option('--threshold', help='FSC threshold to use for resolution calculation.')
+    ] = 0.143,
     tool: Annotated[
         Literal['relion', 'm'],
         typer.Option('--tool', help='Tool to use for refinement.', rich_help_panel = 'Refinement'),
@@ -82,6 +86,7 @@ def refine(
         class_assignments=class_assignments,
         raw_tomogram_dir=raw_tomogram_dir,
         output_dir=resolve_output_dir(output_dir, 'refine'),
+        threshold=threshold,
         tool=tool,
         mask=mask,
         iterations=iterations,

src/stamp/commands/classify.py

@@ -25,7 +25,7 @@
 from stamp.utils.halfset import split_by_half_set
 from stamp.schemas.particles import ClassAssignment, HalfSet, ParticleSet
 from stamp.utils.errors import StampPipelineError
-from stamp.utils.io import match_by_stem, resolve_directory_voxel_size_angstrom, write_sidecar
+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
 from stamp.utils.plotting.core import PlotFormat, central_slice, finish, plot_path
 from stamp.utils.plotting.classify import class_average_grid, scatter_labels
@@ -54,6 +54,7 @@ def run_classify(
     n_workers: int = 1,
     make_plots: bool = True,
     plot_format: str = 'tiff',
+    keep_raw: bool = False,
 ) -> None:
     if voxel_size_angstrom is None:
         voxel_size_angstrom = resolve_directory_voxel_size_angstrom(list(raw_tomogram_dir.glob('*.mrc')))
@@ -76,7 +77,8 @@ def run_classify(
         box_voxels += 1  # odd box keeps the particle exactly centred
 
     subvolumes, 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,
+        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',
     )
     if skipped:
         log.warning(f'Skipped {len(skipped)} particles whose {box_voxels}-voxel box fell outside the volume or had no matching tomogram or had no orientation')
@@ -132,6 +134,12 @@ def run_classify(
     assignments_path = output_dir / 'class_assignments.json'
     assignments_path.write_text(json.dumps([a.model_dump() for a in assignments], indent=2))
 
+    if not keep_raw:
+        raw_dir = output_dir / 'raw'
+        if raw_dir.is_dir():
+            archive_path = archive_and_remove_directory(raw_dir)
+            log.info(f'Archived raw classify output to {archive_path}')
+
     write_sidecar(
         output_dir,
         stage='classify',
@@ -211,6 +219,7 @@ def build_classify_commands(config, output_dir: Path, track: str = 'real') -> li
         argv.append('--no-strict-halfset-independence')
     if not settings.inplane_alignment:
         argv.append('--no-inplane-alignment')
+    argv.append('--keep-raw' if settings.keep_raw else '--no-keep-raw')
     argv.append('--plots' if config.plots.enabled else '--no-plots')
     argv += ['--plot-format', config.plots.format]
     return [ToolCommand(tool='classify', argv=argv, working_directory=target, output_paths=[target / 'class_averages'])]