diff --git a/pyproject.toml b/pyproject.toml index a52da59..a467af2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "stamp" -version = "0.13.0" +version = "0.13.1" description = "A workflow for Sub-Tomogram Averaging Membrane Proteins" readme = "README.md" authors = [ diff --git a/src/stamp/classify/average.py b/src/stamp/classify/average.py index 431596d..4da2cb2 100644 --- a/src/stamp/classify/average.py +++ b/src/stamp/classify/average.py @@ -3,11 +3,12 @@ ''' # Import external dependencies -import mrcfile, numpy as np +import numpy as np from pathlib import Path # Import internal STAMP objects from stamp.schemas.subvolumes import Subvolumes +from stamp.utils.io import new_mrc from stamp.utils.log import log # compute_class_averages: mean subvolume per cluster, streamed in batches @@ -44,8 +45,6 @@ def write_class_averages( suffix = f'_half{half_set_label}' if half_set_label else '' path = output_dir / f'{cluster_id}{suffix}_n{count}.mrc' log.debug(f'Writing class average {cluster_id} (half-set: {half_set_label})') - with mrcfile.new(path, overwrite=True) as mrc: - mrc.set_data(np.transpose(average, (2, 1, 0)).astype(np.float32)) - mrc.voxel_size = voxel_size_angstrom + new_mrc(path, np.transpose(average, (2, 1, 0)).astype(np.float32), voxel_size=voxel_size_angstrom) written.append(path) return written diff --git a/src/stamp/classify/extract.py b/src/stamp/classify/extract.py index 7f6fb08..091c25b 100644 --- a/src/stamp/classify/extract.py +++ b/src/stamp/classify/extract.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import json, mrcfile, numpy as np, os, shutil +import json, numpy as np, os, shutil from dataclasses import dataclass from pathlib import Path from scipy.ndimage import map_coordinates @@ -12,6 +12,7 @@ from stamp.schemas.particles import Particle from stamp.schemas.subvolumes import Subvolumes from stamp.utils.errors import StampValidationError +from stamp.utils.io import new_mrc, open_mrc, read_mrc from stamp.utils.log import log from stamp.utils.parallel import run_parallel @@ -81,19 +82,16 @@ def box_fits_inside( # _subtract_one_tomogram: worker entry point to replace membrane voxels with the background median and write a float32 mrc def _subtract_one_tomogram(job: _SubtractTomogramJob) -> Path: - with mrcfile.open(job.tomogram_path, permissive=True) as mrc: + with open_mrc(job.tomogram_path) as mrc: tomogram = np.array(mrc.data, dtype=np.float32) # writable copy voxel_size = mrc.voxel_size.copy() - with mrcfile.open(job.segmentation_path, permissive=True) as mrc: - membrane = np.asarray(mrc.data) > 0 + membrane = read_mrc(job.segmentation_path) > 0 if membrane.shape != tomogram.shape: raise StampValidationError(f'segmentation shape {membrane.shape} does not match tomogram shape {tomogram.shape} for {job.tomogram_id!r}; they must be the same volume at the same binning') # replace membrane voxels with the background median so the class average is not dominated by the membrane slab tomogram[membrane] = np.median(tomogram[~membrane]) job.output_path.parent.mkdir(parents=True, exist_ok=True) - with mrcfile.new(str(job.output_path), overwrite=True) as out: - out.set_data(tomogram) - out.voxel_size = voxel_size + new_mrc(job.output_path, tomogram, voxel_size=voxel_size) return job.output_path # _marker_path: per-tomogram completion marker, written only after its subvolume cache is in place @@ -133,7 +131,7 @@ def _extract_one_tomogram(job: _ExtractTomogramJob) -> tuple[list[Particle], lis kept: list[Particle] = [] skipped: list[Particle] = [] subvolumes: list[np.ndarray] = [] - with mrcfile.mmap(job.tomogram_path, permissive=True) as mrc: + with open_mrc(job.tomogram_path, mmap=True) as mrc: tomogram = mrc.data if tomogram.dtype != np.float32: tomogram = tomogram.astype(np.float32) diff --git a/src/stamp/commands/classify.py b/src/stamp/commands/classify.py index cff1184..04a50f7 100644 --- a/src/stamp/commands/classify.py +++ b/src/stamp/commands/classify.py @@ -27,7 +27,7 @@ from stamp.schemas.particles import ClassAssignment, HalfSet, ParticleSet from stamp.schemas.subvolumes import Subvolumes from stamp.utils.errors import StampPipelineError -from stamp.utils.io import archive_and_remove_directory, match_by_stem, resolve_directory_voxel_size_angstrom, write_sidecar +from stamp.utils.io import archive_and_remove_directory, glob_mrc, match_by_stem, resolve_directory_voxel_size_angstrom, write_sidecar from stamp.utils.log import log, log_config from stamp.utils.plotting.core import PlotFormat, central_slice, finish, plot_path from stamp.utils.plotting.classify import class_average_grid, scatter_labels @@ -59,7 +59,7 @@ def run_classify( 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'))) + voxel_size_angstrom = resolve_directory_voxel_size_angstrom(glob_mrc(raw_tomogram_dir)) if voxel_size_angstrom is None: raise StampPipelineError('Voxel size not given and not found in any MRC header in --raw-dir.') log.debug(f'Resolved voxel size from MRC headers: {voxel_size_angstrom} Å') @@ -68,12 +68,12 @@ def run_classify( is_decoy = is_decoy_particle_set(particle_set) log.info(f'Loaded {len(particle_set.particles)} {"decoy" if is_decoy else "real"} particles') tomogram_ids = sorted({particle.tomogram_id for particle in particle_set.particles}) - matched, unmatched = match_by_stem([Path(f'{tomogram_id}.mrc') for tomogram_id in tomogram_ids], sorted(raw_tomogram_dir.glob('*.mrc'))) + matched, unmatched = match_by_stem([Path(f'{tomogram_id}.mrc') for tomogram_id in tomogram_ids], glob_mrc(raw_tomogram_dir)) for stem_path in unmatched: log.warning(f'No raw tomogram matching {stem_path.stem}, its particles will be skipped') tomogram_paths = {stem_path.stem: raw_path for stem_path, raw_path in matched.items()} segmentation_paths = ( - {path.stem: str(path) for path in sorted(segmentation_dir.glob('*.mrc'))} if segmentation_dir is not None else {} + {path.stem: str(path) for path in glob_mrc(segmentation_dir)} if segmentation_dir is not None else {} ) box_voxels = int(round(box_angstrom / voxel_size_angstrom)) if box_voxels % 2 == 0: diff --git a/src/stamp/commands/decoy.py b/src/stamp/commands/decoy.py index 5a60570..4884b7d 100644 --- a/src/stamp/commands/decoy.py +++ b/src/stamp/commands/decoy.py @@ -23,7 +23,7 @@ from stamp.schemas.particles import ParticleSet from stamp.utils.checkpoint import directory_signatures, file_signature, sync_checkpoint from stamp.utils.errors import StampPipelineError, StampValidationError -from stamp.utils.io import archive_and_remove_directory, load_tomogram_manifests, resolve_directory_voxel_size_angstrom, write_sidecar +from stamp.utils.io import archive_and_remove_directory, glob_mrc, load_tomogram_manifests, resolve_directory_voxel_size_angstrom, write_sidecar from stamp.utils.log import log from stamp.utils.plotting.picks import plot_positions @@ -58,7 +58,7 @@ def run_decoy( '''Generate a decoy dataset to run through STAMP alongside real data''' if voxel_size_angstrom is None: _HEADER_DIR = raw_tomogram_dir or segmentation_dir - voxel_size_angstrom = resolve_directory_voxel_size_angstrom(list(_HEADER_DIR.glob('*.mrc'))) + voxel_size_angstrom = resolve_directory_voxel_size_angstrom(glob_mrc(_HEADER_DIR)) log.progress(f'Generating decoy dataset ({method})') config_fields = { key: value diff --git a/src/stamp/commands/identify.py b/src/stamp/commands/identify.py index faa14c6..f260904 100644 --- a/src/stamp/commands/identify.py +++ b/src/stamp/commands/identify.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import json, matplotlib.pyplot as plt, mrcfile, numpy as np, re, shutil, tomllib +import json, matplotlib.pyplot as plt, numpy as np, re, shutil, tomllib from collections import defaultdict from dataclasses import dataclass from pathlib import Path @@ -21,7 +21,7 @@ from stamp.identify.simulate import simulate_density, to_comparable from stamp.utils.checkpoint import atomic_write_json, directory_signatures, signatures, sync_checkpoint from stamp.utils.errors import StampPipelineError -from stamp.utils.io import write_sidecar +from stamp.utils.io import glob_mrc, open_mrc, write_sidecar from stamp.utils.log import log from stamp.utils.parallel import run_parallel from stamp.utils.plotting.core import PlotFormat, finish, plot_path @@ -41,11 +41,11 @@ def _classify_was_inplane_aligned(classes: Path) -> bool: def load_class_averages(directory: Path) -> dict[str, tuple[np.ndarray, float]]: grouped: dict[str, list[np.ndarray]] = defaultdict(list) voxel_sizes: dict[str, float] = {} - for path in sorted(directory.glob('*.mrc')): + for path in glob_mrc(directory): match = _CLASS_ID.match(path.stem) if not match: continue - with mrcfile.open(str(path), permissive=True) as mrc: + with open_mrc(path) as mrc: grouped[match.group(1)].append(np.transpose(np.asarray(mrc.data), (2, 1, 0))) voxel_sizes[match.group(1)] = float(mrc.voxel_size.x) if not grouped: diff --git a/src/stamp/commands/pick.py b/src/stamp/commands/pick.py index 1774a67..e9d91b5 100644 --- a/src/stamp/commands/pick.py +++ b/src/stamp/commands/pick.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import json, mrcfile, numpy as np +import json, numpy as np from dataclasses import dataclass from pathlib import Path from skimage import measure @@ -23,7 +23,7 @@ from stamp.schemas.picks import RawPick from stamp.utils.checkpoint import signatures, sync_checkpoint from stamp.utils.errors import StampPipelineError, StampValidationError -from stamp.utils.io import archive_and_remove_directory, load_cached_picks, load_tomogram_manifests, write_sidecar +from stamp.utils.io import archive_and_remove_directory, glob_mrc, load_cached_picks, load_tomogram_manifests, read_mrc, write_sidecar from stamp.utils.log import log from stamp.utils.parallel import run_parallel from stamp.utils.plotting.picks import plot_positions @@ -109,7 +109,7 @@ def _resolve_vesicle_labels(vesicle_labels_mrc: Path | None, manifests: list[Tom log.warning('--vesicle-labels-mrc is a single file but multiple tomograms are being picked; ignoring --vesicle-labels-mrc') return {} - by_stem = {path.stem: path for path in sorted(vesicle_labels_mrc.glob('*.mrc'))} + by_stem = {path.stem: path for path in glob_mrc(vesicle_labels_mrc)} resolved: dict[str, Path] = {} for manifest in manifests: match = by_stem.get(f'{manifest.tomogram_id}_labelled') or next( @@ -132,8 +132,7 @@ class _VesicleSummaryJob: # _summarise_one_tomogram: multiprocessing worker entry point def _summarise_one_tomogram(job: _VesicleSummaryJob) -> list[dict]: - with mrcfile.open(str(job.segmentation_path), permissive=True) as mrc: - segmentation = np.asarray(mrc.data) + segmentation = read_mrc(job.segmentation_path) labels = load_vesicle_labels(job.labels_mrc, segmentation.shape) vertices, _normals = extract_surface(segmentation) _v, faces, _n2, _val = measure.marching_cubes(segmentation.astype(np.float32), level=0.5) diff --git a/src/stamp/commands/plot.py b/src/stamp/commands/plot.py index ae7650f..da048ec 100644 --- a/src/stamp/commands/plot.py +++ b/src/stamp/commands/plot.py @@ -9,6 +9,7 @@ # Import internal STAMP objects from stamp.schemas.particles import Particle, ParticleSet from stamp.utils.errors import StampPipelineError +from stamp.utils.io import glob_mrc from stamp.utils.log import log from stamp.utils.plotting.picks import plot_positions @@ -51,8 +52,8 @@ def run_plots_pick( particles: list[Particle] = [p for p in particle_set.particles if p.tomogram_id in selected] log.info(f'Plotting {len(selected)} of {len(all_tomogram_ids)} tomogram(s)') - segmentation_paths = {path.stem: path for path in sorted(segmentation_dir.glob('*.mrc'))} if segmentation_dir else {} - raw_tomogram_paths = {path.stem: path for path in sorted(raw_tomogram_dir.glob('*.mrc'))} if raw_tomogram_dir else {} + segmentation_paths = {path.stem: path for path in glob_mrc(segmentation_dir)} if segmentation_dir else {} + raw_tomogram_paths = {path.stem: path for path in glob_mrc(raw_tomogram_dir)} if raw_tomogram_dir else {} output_dir.mkdir(parents=True, exist_ok=True) plot_positions(particles, output_dir, pick_plot_style, plot_format, segmentation_paths, raw_tomogram_paths, zstack_movie=pick_zstack_movie, plot_3d_view=pick_plot_3d, max_workers=n_workers) diff --git a/src/stamp/commands/refine.py b/src/stamp/commands/refine.py index 6bdaaa9..7e07ed8 100644 --- a/src/stamp/commands/refine.py +++ b/src/stamp/commands/refine.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import json, matplotlib.pyplot as plt, mrcfile, numpy as np +import json, matplotlib.pyplot as plt, numpy as np from pathlib import Path # Import STAMP objects @@ -22,7 +22,7 @@ from stamp.schemas.particles import ClassAssignment, ParticleSet from stamp.utils.checkpoint import file_signature, signatures, sync_checkpoint from stamp.utils.errors import StampAdapterError, StampPipelineError, StampValidationError -from stamp.utils.io import resolve_directory_voxel_size_angstrom, write_sidecar +from stamp.utils.io import glob_mrc, new_mrc, read_mrc, resolve_directory_voxel_size_angstrom, write_sidecar from stamp.utils.log import log from stamp.utils.plotting.core import finish, plot_path from stamp.utils.plotting.refine import fsc_curve @@ -70,10 +70,8 @@ def _run_half(adapter, runner, particles, reference, workdir, parameters, backen if backend == 'mock': log.debug(f'{workdir.name}: mock backend, synthesising final map from the seed reference') # mock backend runs nothing; synthesise a deterministic map from the seed - with mrcfile.open(str(reference), permissive=True) as mrc: - seed = np.asarray(mrc.data, dtype=np.float32) - with mrcfile.new(final_map, overwrite=True) as mrc: - mrc.set_data(seed) + seed = read_mrc(reference, dtype=np.float32) + new_mrc(final_map, seed) else: produced = adapter.parse_output(result).parsed['final_map'] Path(produced).replace(final_map) @@ -98,7 +96,7 @@ def run_refine( plot_format: str = 'tiff', ) -> None: if voxel_size_angstrom is None: - voxel_size_angstrom = resolve_directory_voxel_size_angstrom(list(raw_tomogram_dir.glob('*.mrc'))) + voxel_size_angstrom = resolve_directory_voxel_size_angstrom(glob_mrc(raw_tomogram_dir)) if voxel_size_angstrom is None: raise StampValidationError('Voxel size not given and not found in any MRC header in --raw-dir.') log.debug(f'Resolved voxel size from MRC headers: {voxel_size_angstrom} Å') @@ -153,14 +151,11 @@ def run_refine( else: log.info(f'{target}: half B loaded from checkpoint') - with mrcfile.open(str(final_a), permissive=True) as mrc: - map_a = np.asarray(mrc.data, dtype=np.float32) - with mrcfile.open(str(final_b), permissive=True) as mrc: - map_b = np.asarray(mrc.data, dtype=np.float32) + map_a = read_mrc(final_a, dtype=np.float32) + map_b = read_mrc(final_b, dtype=np.float32) user_mask = None if mask is not None: - with mrcfile.open(str(mask), permissive=True) as mrc: - user_mask = np.asarray(mrc.data, dtype=np.float32) + user_mask = read_mrc(mask, dtype=np.float32) log.progress(f'{target}: computing FSC') 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: diff --git a/src/stamp/decoy/generate.py b/src/stamp/decoy/generate.py index 71fcc19..bfb1702 100644 --- a/src/stamp/decoy/generate.py +++ b/src/stamp/decoy/generate.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import mrcfile, numpy as np +import numpy as np from dataclasses import dataclass from pathlib import Path from scipy.spatial import cKDTree @@ -25,7 +25,7 @@ from stamp.schemas.picks import RawPick from stamp.utils.checkpoint import atomic_save_npz from stamp.utils.errors import StampValidationError -from stamp.utils.io import cache_tomogram_picks, load_cached_picks +from stamp.utils.io import cache_tomogram_picks, load_cached_picks, new_mrc, read_mrc from stamp.utils.log import log from stamp.utils.parallel import run_parallel, run_parallel_ordered @@ -186,8 +186,7 @@ def _on_error(job: _RejectedSurfaceJob, exc: Exception) -> None: # _extract_surface_for_manifest: worker entry point, loads segmentation and extracts its surface def _extract_surface_for_manifest(manifest: TomogramManifest) -> tuple[tuple[int, ...], np.ndarray, np.ndarray]: - with mrcfile.open(str(manifest.segmentation_path), permissive=True) as mrc: - segmentation = np.asarray(mrc.data) + segmentation = read_mrc(manifest.segmentation_path) vertices, normals = extract_surface(segmentation) return segmentation.shape, vertices, normals @@ -355,10 +354,8 @@ def _generate_one_synthetic_noise_tomogram(job: _SyntheticNoiseJob) -> tuple[Tom volume = rng.normal(0.0, 1.0, size=job.tomogram_shape).astype(np.float32) volume[job.shell > 0] += job.config.density_sign * 2.0 - with mrcfile.new(segmentation_path, overwrite=True) as mrc: - mrc.set_data(job.shell) - with mrcfile.new(tomogram_path, overwrite=True) as mrc: - mrc.set_data(volume) + new_mrc(segmentation_path, job.shell) + new_mrc(tomogram_path, volume) shell_voxels = np.argwhere(job.shell > 0) chosen = rng.choice( @@ -424,10 +421,8 @@ def generate_synthetic_noise_decoys( def _score_surface( manifest: TomogramManifest, config: NativePickerConfig ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - with mrcfile.open(str(manifest.segmentation_path), permissive=True) as mrc: - segmentation = np.asarray(mrc.data) - with mrcfile.open(str(manifest.raw_tomogram_path), permissive=True) as mrc: - tomogram = np.asarray(mrc.data) + segmentation = read_mrc(manifest.segmentation_path) + tomogram = read_mrc(manifest.raw_tomogram_path) if segmentation.shape != tomogram.shape: raise StampValidationError(f'Segmentation and tomogram shapes disagree for {manifest.tomogram_id}') diff --git a/src/stamp/filter/render.py b/src/stamp/filter/render.py index b7b11b3..0079ddf 100644 --- a/src/stamp/filter/render.py +++ b/src/stamp/filter/render.py @@ -3,32 +3,33 @@ ''' # Import external dependencies -import mrcfile, numpy as np +import numpy as np from matplotlib.colors import LinearSegmentedColormap from pathlib import Path # Import internal STAMP objects from stamp.schemas.particles import Particle +from stamp.utils.io import open_mrc from stamp.utils.plotting.core import PROTEIN_CMAP, bare, central_slice # load_central_slice: memory-mapped and downsampled central Z slice of an MRC volume def load_central_slice(path: Path, *, max_side_px: int = 1024) -> tuple[np.ndarray, int]: - with mrcfile.mmap(str(path), mode='r', permissive=True) as mrc: + with open_mrc(path, mmap=True) as mrc: plane = np.asarray(central_slice(mrc.data)) step = max(1, max(plane.shape) // max_side_px) return plane[::step, ::step], step # slice_at: memory-mapped and downsampled slice of an MRC volume at a given Z def slice_at(path: Path, z: int, *, max_side_px: int = 1024) -> tuple[np.ndarray, int]: - with mrcfile.mmap(str(path), mode='r', permissive=True) as mrc: + with open_mrc(path, mmap=True) as mrc: plane = np.asarray(mrc.data[z]) step = max(1, max(plane.shape) // max_side_px) return plane[::step, ::step], step # volume_depth: Z extent of an MRC volume, for slider bounds def volume_depth(path: Path) -> int: - with mrcfile.mmap(str(path), mode='r', permissive=True) as mrc: - return mrc.data.shape[0] + with open_mrc(path, header_only=True) as mrc: + return int(mrc.header.nz) # draw_segmentation: dark-theme segmentation render def draw_segmentation(ax, plane: np.ndarray) -> None: diff --git a/src/stamp/filter/state.py b/src/stamp/filter/state.py index 66256a5..03bd859 100644 --- a/src/stamp/filter/state.py +++ b/src/stamp/filter/state.py @@ -10,6 +10,7 @@ # Import internal STAMP objects from stamp.schemas.particles import Particle, ParticleSet from stamp.utils.errors import StampPipelineError +from stamp.utils.io import glob_mrc # FilterState: a filtering session's data and accept/reject decisions @dataclass @@ -71,7 +72,7 @@ def load_filter_state( def _match(directory: Path | None) -> dict[str, Path]: if directory is None: return {} - return {path.stem: path for path in directory.glob('*.mrc') if path.stem in tomogram_ids} + return {path.stem: path for path in glob_mrc(directory) if path.stem in tomogram_ids} state = FilterState(particle_set, _match(segmentation_dir), _match(raw_tomogram_dir)) diff --git a/src/stamp/picking/native.py b/src/stamp/picking/native.py index 4e53951..ab01e17 100644 --- a/src/stamp/picking/native.py +++ b/src/stamp/picking/native.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import mrcfile, numpy as np +import numpy as np from dataclasses import dataclass from pathlib import Path from typing import Literal @@ -24,6 +24,7 @@ from stamp.picking.vesicles import load_vesicle_labels, vesicle_ids_at from stamp.schemas.picks import RawPick from stamp.utils.errors import StampPipelineError, StampValidationError +from stamp.utils.io import read_mrc from stamp.utils.log import log # PICKER_NAME: name for picker tool @@ -111,10 +112,8 @@ def pick_tomogram( tomogram_id: str, config: NativePickerConfig, ) -> list[RawPick]: - with mrcfile.open(str(segmentation_path), permissive=True) as mrc: - segmentation = np.asarray(mrc.data) - with mrcfile.open(str(tomogram_path), permissive=True) as mrc: - tomogram = np.asarray(mrc.data) + segmentation = read_mrc(segmentation_path) + tomogram = read_mrc(tomogram_path) if segmentation.shape != tomogram.shape: raise StampValidationError(f'Segmentation shape {segmentation.shape} does not match tomogram shape {tomogram.shape} for {tomogram_id}; they must be the same volume at the same binning') diff --git a/src/stamp/picking/vesicles.py b/src/stamp/picking/vesicles.py index 503da62..7d3e2ce 100644 --- a/src/stamp/picking/vesicles.py +++ b/src/stamp/picking/vesicles.py @@ -3,11 +3,13 @@ ''' # Import external dependencies -import mrcfile, numpy as np +import numpy as np from dataclasses import dataclass from pathlib import Path # Import internal STAMP objects +from stamp.utils.errors import StampValidationError +from stamp.utils.io import read_mrc from stamp.utils.log import log # VESICLE_ID_TEMPLATE: '{tomogram_id}:v{label:04d}', built from EValuator's own integer label @@ -15,11 +17,10 @@ # load_vesicle_labels: read an EValuator labelled MRC and check it matches the segmentation shape def load_vesicle_labels(labels_mrc_path: Path, segmentation_shape: tuple[int, ...]) -> np.ndarray: - with mrcfile.open(str(labels_mrc_path), permissive=True) as mrc: - labels = np.asarray(mrc.data).astype(np.int64) + labels = read_mrc(labels_mrc_path, dtype=np.int64) if labels.shape != segmentation_shape: - raise ValueError( - f'Vesicle labels MRC {labels_mrc_path!r} has shape {labels.shape}, ' + raise StampValidationError( + f'Vesicle labels MRC {str(labels_mrc_path)!r} has shape {labels.shape}, ' f'expected {segmentation_shape} to match the segmentation' ) log.debug(f'Loaded vesicle labels from {labels_mrc_path.name}, {int(labels.max())} label(s)') diff --git a/src/stamp/run/orchestrate.py b/src/stamp/run/orchestrate.py index 944959a..8bf3451 100644 --- a/src/stamp/run/orchestrate.py +++ b/src/stamp/run/orchestrate.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import json, mrcfile, tomllib +import json, tomllib from dataclasses import dataclass, field from pathlib import Path @@ -16,6 +16,7 @@ from stamp.commands.refine import _ADAPTERS as REFINE_ADAPTERS, run_refine from stamp.run.state import STAGE_ORDER, is_complete, mark_complete, stage_dir, stages_to_run from stamp.schemas.config import RunConfig +from stamp.utils.io import glob_mrc, open_mrc from stamp.utils.log import log # RunOutcome: completed run information reported @@ -51,11 +52,11 @@ def _real_pick(config: RunConfig, output_dir: Path) -> Path: def _synthetic_shape(config: RunConfig) -> tuple[int, int, int]: if config.decoy.synthetic_shape_voxels is not None: return config.decoy.synthetic_shape_voxels - first = next(iter(sorted(config.run.raw_tomogram_dir.glob('*.mrc'))), None) + first = next(iter(glob_mrc(config.run.raw_tomogram_dir)), None) if first is None: log.debug('No real tomograms found to infer synthetic shape from, defaulting to (200, 200, 200)') return (200, 200, 200) - with mrcfile.open(str(first), permissive=True, header_only=True) as mrc: + with open_mrc(first, header_only=True) as mrc: shape = (int(mrc.header.nz), int(mrc.header.ny), int(mrc.header.nx)) log.debug(f'Inferred synthetic shape {shape} from {first}') return shape diff --git a/src/stamp/utils/io.py b/src/stamp/utils/io.py index efe615f..eac1394 100644 --- a/src/stamp/utils/io.py +++ b/src/stamp/utils/io.py @@ -3,7 +3,9 @@ ''' # Import external dependencies -import hashlib, json, mrcfile, shutil, subprocess, tarfile, tomli_w, tomllib +import hashlib, json, mrcfile, numpy as np, shutil, subprocess, tarfile, tomli_w, tomllib +from collections.abc import Iterator +from contextlib import contextmanager from datetime import datetime, timezone from importlib.metadata import PackageNotFoundError, version from pathlib import Path @@ -31,6 +33,58 @@ def _is_writable(directory: Path): directory = _resolve_abspath(directory) return access(directory, W_OK) +# _mrc_error: translate an OSError/ValueError from mrcfile into a StampValidationError naming the path +def _mrc_error(action: str, path: Path, exc: Exception) -> Exception: + from stamp.utils.errors import StampValidationError + if isinstance(exc, PermissionError): + return StampValidationError(f'No permission to {action} MRC file {str(path)!r}') + if isinstance(exc, FileNotFoundError): + return StampValidationError(f'MRC file not found: {str(path)!r}') + return StampValidationError(f'Could not {action} MRC file {str(path)!r}: {exc}') + +# open_mrc: context manager opening an MRC with IO failures raised as StampValidationError +@contextmanager +def open_mrc(path: Path | str, *, mmap: bool = False, header_only: bool = False) -> Iterator[Any]: + path = Path(path) + # only wrap open so any errors raised in caller's with body aren't shown as IO errors instead + try: + if mmap: + mrc = mrcfile.mmap(str(path), mode='r', permissive=True) + else: + mrc = mrcfile.open(str(path), mode='r', permissive=True, header_only=header_only) + except (OSError, ValueError) as exc: + raise _mrc_error('read', path, exc) from exc + with mrc: + yield mrc + +# read_mrc: read an MRC file volume +def read_mrc(path: Path | str, *, dtype: type | None = None) -> np.ndarray: + with open_mrc(path) as mrc: + data = np.asarray(mrc.data) + return data if dtype is None else data.astype(dtype, copy=False) + +# new_mrc: write an array (and optional voxel size) to a new MRC, raising IO failures as StampValidationError +def new_mrc(path: Path | str, data: np.ndarray, *, voxel_size: Any = None) -> None: + path = Path(path) + try: + with mrcfile.new(str(path), overwrite=True) as mrc: + mrc.set_data(data) + if voxel_size is not None: + mrc.voxel_size = voxel_size + except (OSError, ValueError) as exc: + raise _mrc_error('write', path, exc) from exc + +# glob_mrc: sorted *.mrc paths in a directory, raising StampConfigError if it is missing or unreadable +def glob_mrc(directory: Path) -> list[Path]: + from os import access, R_OK, X_OK + from stamp.utils.errors import StampConfigError + if not directory.is_dir(): + raise StampConfigError(f'Not a directory: {str(directory)!r}') + # X_OK is needed to list and open files inside a directory + if not access(directory, R_OK | X_OK): + raise StampConfigError(f'No permission to read directory {str(directory)!r}') + return sorted(directory.glob('*.mrc')) + # toml_none_to_empty: map any None instances to an empty string for TOML serialisation def toml_none_to_empty(obj): if isinstance(obj, dict): @@ -187,8 +241,8 @@ def load_tomogram_manifests( ) -> list[TomogramManifest]: from stamp.utils.errors import StampPipelineError from stamp.utils.log import log - segmentation_paths = sorted(segmentation_dir.glob('*.mrc')) - raw_paths = sorted(raw_tomogram_dir.glob('*.mrc')) + segmentation_paths = glob_mrc(segmentation_dir) + raw_paths = glob_mrc(raw_tomogram_dir) matched, unmatched = match_by_stem(segmentation_paths, raw_paths) for segmentation_path in unmatched: log.warning(f'No raw tomogram matching {segmentation_path.stem}, skipping') @@ -209,7 +263,7 @@ def load_tomogram_manifests( # read_voxel_size_angstrom: voxel size from an MRC header, Angstrom, or None if absent/zero/non-cubic def read_voxel_size_angstrom(path: Path) -> float | None: from stamp.utils.log import log - with mrcfile.open(str(path), header_only=True, permissive=True) as mrc: + with open_mrc(path, header_only=True) as mrc: voxel_size = mrc.voxel_size x, y, z = float(voxel_size.x), float(voxel_size.y), float(voxel_size.z) if x <= 0 or y <= 0 or z <= 0: diff --git a/src/stamp/utils/plotting/picks.py b/src/stamp/utils/plotting/picks.py index 3f4b306..af4c657 100644 --- a/src/stamp/utils/plotting/picks.py +++ b/src/stamp/utils/plotting/picks.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import matplotlib.pyplot as plt, mrcfile, numpy as np +import matplotlib.pyplot as plt, numpy as np from dataclasses import dataclass from matplotlib.animation import FFMpegWriter, FuncAnimation, PillowWriter from mpl_toolkits.mplot3d.art3d import Poly3DCollection @@ -13,6 +13,7 @@ # Import internal STAMP objects from stamp.utils.plotting.core import GREEN, GREY, PlotFormat, central_slice, finish, plot_path, show_segmentation +from stamp.utils.io import read_mrc from stamp.utils.log import log from stamp.utils.parallel import run_parallel @@ -31,8 +32,7 @@ def _ensure_agg_backend() -> None: # _render_movie_job: multiprocessing worker entry point def _render_movie_job(job: _ZstackJob) -> None: _ensure_agg_backend() - with mrcfile.open(str(job.volume_path), permissive=True) as mrc: - volume = np.asarray(mrc.data) + volume = read_mrc(job.volume_path) _zstack_movie_for_background( job.tomogram_id, job.positions, volume, job.output_dir, job.stem_suffix, _RENDER_PLANE_BY_KIND[job.render_kind], job.fps, job.slab_voxels, @@ -49,8 +49,7 @@ class _SegmentedJob: # _render_segmented_job: multiprocessing worker entry point for plot_segmented def _render_segmented_job(job: _SegmentedJob) -> None: _ensure_agg_backend() - with mrcfile.open(str(job.segmentation_path), permissive=True) as mrc: - segmentation = central_slice(np.asarray(mrc.data)) + segmentation = central_slice(read_mrc(job.segmentation_path)) fig, ax = plt.subplots(figsize=(4.2, 4)) show_segmentation(ax, segmentation, field_px=segmentation.shape[-1], picks=job.picks) ax.set_title(job.tomogram_id, fontsize=9) @@ -96,8 +95,7 @@ class _RawJob: # _render_raw_job: multiprocessing worker entry point for plot_raw def _render_raw_job(job: _RawJob) -> None: _ensure_agg_backend() - with mrcfile.open(str(job.raw_path), permissive=True) as mrc: - raw_slice = central_slice(np.asarray(mrc.data)) + raw_slice = central_slice(read_mrc(job.raw_path)) fig, ax = plt.subplots(figsize=(4.2, 4)) ax.imshow(raw_slice, cmap='Greys_r') if len(job.picks): @@ -264,8 +262,7 @@ def _render_3d_job(job: _Scatter3DJob) -> None: fig = plt.figure(figsize=(6, 6)) ax = fig.add_subplot(projection='3d') if job.segmentation_path is not None: - with mrcfile.open(str(job.segmentation_path), permissive=True) as mrc: - segmentation = np.asarray(mrc.data) + segmentation = read_mrc(job.segmentation_path) # downsample so marching_cubes/rendering stays fast on full-size tomograms step = max(1, max(segmentation.shape) // 128) membrane = (segmentation[::step, ::step, ::step] == 1).astype(np.float32) diff --git a/tests/unit/test_picking.py b/tests/unit/test_picking.py index 68be1f3..d82ee64 100644 --- a/tests/unit/test_picking.py +++ b/tests/unit/test_picking.py @@ -618,9 +618,9 @@ def test_build_particle_set_assigns_ids_and_half_sets(self) -> None: # TestVesicle: class containing unit tests for picking/vesicle.py class TestVesicle: def test_load_vesicle_labels_rejects_shape_mismatch(self, tmp_path: Path) -> None: - '''Mismatched shapes raise ValueError rather than silently misattributing''' + '''Mismatched shapes raise StampValidationError rather than silently misattributing''' _write_mrc(tmp_path / 'labels.mrc', np.zeros((10, 10, 10))) - with pytest.raises(ValueError): + with pytest.raises(StampValidationError): load_vesicle_labels(tmp_path / 'labels.mrc', (20, 20, 20)) def test_load_vesicle_labels_reads_integer_labels(self, tmp_path: Path) -> None: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 9106966..03a37f5 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import mrcfile, numpy as np, pytest, subprocess, tomllib +import mrcfile, numpy as np, os, pytest, subprocess, tomllib from pathlib import Path # Import internal functions and schema @@ -11,7 +11,7 @@ from stamp.schemas.particles import HalfSet, Particle from stamp.utils import io as io_utils from stamp.schemas.provenance import ProvenanceSidecar -from stamp.utils.errors import StampPipelineError +from stamp.utils.errors import StampConfigError, StampPipelineError, StampValidationError from stamp.utils.log import configure_logging from stamp.utils.parallel import run_parallel, run_parallel_ordered @@ -272,6 +272,76 @@ 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 TestIoMrc: + def test_round_trip(self, tmp_path: Path) -> None: + path = tmp_path / 'a.mrc' + data = np.arange(24, dtype=np.float32).reshape(2, 3, 4) + io_utils.new_mrc(path, data, voxel_size=3.0) + assert np.array_equal(io_utils.read_mrc(path), data) + with io_utils.open_mrc(path, header_only=True) as mrc: + assert float(mrc.voxel_size.x) == 3.0 + with io_utils.open_mrc(path, mmap=True) as mrc: + assert mrc.data.shape == (2, 3, 4) + + def test_read_dtype(self, tmp_path: Path) -> None: + path = tmp_path / 'a.mrc' + io_utils.new_mrc(path, np.ones((2, 2, 2), dtype=np.float32)) + assert io_utils.read_mrc(path, dtype=np.float32).dtype == np.float32 + assert io_utils.read_mrc(path, dtype=np.float64).dtype == np.float64 + + def test_missing_file(self, tmp_path: Path) -> None: + with pytest.raises(StampValidationError, match='not found'): + io_utils.read_mrc(tmp_path / 'nope.mrc') + + def test_corrupt_file(self, tmp_path: Path) -> None: + path = tmp_path / 'bad.mrc' + path.write_bytes(b'not an mrc') + with pytest.raises(StampValidationError): + io_utils.read_mrc(path) + + @pytest.mark.skipif(os.geteuid() == 0, reason='root ignores file modes') + def test_unreadable_file(self, tmp_path: Path) -> None: + path = tmp_path / 'a.mrc' + io_utils.new_mrc(path, np.ones((2, 2, 2), dtype=np.float32)) + path.chmod(0o000) + try: + with pytest.raises(StampValidationError, match='No permission'): + io_utils.read_mrc(path) + finally: + path.chmod(0o600) + + @pytest.mark.skipif(os.geteuid() == 0, reason='root ignores file modes') + def test_unwritable_directory(self, tmp_path: Path) -> None: + tmp_path.chmod(0o500) + try: + with pytest.raises(StampValidationError, match='No permission'): + io_utils.new_mrc(tmp_path / 'a.mrc', np.ones((2, 2, 2), dtype=np.float32)) + finally: + tmp_path.chmod(0o700) + + def test_body_error_not_remapped(self, tmp_path: Path) -> None: + path = tmp_path / 'a.mrc' + io_utils.new_mrc(path, np.ones((2, 2, 2), dtype=np.float32)) + with pytest.raises(OSError, match='boom'): + with io_utils.open_mrc(path): + raise OSError('boom') + + def test_glob_mrc(self, tmp_path: Path) -> None: + (tmp_path / 'b.mrc').touch() + (tmp_path / 'a.mrc').touch() + assert [p.name for p in io_utils.glob_mrc(tmp_path)] == ['a.mrc', 'b.mrc'] + with pytest.raises(StampConfigError, match='Not a directory'): + io_utils.glob_mrc(tmp_path / 'missing') + + @pytest.mark.skipif(os.geteuid() == 0, reason='root ignores file modes') + def test_glob_mrc_unreadable_directory(self, tmp_path: Path) -> None: + tmp_path.chmod(0o000) + try: + with pytest.raises(StampConfigError, match='No permission'): + io_utils.glob_mrc(tmp_path) + finally: + tmp_path.chmod(0o700) + class TestLog: def test_configure_logging_writes_to_the_given_directory(self, tmp_path) -> None: '''Logging is configured to the specified directory.''' @@ -376,3 +446,5 @@ def test_round_trip_and_truncated_file_is_skipped(self, tmp_path: Path) -> None: (tmp_path / 'cache' / 't1.json').write_text('[{"tomogram_id": ') # truncated loaded = load_cached_picks(tmp_path / 'cache', ['t0', 't1']) assert [p.tomogram_id for p in loaded] == ['t0'] + + diff --git a/uv.lock b/uv.lock index 38fa4a2..047d694 100644 --- a/uv.lock +++ b/uv.lock @@ -850,7 +850,7 @@ wheels = [ [[package]] name = "stamp" -version = "0.13.0" +version = "0.13.1" source = { editable = "." } dependencies = [ { name = "fast-hdbscan" },