STAMP v0.13.1 release

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

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

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)

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:

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

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:

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)

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)

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:

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}')

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:

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