From 3b2fb5efefc1b0665e66ed5a80a3f435bb09bcfa Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 11:48:42 +0100 Subject: [PATCH 01/25] feat(utils): add pick caching utilities - Modified utils/io.py to add utility functions for saving a tomogram's picks to a cache directory, and reading a cache directory back into memory. --- src/stamp/utils/io.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/stamp/utils/io.py b/src/stamp/utils/io.py index 4a00c71..78e9c19 100644 --- a/src/stamp/utils/io.py +++ b/src/stamp/utils/io.py @@ -14,6 +14,7 @@ # Import STAMP objects from stamp.schemas.manifest import TomogramManifest +from stamp.schemas.picks import RawPick from stamp.schemas.provenance import ProvenanceSidecar # _CACHE_NAME: per-directory checksum cache keyed on (path, size, mtime) @@ -322,3 +323,18 @@ def print_toml(path: Path) -> None: console.print() console.print(build_toml_tree(data, path.name)) console.print() + +# cache_tomogram_picks: write one tomogram's picks to /.json +def cache_tomogram_picks(cache_dir: Path, tomogram_id: str, picks: list[RawPick]) -> None: + cache_dir.mkdir(parents=True, exist_ok=True) + (cache_dir / f'{tomogram_id}.json').write_text(json.dumps([pick.model_dump() for pick in picks], indent=2)) + +# load_cached_picks: read back per-tomogram JSON pick caches written by cache_tomogram_picks +def load_cached_picks(cache_dir: Path, tomogram_ids: list[str]) -> list[RawPick]: + picks: list[RawPick] = [] + for tomogram_id in tomogram_ids: + cache_path = cache_dir / f'{tomogram_id}.json' + if not cache_path.exists(): + continue + picks.extend(RawPick(**raw) for raw in json.loads(cache_path.read_text())) + return picks From 8a032789988e668eb9534522190b8467c213235b Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 12:35:06 +0100 Subject: [PATCH 02/25] feat(decoy): flush per-tomogram outputs to disk - Modified decoy/generate.py to update generation functions to cache each tomogram's picks to a JSON file, then read these back. - Modified commands/decoy.py to pass output directory through to decoy generation command, and add keep raw logic. - Modified cli/decoy.py to add --keep-raw option for disabling raw pick archiving. --- src/stamp/cli/decoy.py | 5 ++++ src/stamp/commands/decoy.py | 15 ++++++++-- src/stamp/decoy/generate.py | 58 ++++++++++++++++++++++++++----------- 3 files changed, 59 insertions(+), 19 deletions(-) diff --git a/src/stamp/cli/decoy.py b/src/stamp/cli/decoy.py index b9f45a4..d1178e0 100644 --- a/src/stamp/cli/decoy.py +++ b/src/stamp/cli/decoy.py @@ -84,6 +84,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 +142,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, diff --git a/src/stamp/commands/decoy.py b/src/stamp/commands/decoy.py index d869168..6afa82a 100644 --- a/src/stamp/commands/decoy.py +++ b/src/stamp/commands/decoy.py @@ -46,6 +46,7 @@ def run_decoy( synthetic_shape, seed, n_workers: int = 1, + keep_raw: bool = False, make_plots: bool = False, pick_plot_style: str = 'segmented', plot_format: str = 'tiff', @@ -97,6 +98,7 @@ def run_decoy( min_distance_from_real_angstrom=min_distance_from_real_angstrom, seed=seed, n_workers=n_workers, + output_dir=output_dir, ) else: decoy_set = generate_shifted_decoys( @@ -109,6 +111,7 @@ def run_decoy( min_pick_distance=min_distance_from_picks_angstrom, seed=seed, n_workers=n_workers, + output_dir=output_dir, ) if decoy_set is None or not decoy_set.particles: @@ -120,6 +123,12 @@ def run_decoy( decoy_path = output_dir / 'decoy_particle_set.json' decoy_path.write_text(decoy_set.model_dump_json(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 decoy output to {archive_path}') + write_sidecar( output_dir, stage='decoy', @@ -165,8 +174,10 @@ def build_decoy_commands(config, output_dir: Path) -> list[ToolCommand]: '--seed', str(config.stage.pick.half_set_seed), '--n-workers', str(config.decoy.n_workers), ] - argv.append('--plots' if config.plots.enabled else '--no-plots') - argv += ['--pick-plot-style', config.plots.pick_style, '--plot-format', config.plots.format] + argv.append('--keep-raw' if config.stage.pick.keep_raw else '--no-keep-raw') + if config.plots_enabled: + argv.append('--plots') + argv += ['--pick-plot-style', config.plots.pick_style, '--plot-format', config.plots.format] if config.plots.pick_zstack_movie: argv.append('--pick-zstack-movie') if config.plots.pick_plot_3d: diff --git a/src/stamp/decoy/generate.py b/src/stamp/decoy/generate.py index 97516a2..e87da5b 100644 --- a/src/stamp/decoy/generate.py +++ b/src/stamp/decoy/generate.py @@ -56,6 +56,7 @@ def generate_rejected_surface_decoys( seed: int, *, n_workers: int = 1, + output_dir: Path | None = None, ) -> ParticleSet | None: log.progress(f'Generating rejected-surface decoys ({n_decoys_per_tomogram}/tomogram)') rng = np.random.default_rng(seed) @@ -65,6 +66,8 @@ def generate_rejected_surface_decoys( min_distance_voxels = config.to_voxels(min_distance_from_real_angstrom) raw_picks: list[RawPick] = [] + cache_dir = output_dir / 'raw' / METHOD_REJECTED_SURFACE if output_dir is not None else None + tomogram_ids = [] # run surface scoring in a pool but keep rng-driven sampling sequential and in manifest order so results stay reproducible for a given seed result_by_tomogram: dict[str, _ScoreSurfaceResult] = {} @@ -120,21 +123,27 @@ def _on_error(job: _RejectedSurfaceJob, exc: Exception) -> None: lowest_first = np.argsort(scores)[:candidate_pool] chosen = rng.choice(lowest_first, size=n_to_take, replace=False) - for index in chosen: - normal_xyz = normals_zyx[index][::-1] - raw_picks.append( - RawPick( - tomogram_id=manifest.tomogram_id, - position=tuple(float(c) for c in points_xyz[index]), - orientation=quaternion_from_reference_to(np.asarray(normal_xyz)), - confidence=float(scores[index]), - source_picker=f'{DECOY_SOURCE_PREFIX}{METHOD_REJECTED_SURFACE}', - offset_window_angstrom=config.offset_windows_angstrom[int(winning_window[index])], - mean_score=float(mean_scores[index]), - profile_score=float(profile_scores[index]), - ) + tomogram_picks = [ + RawPick( + tomogram_id=manifest.tomogram_id, + position=tuple(float(c) for c in points_xyz[index]), + orientation=quaternion_from_reference_to(np.asarray(normals_zyx[index][::-1])), + confidence=float(scores[index]), + source_picker=f'{DECOY_SOURCE_PREFIX}{METHOD_REJECTED_SURFACE}', + offset_window_angstrom=config.offset_windows_angstrom[int(winning_window[index])], + mean_score=float(mean_scores[index]), + profile_score=float(profile_scores[index]), ) + for index in chosen + ] + if cache_dir is not None: + cache_tomogram_picks(cache_dir, manifest.tomogram_id, tomogram_picks) + tomogram_ids.append(manifest.tomogram_id) + else: + raw_picks.extend(tomogram_picks) + if cache_dir is not None: + raw_picks = load_cached_picks(cache_dir, tomogram_ids) decoys = _finalise(raw_picks, seed, METHOD_REJECTED_SURFACE) log.info(f'Generated {len(raw_picks)} decoy particles') return decoys @@ -159,6 +168,7 @@ def generate_shifted_decoys( max_attempts_per_particle: int = 200, *, n_workers: int = 1, + output_dir: Path | None = None, ) -> ParticleSet | None: log.progress('Generating shifted decoys') rng = np.random.default_rng(seed) @@ -188,6 +198,8 @@ def _on_error(manifest: TomogramManifest, exc: Exception) -> None: min_pick_distance_voxels = config.to_voxels(min_pick_distance) raw_picks: list[RawPick] = [] + cache_dir = output_dir / 'raw' / METHOD_SHIFTED if output_dir is not None else None + tomogram_ids: list[str] = [] for tomogram_id, positions in real_by_tomogram.items(): shape = shape_by_tomogram.get(tomogram_id) surface = surface_by_tomogram.get(tomogram_id) @@ -197,6 +209,7 @@ def _on_error(manifest: TomogramManifest, exc: Exception) -> None: pick_tree = cKDTree(np.array(positions)) extent_xyz = np.array(shape)[::-1] - 1 + tomogram_picks: list[RawPick] = [] for position in positions: origin = np.array(position) for _attempt in range(max_attempts_per_particle): @@ -213,7 +226,7 @@ def _on_error(manifest: TomogramManifest, exc: Exception) -> None: if pick_tree.query(candidate, k=1)[0] < min_pick_distance_voxels: continue - raw_picks.append( + tomogram_picks.append( RawPick( tomogram_id=tomogram_id, position=tuple(float(c) for c in candidate), @@ -226,6 +239,14 @@ def _on_error(manifest: TomogramManifest, exc: Exception) -> None: else: log.debug(f'{tomogram_id}: exhausted {max_attempts_per_particle} attempts placing a shifted decoy, skipping one particle') + if cache_dir is not None: + cache_tomogram_picks(cache_dir, tomogram_id, tomogram_picks) + tomogram_ids.append(tomogram_id) + else: + raw_picks.extend(tomogram_picks) + + if cache_dir is not None: + raw_picks = load_cached_picks(cache_dir, tomogram_ids) decoys = _finalise(raw_picks, seed, METHOD_SHIFTED) log.info(f'Generated {len(raw_picks)} decoy particles') return decoys @@ -242,6 +263,7 @@ class _SyntheticNoiseJob: config: NativePickerConfig seed: int n_decoys_per_tomogram: int + cache_dir: Path | None = None # _generate_one_synthetic_noise_tomogram: multiprocessing worker entry point def _generate_one_synthetic_noise_tomogram(job: _SyntheticNoiseJob) -> tuple[TomogramManifest, list[RawPick]]: @@ -284,6 +306,8 @@ def _generate_one_synthetic_noise_tomogram(job: _SyntheticNoiseJob) -> tuple[Tom source_picker=f'{DECOY_SOURCE_PREFIX}{METHOD_SYNTHETIC_NOISE}', ) ) + if job.cache_dir is not None: + cache_tomogram_picks(job.cache_dir, tomogram_id, raw_picks) return manifest, raw_picks # generate_synthetic_noise_decoys: Gaussian-noise volumes with a membrane-like shell and shell-sampled decoys @@ -310,12 +334,12 @@ def generate_synthetic_noise_decoys( ) distance = np.linalg.norm(grid - centre, axis=-1) shell = ((distance > radius - 1.5) & (distance < radius + 1.5)).astype(np.float32) - - jobs = [_SyntheticNoiseJob(index, tomogram_shape, shell, centre, segmentation_dir, tomogram_dir, config, seed, n_decoys_per_tomogram) for index in range(n_tomograms)] + cache_dir = output_dir / 'raw' / METHOD_SYNTHETIC_NOISE + jobs = [_SyntheticNoiseJob(index, tomogram_shape, shell, centre, segmentation_dir, tomogram_dir, config, seed, n_decoys_per_tomogram, cache_dir) for index in range(n_tomograms)] per_tomogram_results = run_parallel_ordered(jobs, _generate_one_synthetic_noise_tomogram, max_workers=n_workers, label='decoy-synthetic-noise') manifests = [manifest for manifest, _ in per_tomogram_results] - raw_picks = [pick for _, picks in per_tomogram_results for pick in picks] + raw_picks = load_cached_picks(cache_dir, [manifest.tomogram_id for manifest in manifests]) decoys = _finalise(raw_picks, seed, METHOD_SYNTHETIC_NOISE) log.info(f'Generated {len(raw_picks)} decoy particles') From 90b6a17328de7b80ee77b0d3f95c64b9a00b53f5 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 12:45:07 +0100 Subject: [PATCH 03/25] refactor(picking): refactor pick caching - Modified commands/pick.py to use shared pick caching utility. - Modified adapters/native.py to use shared pick caching utility. --- src/stamp/adapters/native.py | 9 +++------ src/stamp/commands/pick.py | 4 ++-- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/src/stamp/adapters/native.py b/src/stamp/adapters/native.py index 49391de..af3d017 100644 --- a/src/stamp/adapters/native.py +++ b/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) diff --git a/src/stamp/commands/pick.py b/src/stamp/commands/pick.py index e1b35d1..fbe0326 100644 --- a/src/stamp/commands/pick.py +++ b/src/stamp/commands/pick.py @@ -22,7 +22,7 @@ from stamp.schemas.manifest import TomogramManifest from stamp.schemas.picks import RawPick from stamp.utils.errors import StampPipelineError, StampValidationError -from stamp.utils.io import archive_and_remove_directory, load_tomogram_manifests, write_sidecar +from stamp.utils.io import archive_and_remove_directory, load_cached_picks, load_tomogram_manifests, write_sidecar from stamp.utils.log import log from stamp.utils.parallel import run_parallel from stamp.utils.plotting.picks import plot_positions @@ -229,7 +229,7 @@ def run_pick( parameters.setdefault('normalise_per_vesicle', normalise_per_vesicle) log.progress(f'Running {picker_name}...') - picks_by_picker[picker_name] = _run_picker(adapter, manifests, picker_output_dir, parameters, runner, backend) + picks_by_picker[picker_name] = (load_cached_picks(picker_output_dir, [m.tomogram_id for m in manifests]) if getattr(adapter, 'runs_in_process', False) and backend != 'mock' else _run_picker(adapter, manifests, picker_output_dir, parameters, runner, backend)) log.info(f'{picker_name}: {len(picks_by_picker[picker_name])} raw picks') log.progress(f'Reconciling picks...') From bd6dd3574ee18cdc28566585c29b48f20fc5b729 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 13:07:30 +0100 Subject: [PATCH 04/25] feat(classify): cache classification subvolumes - Modified classify/extract.py to cache each tomogram's subvolume stack to a cache directory. - Modified commands/classify.py to pass cache directory through to extraction functions, and add cache directory archiving logic. - Modified schemas/config.py to add keep_raw field to ClassifyStage configuration table. - Modified cli/classify to add a --keep-raw option for disabling archiving of raw subvolumes. --- src/stamp/classify/extract.py | 27 ++++++++++++++++++--------- src/stamp/cli/classify.py | 5 +++++ src/stamp/commands/classify.py | 13 +++++++++++-- src/stamp/schemas/config.py | 1 + 4 files changed, 35 insertions(+), 11 deletions(-) diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py index 9997747..8c623a5 100644 --- a/src/stamp/classify/extract.py +++ b/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]]] = {} + 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: - result_by_tomogram[job.tomogram_id] = result + tomogram_kept, tomogram_skipped, _subvolumes = result + result_by_tomogram[job.tomogram_id] = (tomogram_kept, tomogram_skipped) def _on_error(job: _ExtractTomogramJob, exc: Exception) -> None: if isinstance(exc, StampValidationError): @@ -135,16 +141,19 @@ 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 = 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')) 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 diff --git a/src/stamp/cli/classify.py b/src/stamp/cli/classify.py index 95e4dc8..cdc7c65 100644 --- a/src/stamp/cli/classify.py +++ b/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, ) diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py index e39ffb1..a539f0e 100644 --- a/src/stamp/commands/classify.py +++ b/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'])] diff --git a/src/stamp/schemas/config.py b/src/stamp/schemas/config.py index fa07575..d188b47 100644 --- a/src/stamp/schemas/config.py +++ b/src/stamp/schemas/config.py @@ -76,6 +76,7 @@ class ClassifyStage(_Strict): inplane_iterations: int = 3 random_state: int = 0 n_workers: int = 1 + keep_raw: bool = False class IdentifyStage(_Strict): candidates: Path From 29e1963cfeeaa460ea4851a344f45395f0eaaae4 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 14:20:55 +0100 Subject: [PATCH 05/25] feat(decoy): add raw pick option to schema - Modified schemas/config.py to add keep raw option to DecoySettings configuration table. - Modified commands/decoy.py to use above config key when constructing decoy command. --- src/stamp/commands/decoy.py | 2 +- src/stamp/schemas/config.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/stamp/commands/decoy.py b/src/stamp/commands/decoy.py index 6afa82a..304d3d7 100644 --- a/src/stamp/commands/decoy.py +++ b/src/stamp/commands/decoy.py @@ -174,7 +174,7 @@ def build_decoy_commands(config, output_dir: Path) -> list[ToolCommand]: '--seed', str(config.stage.pick.half_set_seed), '--n-workers', str(config.decoy.n_workers), ] - argv.append('--keep-raw' if config.stage.pick.keep_raw else '--no-keep-raw') + argv.append('--keep-raw' if config.decoy.keep_raw else '--no-keep-raw') if config.plots_enabled: argv.append('--plots') argv += ['--pick-plot-style', config.plots.pick_style, '--plot-format', config.plots.format] diff --git a/src/stamp/schemas/config.py b/src/stamp/schemas/config.py index d188b47..30d8d47 100644 --- a/src/stamp/schemas/config.py +++ b/src/stamp/schemas/config.py @@ -52,6 +52,7 @@ class DecoySettings(_Strict): n_synthetic_tomograms: int = 3 synthetic_shape_voxels: tuple[int, int, int] | None = None # None = match the first real tomogram n_workers: int = 1 + keep_raw: bool = False # Per-stage tables mirror each command's options class PickStage(_Strict): From 945d2e53ca5ab2a28b7bf4ee7a51e493fcfdf0e2 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 14:31:23 +0100 Subject: [PATCH 06/25] refactor(schemas): allow decoy per tomo to be None - Modified schemas/config.py to allow n_decoys_per_tomogram to be None as well as integer. --- src/stamp/schemas/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/stamp/schemas/config.py b/src/stamp/schemas/config.py index 30d8d47..9b5e61e 100644 --- a/src/stamp/schemas/config.py +++ b/src/stamp/schemas/config.py @@ -44,7 +44,7 @@ class PlotSettings(_Strict): class DecoySettings(_Strict): enabled: bool = True method: Literal['rejected-surface', 'shifted', 'synthetic-noise'] = 'rejected-surface' - n_decoys_per_tomogram: int = 50 + n_decoys_per_tomogram: int | None = None # None = adaptive (target same N as real picks) min_distance_from_real_angstrom: float = 100.0 min_distance_from_picks_angstrom: float = 60.0 min_shift_angstrom: float = 200.0 From a4a20d42f04dc23414547d1984386a9f35fcf745 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 14:50:08 +0100 Subject: [PATCH 07/25] refactor(decoy): allow decoy per tomo to be None - Modified cli/decoy.py to allow --n_decoys_per_tomogram to be None as well as integer, and to default to None. --- src/stamp/cli/decoy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/stamp/cli/decoy.py b/src/stamp/cli/decoy.py index d1178e0..d3408cf 100644 --- a/src/stamp/cli/decoy.py +++ b/src/stamp/cli/decoy.py @@ -49,9 +49,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'), From c93b70ae589289e04b42252d5dbb0309de41b402 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 14:53:18 +0100 Subject: [PATCH 08/25] feat(decoy): add adaptive targeting - Modified decoy/generate.py to add function for calculating the number of target decoys to generate for rejected surface method, and to raise error for synthetic noise error if omitted. - Modified commands/decoy.py to wrap synthetic noise generation in try...except to catch any validation errors. --- src/stamp/commands/decoy.py | 29 ++++++++++++++++------------- src/stamp/decoy/generate.py | 15 +++++++++++---- 2 files changed, 27 insertions(+), 17 deletions(-) diff --git a/src/stamp/commands/decoy.py b/src/stamp/commands/decoy.py index 304d3d7..5da1795 100644 --- a/src/stamp/commands/decoy.py +++ b/src/stamp/commands/decoy.py @@ -20,7 +20,7 @@ from stamp.picking.native import NativePickerConfig from stamp.run.state import stage_dir from stamp.schemas.particles import ParticleSet -from stamp.utils.errors import StampPipelineError +from stamp.utils.errors import StampPipelineError, StampValidationError from stamp.utils.io import load_tomogram_manifests, resolve_directory_voxel_size_angstrom, write_sidecar from stamp.utils.log import log from stamp.utils.plotting.picks import plot_positions @@ -71,18 +71,21 @@ def run_decoy( shape = tuple(int(value) for value in synthetic_shape.split(',')) if len(shape) != 3: raise StampPipelineError('--synthetic-shape must be three integers') - decoy_set, decoy_manifests = generate_synthetic_noise_decoys( - tomogram_shape=shape, - n_tomograms=n_synthetic_tomograms, - n_decoys_per_tomogram=n_decoys_per_tomogram, - output_dir=output_dir, - config=config, - seed=seed, - n_workers=n_workers, - ) - (output_dir / 'decoy_manifests.json').write_text( - json.dumps([m.model_dump(mode='json') for m in decoy_manifests], indent=2) - ) + try: + decoy_set, decoy_manifests = generate_synthetic_noise_decoys( + tomogram_shape=shape, + n_tomograms=n_synthetic_tomograms, + n_decoys_per_tomogram=n_decoys_per_tomogram, + output_dir=output_dir, + config=config, + seed=seed, + n_workers=n_workers, + ) + (output_dir / 'decoy_manifests.json').write_text( + json.dumps([m.model_dump(mode='json') for m in decoy_manifests], indent=2) + ) + except StampValidationError as exc: + raise StampPipelineError(str(exc)) from exc else: real_set = ParticleSet.model_validate(json.loads(real_particle_set.read_text())) manifests = load_tomogram_manifests(segmentation_dir, raw_tomogram_dir, voxel_size_angstrom) diff --git a/src/stamp/decoy/generate.py b/src/stamp/decoy/generate.py index e87da5b..dd1b391 100644 --- a/src/stamp/decoy/generate.py +++ b/src/stamp/decoy/generate.py @@ -42,6 +42,10 @@ class _RejectedSurfaceJob: # _ScoreSurfaceResult: type alias for _score_surface's return shape _ScoreSurfaceResult = tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] +# _adaptive_decoy_count: target decoy count derived from a tomogram's real pick count +def _adaptive_decoy_count(n_real_picks: int, *, target_ratio: float = 1.0) -> int: + return max(1, round(n_real_picks * target_ratio)) + # _score_one_tomogram: worker entry point, isolates StampValidationError per tomogram for on_error def _score_one_tomogram(job: _RejectedSurfaceJob) -> _ScoreSurfaceResult: return _score_surface(job.manifest, job.config) @@ -51,14 +55,14 @@ def generate_rejected_surface_decoys( real_particle_set: ParticleSet, manifests: list[TomogramManifest], config: NativePickerConfig, - n_decoys_per_tomogram: int, + n_decoys_per_tomogram: int | None, min_distance_from_real_angstrom: float, seed: int, *, n_workers: int = 1, output_dir: Path | None = None, ) -> ParticleSet | None: - log.progress(f'Generating rejected-surface decoys ({n_decoys_per_tomogram}/tomogram)') + log.progress(f'Generating rejected-surface decoys ({n_decoys_per_tomogram or "adaptive n"}/tomogram)') rng = np.random.default_rng(seed) real_by_tomogram: dict[str, list[tuple[float, float, float]]] = {} for particle in real_particle_set.particles: @@ -118,7 +122,8 @@ def _on_error(job: _RejectedSurfaceJob, exc: Exception) -> None: continue # prefer the most confidently empty positions, but jitter the choice so decoys aren't all clustered in one flat region - n_to_take = min(n_decoys_per_tomogram, points_xyz.shape[0]) + target = n_decoys_per_tomogram if n_decoys_per_tomogram is not None else _adaptive_decoy_count(len(real_positions)) + n_to_take = min(target, points_xyz.shape[0]) candidate_pool = min(points_xyz.shape[0], n_to_take * 5) lowest_first = np.argsort(scores)[:candidate_pool] chosen = rng.choice(lowest_first, size=n_to_take, replace=False) @@ -314,13 +319,15 @@ def _generate_one_synthetic_noise_tomogram(job: _SyntheticNoiseJob) -> tuple[Tom def generate_synthetic_noise_decoys( tomogram_shape: tuple[int, int, int], n_tomograms: int, - n_decoys_per_tomogram: int, + n_decoys_per_tomogram: int | None, output_dir: Path, config: NativePickerConfig, seed: int, *, n_workers: int = 1, ) -> tuple[ParticleSet | None, list[TomogramManifest]]: + if n_decoys_per_tomogram is None: + raise StampValidationError('n_decoys_per_tomogram must be set explicitly for synthetic-noise decoys') log.progress(f'Generating synthetic-noise decoys ({n_tomograms} tomograms, {n_decoys_per_tomogram}/tomogram)') segmentation_dir = output_dir / 'segmentations' tomogram_dir = output_dir / 'tomograms' From 23c846872ab535220f762b7e0187d2f039bc87e7 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:03:50 +0100 Subject: [PATCH 09/25] test(decoy): add tests for adaptive targeting - Modified tests/unit/test_decoy.py to add unit tests for adaptive targeting, and to split tests into two classes. --- tests/unit/test_decoy.py | 56 +++++++++++++++++++++++++++++++++++++--- 1 file changed, 53 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_decoy.py b/tests/unit/test_decoy.py index 66df209..9cf98e7 100644 --- a/tests/unit/test_decoy.py +++ b/tests/unit/test_decoy.py @@ -8,6 +8,7 @@ # Import internal STAMP objects from stamp.decoy.generate import ( + _adaptive_decoy_count, generate_rejected_surface_decoys, generate_shifted_decoys, generate_synthetic_noise_decoys, @@ -20,7 +21,7 @@ from stamp.picking.native import NativePickerConfig, pick_tomogram from stamp.schemas.manifest import TomogramManifest from stamp.schemas.particles import HalfSet, Particle, ParticleSet -from stamp.utils.errors import StampPipelineError +from stamp.utils.errors import StampPipelineError, StampValidationError # _write_vesicle_pair: write a matched segmentation/tomogram vesicle pair and return its manifest def _write_vesicle_pair(tmp_path: Path, tomogram_id: str = 'tomo000') -> TomogramManifest: @@ -70,8 +71,8 @@ def _real_set_from_picker(manifest: TomogramManifest, config: NativePickerConfig contributing_pickers=['stamp-native'], ) -# TestDecoy: class containing unit tests for test_decoy.py -class TestDecoy: +# TestDecoyGenerate: class containing unit tests for decoy/generate.py +class TestDecoyGenerate: def test_rejected_surface_decoys_avoid_real_picks(self, tmp_path: Path) -> None: '''Every decoy is kept clear of real picks''' manifest = _write_vesicle_pair(tmp_path) @@ -177,6 +178,55 @@ def test_synthetic_decoys_are_oriented(self, tmp_path: Path) -> None: ) assert all(p.orientation is not None for p in decoy_set.particles) + def test__adaptive_decoy_count_matches_real_picks(self): + '''Default ratio targets the same count as real picks.''' + assert _adaptive_decoy_count(20) == 20 + assert _adaptive_decoy_count(7) == 7 + + def test__adaptive_decoy_count_floors_at_one(self): + '''Zero picks still requests at least one decoy.''' + assert _adaptive_decoy_count(0) >= 1 + + def test__adaptive_decoy_count_respects_custom_ratio(self): + '''target_ratio scales the target decoys returned.''' + assert _adaptive_decoy_count(10, target_ratio=0.5) == 5 + + def test_generate_rejected_surface_decoys_adaptive_default(self, tmp_path): + '''Derive count from real picks per tomogram.''' + manifest = _write_vesicle_pair(tmp_path) + config = NativePickerConfig(voxel_size_angstrom=10.0, n_mad=2.5, offset_windows_angstrom=((20.0, 80.0),)) + real_set = _real_set_from_picker(manifest, config) + assert real_set.particles, 'fixture produced no real picks' + + decoy_set = generate_rejected_surface_decoys( + real_particle_set=real_set, + manifests=[manifest], + config=config, + n_decoys_per_tomogram=None, + min_distance_from_real_angstrom=100.0, + seed=1, + ) + assert decoy_set is not None + assert len(decoy_set.particles) <= len(real_set.particles) + + def test_generate_synthetic_noise_decoys_requires_explicit_count(self, tmp_path): + '''n_decoys_per_tomogram=None raises for synthetic noise decoys.''' + manifest = _write_vesicle_pair(tmp_path) + config = NativePickerConfig(voxel_size_angstrom=10.0, n_mad=2.5, offset_windows_angstrom=((20.0, 80.0),)) + real_set = _real_set_from_picker(manifest, config) + assert real_set.particles, 'fixture produced no real picks' + with pytest.raises(StampValidationError): + generate_synthetic_noise_decoys( + tomogram_shape=(40, 40, 40), + n_tomograms=1, + n_decoys_per_tomogram=None, + output_dir=tmp_path, + config=config, + seed=1, + ) + +# TestDecoyValidate: class containing unit tests for decoy/validate.py +class TestDecoyValidate: def test_is_decoy_particle_set_detects_mixture(self) -> None: '''A mixed decoy/real set raises DecoyContaminationError''' def _particle(particle_id: str, picker: str) -> Particle: From fe138f9ea12fbf128ddf0a2dab136c8d6af6a074 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:04:11 +0100 Subject: [PATCH 10/25] docs: update configuration documentation --- docs/configuration.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 4315d7b..7bfeb64 100644 --- a/docs/configuration.md +++ b/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` From bfd99238f5745928f8bd8d7711ffb074a87d66fc Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:06:44 +0100 Subject: [PATCH 11/25] fix(decoy): fix missing imports --- src/stamp/decoy/generate.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/stamp/decoy/generate.py b/src/stamp/decoy/generate.py index dd1b391..72aa2a9 100644 --- a/src/stamp/decoy/generate.py +++ b/src/stamp/decoy/generate.py @@ -24,6 +24,7 @@ from stamp.schemas.particles import Particle, ParticleSet from stamp.schemas.picks import RawPick from stamp.utils.errors import 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, run_parallel_ordered From 89842761dc3ff73658efbf4104c9e1bf88c7e63a Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:07:01 +0100 Subject: [PATCH 12/25] fix(schemas): fix circular import loop --- src/stamp/schemas/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/stamp/schemas/utils.py b/src/stamp/schemas/utils.py index 6c61950..8b1a88f 100644 --- a/src/stamp/schemas/utils.py +++ b/src/stamp/schemas/utils.py @@ -2,11 +2,9 @@ STAMP: field validators shared across schema modules ''' -# Import internal STAMP objects -from stamp.utils.errors import StampPipelineError - # assert_unit_quaternion: raise unless (w, x, y, z) is unit-length within tolerance def assert_unit_quaternion(value: tuple[float, float, float, float]) -> None: norm = sum(component**2 for component in value) ** 0.5 if not (0.99 <= norm <= 1.01): + from stamp.utils.errors import StampPipelineError raise StampPipelineError(f'orientation quaternion must be unit-length, got norm={norm:.4f}') From f61e964438a11b57e8de833f9056d70b9042eb14 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:22:42 +0100 Subject: [PATCH 13/25] fix(refine): add configurable FSC threshold - Modified schemas/config.py to add threshold to RefineStage configuration table. - Modified cli/refine.py to add --threshold option and pass it through to run_refine. - Modified commands/refine.py to pass threshold through to refinement functions. - Modified refine/fsc.py to use threshold instead of a hard-coded 0.143 when drawing the FSC threshold line. --- src/stamp/cli/refine.py | 5 +++++ src/stamp/commands/refine.py | 6 ++++-- src/stamp/refine/fsc.py | 8 ++++++-- src/stamp/schemas/config.py | 1 + 4 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/stamp/cli/refine.py b/src/stamp/cli/refine.py index 4afb93a..14667d7 100644 --- a/src/stamp/cli/refine.py +++ b/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, diff --git a/src/stamp/commands/refine.py b/src/stamp/commands/refine.py index 978e994..30a3d5e 100644 --- a/src/stamp/commands/refine.py +++ b/src/stamp/commands/refine.py @@ -77,6 +77,7 @@ def run_refine( class_assignments: Path, raw_tomogram_dir: Path, output_dir: Path, + threshold: float, tool: str, mask: Path | None, iterations: int, @@ -128,7 +129,7 @@ def run_refine( if mask is not None: with mrcfile.open(str(mask), permissive=True) as mrc: user_mask = np.asarray(mrc.data, dtype=np.float32) - fsc = compute_fsc(map_a, map_b, voxel_size_angstrom, mask=user_mask if user_mask is not None else soft_sphere_mask(map_a.shape)) + fsc = compute_fsc(map_a, map_b, voxel_size_angstrom, mask=user_mask if user_mask is not None else soft_sphere_mask(map_a.shape), threshold=threshold) except (StampValidationError, StampAdapterError) as exc: log.error(f'{target}: {exc}') failed_targets.append(target) @@ -136,7 +137,7 @@ def run_refine( final_a.replace(tree['combined'] / 'final_A.mrc') final_b.replace(tree['combined'] / 'final_B.mrc') - write_fsc_files(fsc, tree['combined']) + write_fsc_files(fsc, tree['combined'], threshold=threshold) log.info(f'{target}: resolution {fsc.resolution_angstrom:.1f} A @ FSC=0.143') if make_plots: @@ -182,6 +183,7 @@ def build_refine_commands(config, output_dir: Path) -> list[ToolCommand]: '--class-assignments', str(stage_dir(output_dir, 'real', 'classify') / 'class_assignments.json'), '--raw-dir', str(config.run.raw_tomogram_dir), '--out-dir', str(target), + '--threshold', str(config.stage.refine.threshold), '--voxel-size-a', str(config.run.voxel_size_angstrom), '--tool', config.stage.refine.tool, '--iterations', str(config.stage.refine.iterations), diff --git a/src/stamp/refine/fsc.py b/src/stamp/refine/fsc.py index 1e16823..782d6bd 100644 --- a/src/stamp/refine/fsc.py +++ b/src/stamp/refine/fsc.py @@ -119,7 +119,11 @@ def compute_fsc( ) # write_fsc_files: fsc.txt (columns) and a minimal fsc.svg line plot -def write_fsc_files(result: FSCResult, output_dir: Path) -> tuple[Path, Path]: +def write_fsc_files( + result: FSCResult, + output_dir: Path, + threshold: float = 0.143, +) -> tuple[Path, Path]: output_dir.mkdir(parents=True, exist_ok=True) text_path = output_dir / 'fsc.txt' rows = ['# freq_per_A fsc_unmasked fsc_masked'] @@ -135,7 +139,7 @@ def write_fsc_files(result: FSCResult, output_dir: Path) -> tuple[Path, Path]: f'{px:.1f},{height * (1 - max(min(value, 1.0), -0.2) / 1.2):.1f}' for px, value in zip(x_scaled, result.fsc_masked) ) - threshold_y = height * (1 - 0.143 / 1.2) + threshold_y = height * (1 - threshold / 1.2) svg_path.write_text( f'' f'' diff --git a/src/stamp/schemas/config.py b/src/stamp/schemas/config.py index 9b5e61e..e62b810 100644 --- a/src/stamp/schemas/config.py +++ b/src/stamp/schemas/config.py @@ -90,6 +90,7 @@ class IdentifyStage(_Strict): class RefineStage(_Strict): tool: Literal['relion', 'm'] = 'relion' class_id: str = 'all' + threshold: float = 0.143 iterations: int = 5 mask: Path | None = None backend: Literal['local', 'mock', 'cluster'] | None = None From b97bc96c4c63522b1ff9fedcf2012faaafa0871a Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:29:54 +0100 Subject: [PATCH 14/25] fix(decoy): fix missing import --- src/stamp/cli/decoy.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/stamp/cli/decoy.py b/src/stamp/cli/decoy.py index d3408cf..7b3531f 100644 --- a/src/stamp/cli/decoy.py +++ b/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 From 93012f84daca0c706f07dfe90ad26df8f184b140 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:33:45 +0100 Subject: [PATCH 15/25] fix(cli): register internal command group --- src/stamp/cli/cli.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/stamp/cli/cli.py b/src/stamp/cli/cli.py index 31d9652..88cc16c 100644 --- a/src/stamp/cli/cli.py +++ b/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() From c690bbfaf36419b727d5d92691d170a570bdf469 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:36:28 +0100 Subject: [PATCH 16/25] fix: fix command building - Modified commands/decoy.py to fix incorrect options and duplicated output directory resolution. - Modified commands/pick.py to fix incorrect options and duplicated output directory resolution. --- src/stamp/commands/decoy.py | 11 ++++++----- src/stamp/commands/pick.py | 10 ++++++---- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/stamp/commands/decoy.py b/src/stamp/commands/decoy.py index 5da1795..fd38252 100644 --- a/src/stamp/commands/decoy.py +++ b/src/stamp/commands/decoy.py @@ -168,21 +168,22 @@ def build_decoy_commands(config, output_dir: Path) -> list[ToolCommand]: real_particles = stage_dir(output_dir, 'real', 'pick') / 'particle_set.json' argv = [ 'stamp', 'decoy', - '--out-dir', str(target), + '--out-dir', str(output_dir), '--voxel-size-a', str(config.run.voxel_size_angstrom), '--method', config.decoy.method, '--real-particle-set', str(real_particles), '--seg-dir', str(config.run.segmentation_dir), '--raw-dir', str(config.run.raw_tomogram_dir), '--seed', str(config.stage.pick.half_set_seed), - '--n-workers', str(config.decoy.n_workers), + '--n-processes', str(config.decoy.n_workers), ] - argv.append('--keep-raw' if config.decoy.keep_raw else '--no-keep-raw') - if config.plots_enabled: + if config.decoy.keep_raw: + argv.append('--keep-raw') + if config.plots.enabled: argv.append('--plots') argv += ['--pick-plot-style', config.plots.pick_style, '--plot-format', config.plots.format] if config.plots.pick_zstack_movie: argv.append('--pick-zstack-movie') if config.plots.pick_plot_3d: argv.append('--pick-plot-3d') - return [ToolCommand(tool='decoy', argv=argv, working_directory=target, output_paths=[target / 'decoy_particle_set.json'])] + return [ToolCommand(tool='decoy', argv=argv, working_directory=target, output_paths=[target / 'decoy_particle_set.json'])] \ No newline at end of file diff --git a/src/stamp/commands/pick.py b/src/stamp/commands/pick.py index fbe0326..1b0577e 100644 --- a/src/stamp/commands/pick.py +++ b/src/stamp/commands/pick.py @@ -304,7 +304,7 @@ def build_pick_commands(config, output_dir: Path) -> list[ToolCommand]: 'stamp', 'pick', ','.join(config.stage.pick.pickers), '--seg-dir', str(config.run.segmentation_dir), '--raw-dir', str(config.run.raw_tomogram_dir), - '--out-dir', str(target), + '--out-dir', str(output_dir), '--voxel-size-a', str(config.run.voxel_size_angstrom), '--consensus-rule', config.stage.pick.consensus_rule, '--distance-threshold', str(config.stage.pick.distance_threshold), @@ -312,11 +312,13 @@ def build_pick_commands(config, output_dir: Path) -> list[ToolCommand]: '--backend', 'local', '--n-processes', str(config.stage.pick.n_workers), ] - argv.append('--plots' if config.plots.enabled else '--no-plots') - argv += ['--pick-plot-style', config.plots.pick_style, '--plot-format', config.plots.format] + if config.plots.enabled: + argv.append('--plots') + argv += ['--pick-plot-style', config.plots.pick_style, '--plot-format', config.plots.format] if config.plots.pick_zstack_movie: argv.append('--pick-zstack-movie') if config.plots.pick_plot_3d: argv.append('--pick-plot-3d') - argv.append('--keep-raw' if config.stage.pick.keep_raw else '--no-keep-raw') + if config.stage.pick.keep_raw: + argv.append('--keep-raw') return [ToolCommand(tool='pick', argv=argv, working_directory=target, output_paths=[target / 'particle_set.json'])] \ No newline at end of file From aa4c3ad60dfa4ad3d1a1ef39c787c85fdd64cd13 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:39:00 +0100 Subject: [PATCH 17/25] fix(utils): fix voxel size resolution falsy error - Modified utils/io.py to fix non-None but falsy voxel sizes falling back to resolved_voxel_size which wouldn't be assigned. --- src/stamp/utils/io.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/stamp/utils/io.py b/src/stamp/utils/io.py index 78e9c19..709e715 100644 --- a/src/stamp/utils/io.py +++ b/src/stamp/utils/io.py @@ -193,15 +193,15 @@ def load_tomogram_manifests( for segmentation_path in unmatched: log.warning(f'No raw tomogram matching {segmentation_path.stem}, skipping') if voxel_size_angstrom is None: - resolved_voxel_size = resolve_directory_voxel_size_angstrom(list(matched.values())) - if resolved_voxel_size is None: + voxel_size_angstrom = resolve_directory_voxel_size_angstrom(list(matched.values())) + if voxel_size_angstrom is None: raise StampPipelineError(f'No voxel size in any raw tomogram header under {raw_tomogram_dir} and --voxel-size-a not given') return [ TomogramManifest( tomogram_id=segmentation_path.stem, segmentation_path=segmentation_path, raw_tomogram_path=raw_path, - voxel_size_angstrom=voxel_size_angstrom or resolved_voxel_size, + voxel_size_angstrom=voxel_size_angstrom, ) for segmentation_path, raw_path in matched.items() ] From c407feb106b51367633faf3c9cdf8d8d98852a9b Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:43:26 +0100 Subject: [PATCH 18/25] chore: remove redundant function --- src/stamp/utils/log.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/src/stamp/utils/log.py b/src/stamp/utils/log.py index efa3793..cbeb5d5 100644 --- a/src/stamp/utils/log.py +++ b/src/stamp/utils/log.py @@ -95,21 +95,6 @@ def resolve_log_directory(specified_dir: Path | None): # If all options exhausted, raise an error raise PermissionError('No writable directory found for log file') -# -- resolve_log_path: returns a tuple of Path and bool, indicating the path to the final log file and whether the header needs to be written to it -def resolve_log_path( - directory: Path, - mode: str -) -> Path: - path = directory / DEFAULT_LOG_NAME - if mode == 'overwrite': - return path - if mode == 'new': - return _next_available_path(path) - if mode == 'append': - write_header = not path.exists() - return path - raise ValueError(f'Unknown mode {mode!r}, expected one of: overwrite, new, append') - # _worker_log_config: (level_name, log_path) as configured by main STAMP process _worker_log_config: tuple[str, Path | None] | None = None From 9058fab67243712af5029cc75c55524c5bb15696 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 15:56:58 +0100 Subject: [PATCH 19/25] test(utils): add log configuration tests - Modified tests/unit/test_utils.py to add unit tests for logging configuration. --- tests/unit/test_utils.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index c11efd0..544d209 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -12,6 +12,7 @@ from stamp.utils import io as io_utils from stamp.schemas.provenance import ProvenanceSidecar from stamp.utils.errors import StampPipelineError +from stamp.utils.log import configure_logging from stamp.utils.parallel import run_parallel, run_parallel_ordered # _particle: returns a valid Particle instance @@ -271,6 +272,28 @@ def test_no_headers_returns_none(self, tmp_path: Path) -> None: mrc.set_data(np.zeros((4, 4, 4), dtype=np.float32)) assert io_utils.resolve_directory_voxel_size_angstrom([path]) is None +class TestLog: + def test_configure_logging_writes_to_the_given_directory(self, tmp_path) -> None: + '''Logging is configured to the specified directory.''' + log_path, level_name = configure_logging(directory=tmp_path, mode='append', quiet=0, verbosity=0) + assert log_path == tmp_path / 'stamp.log' + assert log_path.is_file() + assert level_name == 'INFO' + + def test_configure_logging_append_mode_appends_to_existing_log(self, tmp_path) -> None: + '''Configuring logging with mode=append appends messages to existing files.''' + log_path = tmp_path / 'stamp.log' + log_path.write_text('existing content\n' * 100) + configure_logging(directory=tmp_path, mode='append', quiet=0, verbosity=0) + assert 'existing content' in log_path.read_text() + + def test_configure_logging_overwrite_mode_truncates_existing_log(self, tmp_path) -> None: + '''Configuring logging with mode=overwrite overwrites existing files.''' + log_path = tmp_path / 'stamp.log' + log_path.write_text('existing content\n' * 100) + configure_logging(directory=tmp_path, mode='overwrite', quiet=0, verbosity=0) + assert 'existing content' not in log_path.read_text() + class TestParallel: def test_sequential_default_matches_map(self): assert sorted(run_parallel([1, 2, 3], _square, max_workers=1, label='sq')) == [1, 4, 9] From f92a5cd092757976a33a6eec29e0e194fd5c9753 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 16:56:54 +0100 Subject: [PATCH 20/25] test: save log files to temporary directories --- tests/conftest.py | 2 +- tests/integration/test_classify_command.py | 2 ++ tests/integration/test_decoy_command.py | 6 ++++++ tests/integration/test_identify_command.py | 1 + tests/integration/test_pick_command.py | 5 +++++ tests/integration/test_plot_command.py | 8 ++++---- tests/integration/test_refine_command.py | 1 + tests/integration/test_run_command.py | 4 ++-- 8 files changed, 22 insertions(+), 7 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 2ce2cde..13e8e5a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -104,7 +104,7 @@ def _write_config(root: Path, **overrides) -> Path: def _run_mock_pipeline(root: Path, args: list[str] | None = None) -> Path: _make_dataset(root) config = _write_config(root) - result = _runner.invoke(stamp_app, ['run', '--config', str(config), *(args or [])]) + result = _runner.invoke(stamp_app, ['-d', str(root), 'run', '--config', str(config), *(args or [])]) assert result.exit_code == 0, result.output return root / 'stamp' diff --git a/tests/integration/test_classify_command.py b/tests/integration/test_classify_command.py index a3e7c61..c4b9d7c 100644 --- a/tests/integration/test_classify_command.py +++ b/tests/integration/test_classify_command.py @@ -67,6 +67,7 @@ def test_classify_end_to_end(self, tmp_path: Path) -> None: result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'classify', '--particles', str(particle_set_path), '--raw-dir', str(raw_dir), @@ -97,6 +98,7 @@ def test_classify_strict_halfset_mode(self, tmp_path: Path) -> None: stamp_app, [ '-vv', + '-d', str(tmp_path), 'classify', '--particles', str(particle_set_path), '--raw-dir', str(raw_dir), diff --git a/tests/integration/test_decoy_command.py b/tests/integration/test_decoy_command.py index 5e1d4f1..8ffb183 100644 --- a/tests/integration/test_decoy_command.py +++ b/tests/integration/test_decoy_command.py @@ -40,6 +40,7 @@ def _run_pick(tmp_path: Path) -> tuple[Path, Path, Path]: result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'pick', '--seg-dir', str(seg_dir), '--raw-dir', str(raw_dir), @@ -63,6 +64,7 @@ def test_decoy_rejected_surface_end_to_end(self, tmp_path: Path) -> None: result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'decoy', '--method', 'rejected-surface', '--real-particle-set', str(particle_set_path), @@ -89,6 +91,7 @@ def test_decoy_shifted_end_to_end(self, tmp_path: Path) -> None: result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'decoy', '--method', 'shifted', '--real-particle-set', str(particle_set_path), @@ -109,6 +112,7 @@ def test_decoy_synthetic_noise_end_to_end(self, tmp_path: Path) -> None: result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'decoy', '--method', 'synthetic-noise', '--out-dir', tmp_path, @@ -152,6 +156,7 @@ def test_decoy_rejects_undersized_set(self, tmp_path: Path) -> None: pick_result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'pick', '--seg-dir', str(seg_dir), '--raw-dir', str(raw_dir), @@ -171,6 +176,7 @@ def test_decoy_rejects_undersized_set(self, tmp_path: Path) -> None: result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'decoy', '--method', 'rejected-surface', '--real-particle-set', str(particle_set_path), diff --git a/tests/integration/test_identify_command.py b/tests/integration/test_identify_command.py index 1a3e6c2..5b57d4e 100644 --- a/tests/integration/test_identify_command.py +++ b/tests/integration/test_identify_command.py @@ -48,6 +48,7 @@ def test_identify_end_to_end(self, tmp_path): ) result = runner.invoke(stamp_app, [ + '-d', str(tmp_path), 'identify', '--classes', str(classes), '--candidates', str(tmp_path / 'candidates.yaml'), diff --git a/tests/integration/test_pick_command.py b/tests/integration/test_pick_command.py index acbdc78..3755a94 100644 --- a/tests/integration/test_pick_command.py +++ b/tests/integration/test_pick_command.py @@ -45,6 +45,7 @@ def test_pick_runs_native_picker_end_to_end(self, tmp_path: Path) -> None: result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'pick', '--seg-dir', str(seg_dir), '--raw-dir', str(raw_dir), @@ -79,6 +80,7 @@ def test_pick_skips_segmentations_without_raw_match(self, tmp_path: Path) -> Non result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'pick', '--seg-dir', str(seg_dir), '--raw-dir', str(raw_dir), @@ -112,6 +114,7 @@ def test_pick_matches_segmentation_by_stem_prefix(self, tmp_path: Path) -> None: result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'pick', '--seg-dir', str(seg_dir), '--raw-dir', str(raw_dir), @@ -132,6 +135,7 @@ def test_pick_reads_voxel_size_from_header_when_omitted(self, tmp_path: Path) -> result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'pick', '--seg-dir', str(seg_dir), '--raw-dir', str(raw_dir), @@ -152,6 +156,7 @@ def test_pick_warns_and_uses_most_common_voxel_size(self, tmp_path: Path) -> Non result = runner.invoke( stamp_app, [ + '-d', str(tmp_path), 'pick', '--seg-dir', str(seg_dir), '--raw-dir', str(raw_dir), diff --git a/tests/integration/test_plot_command.py b/tests/integration/test_plot_command.py index d32d8ff..7a3b47d 100644 --- a/tests/integration/test_plot_command.py +++ b/tests/integration/test_plot_command.py @@ -40,7 +40,7 @@ def test_plots_all_tomograms(self, tmp_path: Path) -> None: _write_raw(raw_dir, 'tomo000') _write_raw(raw_dir, 'tomo001') out_dir = tmp_path / 'plots' - result = runner.invoke(stamp_app, ['plot', 'pick', str(particle_set_path), '--raw-dir', str(raw_dir), '--out-dir', str(tmp_path), '--pick-plot-style', 'none', '--no-pick-zstack-movie', '--plot-format', 'png']) + result = runner.invoke(stamp_app, ['-d', str(tmp_path),'plot', 'pick', str(particle_set_path), '--raw-dir', str(raw_dir), '--out-dir', str(tmp_path), '--pick-plot-style', 'none', '--no-pick-zstack-movie', '--plot-format', 'png']) assert result.exit_code == 0, result.output assert list(out_dir.glob('*.png')) @@ -52,7 +52,7 @@ def test_restricts_by_tomogram_ids(self, tmp_path: Path) -> None: _write_raw(raw_dir, 'tomo000') _write_raw(raw_dir, 'tomo001') out_dir = tmp_path / 'plots' - result = runner.invoke(stamp_app, ['plot', 'pick', str(particle_set_path), '--raw-dir', str(raw_dir), '--out-dir', str(tmp_path), '--pick-plot-style', 'none', '--no-pick-zstack-movie', '--plot-format', 'png', '--tomogram-ids', 'tomo000']) + result = runner.invoke(stamp_app, ['-d', str(tmp_path),'plot', 'pick', str(particle_set_path), '--raw-dir', str(raw_dir), '--out-dir', str(tmp_path), '--pick-plot-style', 'none', '--no-pick-zstack-movie', '--plot-format', 'png', '--tomogram-ids', 'tomo000']) assert result.exit_code == 0, result.output assert 'Plotting 1 of 2 tomogram(s)' in result.output @@ -63,7 +63,7 @@ def test_unknown_tomogram_id_errors(self, tmp_path: Path) -> None: raw_dir = tmp_path / 'raw' _write_raw(raw_dir, 'tomo000') _write_raw(raw_dir, 'tomo001') - result = runner.invoke(stamp_app, ['plot', 'pick', str(particle_set_path), '--raw-dir', str(raw_dir), '--out-dir', str(tmp_path), '--tomogram-ids', 'not-a-tomogram']) + result = runner.invoke(stamp_app, ['-d', str(tmp_path),'plot', 'pick', str(particle_set_path), '--raw-dir', str(raw_dir), '--out-dir', str(tmp_path), '--tomogram-ids', 'not-a-tomogram']) assert result.exit_code != 0 # check --tomogram-ids and --n-tomograms are mutually exclusive @@ -72,5 +72,5 @@ def test_tomogram_ids_and_n_tomograms_mutually_exclusive(self, tmp_path: Path) - _write_particle_set(particle_set_path) raw_dir = tmp_path / 'raw' _write_raw(raw_dir, 'tomo000') - result = runner.invoke(stamp_app, ['plot', 'pick', str(particle_set_path), '--raw-dir', str(raw_dir), '--tomogram-ids', 'tomo000', '--n-tomograms', '1']) + result = runner.invoke(stamp_app, ['-d', str(tmp_path),'plot', 'pick', str(particle_set_path), '--raw-dir', str(raw_dir), '--tomogram-ids', 'tomo000', '--n-tomograms', '1']) assert result.exit_code != 0 diff --git a/tests/integration/test_refine_command.py b/tests/integration/test_refine_command.py index 3847d5b..847466b 100644 --- a/tests/integration/test_refine_command.py +++ b/tests/integration/test_refine_command.py @@ -62,6 +62,7 @@ def _fixture(tmp_path, mixed=False): def _invoke(tmp_path, stage_d): return runner.invoke(stamp_app, [ + '-d', str(tmp_path), 'refine', '--class-id', 'c00', '--identification', str(tmp_path / 'identification.json'), diff --git a/tests/integration/test_run_command.py b/tests/integration/test_run_command.py index 664e0cd..8441ce6 100644 --- a/tests/integration/test_run_command.py +++ b/tests/integration/test_run_command.py @@ -31,12 +31,12 @@ def test_run_produces_both_tracks_and_report(self, tmp_path, run_mock_pipeline): def test_second_invocation_skips_completed_stages(self, tmp_path, make_dataset, write_config): make_dataset(tmp_path) config = write_config(tmp_path) - first = runner.invoke(stamp_app, ['run', '--config', str(config)]) + first = runner.invoke(stamp_app, ['-d', str(tmp_path), 'run', '--config', str(config)]) assert first.exit_code == 0 state = json.loads((tmp_path / 'stamp' / 'run_state.json').read_text()) assert state['real']['pick'] is True - second = runner.invoke(stamp_app, ['run', '--config', str(config)]) + second = runner.invoke(stamp_app, ['-d', str(tmp_path), 'run', '--config', str(config)]) assert second.exit_code == 0 assert 'pick' not in second.output.lower() or 'skip' in second.output.lower() From 90196ba316ca4c6118b576252fe0267e6a13e18f Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 17:05:40 +0100 Subject: [PATCH 21/25] fix(run): add missing threshold config value --- src/stamp/run/orchestrate.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/stamp/run/orchestrate.py b/src/stamp/run/orchestrate.py index 92dd659..bfdfb39 100644 --- a/src/stamp/run/orchestrate.py +++ b/src/stamp/run/orchestrate.py @@ -213,6 +213,7 @@ def run_pipeline( class_assignments=stage_dir(output_dir, 'real', 'classify') / 'class_assignments.json', raw_tomogram_dir=config.run.raw_tomogram_dir, output_dir=refine_dir, + threshold=config.stage.refine.threshold, tool=config.stage.refine.tool, mask=config.stage.refine.mask, iterations=config.stage.refine.iterations, From 1b23ae1037a7717508ae68b409a84a9d5d1b8a08 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 17:09:09 +0100 Subject: [PATCH 22/25] fix(decoy): fix missing import --- src/stamp/commands/decoy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/stamp/commands/decoy.py b/src/stamp/commands/decoy.py index fd38252..2323b3e 100644 --- a/src/stamp/commands/decoy.py +++ b/src/stamp/commands/decoy.py @@ -21,7 +21,7 @@ from stamp.run.state import stage_dir from stamp.schemas.particles import ParticleSet from stamp.utils.errors import StampPipelineError, StampValidationError -from stamp.utils.io import load_tomogram_manifests, resolve_directory_voxel_size_angstrom, write_sidecar +from stamp.utils.io import archive_and_remove_directory, load_tomogram_manifests, resolve_directory_voxel_size_angstrom, write_sidecar from stamp.utils.log import log from stamp.utils.plotting.picks import plot_positions From a7bc5de65105e44f625c32e9319844d4ed1d40d5 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 17:09:41 +0100 Subject: [PATCH 23/25] fix(pick): remove cache load before picking --- src/stamp/commands/pick.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/stamp/commands/pick.py b/src/stamp/commands/pick.py index 1b0577e..d074880 100644 --- a/src/stamp/commands/pick.py +++ b/src/stamp/commands/pick.py @@ -229,7 +229,7 @@ def run_pick( parameters.setdefault('normalise_per_vesicle', normalise_per_vesicle) log.progress(f'Running {picker_name}...') - picks_by_picker[picker_name] = (load_cached_picks(picker_output_dir, [m.tomogram_id for m in manifests]) if getattr(adapter, 'runs_in_process', False) and backend != 'mock' else _run_picker(adapter, manifests, picker_output_dir, parameters, runner, backend)) + picks_by_picker[picker_name] = _run_picker(adapter, manifests, picker_output_dir, parameters, runner, backend) log.info(f'{picker_name}: {len(picks_by_picker[picker_name])} raw picks') log.progress(f'Reconciling picks...') From be52b2de1ee8eb523ea0791bf209fb898178ea9d Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 17:10:00 +0100 Subject: [PATCH 24/25] fix(classify): fix regressed subvolume bug --- src/stamp/classify/extract.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py index 8c623a5..efe5137 100644 --- a/src/stamp/classify/extract.py +++ b/src/stamp/classify/extract.py @@ -126,11 +126,11 @@ 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]]] = {} + 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: - tomogram_kept, tomogram_skipped, _subvolumes = result - result_by_tomogram[job.tomogram_id] = (tomogram_kept, tomogram_skipped) + 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): @@ -146,12 +146,14 @@ def _on_error(job: _ExtractTomogramJob, exc: Exception) -> None: result = result_by_tomogram.get(job.tomogram_id) if result is None: continue - tomogram_kept, tomogram_skipped = result + tomogram_kept, tomogram_skipped, subvolumes = 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)) log.debug(f'extract_particle_set: {len(kept)} kept, {len(skipped)} skipped') if not parts: From b7bea935b0e989a32d919686c335b2e2c6927f2b Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 24 Sep 2026 17:10:32 +0100 Subject: [PATCH 25/25] chore: update stamp version to v0.10.0 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 750059b..3fc42cd 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/uv.lock b/uv.lock index b10c48a..5303fe1 100644 --- a/uv.lock +++ b/uv.lock @@ -786,7 +786,7 @@ wheels = [ [[package]] name = "stamp" -version = "0.9.2" +version = "0.10.0" source = { editable = "." } dependencies = [ { name = "loguru" },