Source code for neunorm.pipelines.mars_tpx3

"""
MARS TPX3 normalization pipeline.
"""

from datetime import datetime
from pathlib import Path
from typing import Optional, Sequence

import numpy as np
import scipp as sc
from loguru import logger

from neunorm import __version__
from neunorm.data_models.roi import ROILike, as_roi_bounds
from neunorm.exporters.hdf5_writer import write_hdf5
from neunorm.exporters.tiff_writer import write_tiff_stack
from neunorm.filters.gamma_filter import apply_gamma_filter
from neunorm.loaders.event_loader import load_event_nexus
from neunorm.processing.normalizer import BackgroundROILike, as_roi_bounds_list, normalize_transmission
from neunorm.processing.reference_preparer import prepare_reference
from neunorm.processing.roi_clipper import apply_roi
from neunorm.processing.run_combiner import combine_runs
from neunorm.tof.event_converter import convert_events_to_2d_histogram
from neunorm.tof.pixel_detector import detect_dead_pixels, detect_hot_pixels


[docs] def run_mars_tpx3_pipeline( # noqa: C901 sample_paths: Sequence[Sequence[str | Path]], ob_paths: Sequence[Sequence[str | Path]], output_path: Path, roi: Optional[ROILike] = None, gamma_filter: bool = True, detector_shape: tuple[int, int] = (514, 514), background_roi: Optional[BackgroundROILike] = None, ) -> sc.DataArray: """Execute MARS TPX3 normalization pipeline. Pipeline Step 1. Load event data 2. Convert events to 2D histogram 3. Run combine (optional) 4. ROI clip (optional) 5. Dead pixel detection 6. Hot pixel detection 7. Gamma filtering 8. Normalization 9. Output Parameters ---------- sample_paths : Sequence[Sequence[str | Path]] List of lists of paths to sample HDF5 files. Each inner list corresponds to one run and will be combined before processing. ob_paths : Sequence[Sequence[str | Path]] List of lists of paths to open beam HDF5 files Each inner list corresponds to one run and will be combined before processing. output_path : Path Path to save the output file (HDF5 or TIFF) roi : Optional[tuple] Region of interest to crop to — an ``ROI`` or a bare ``(x0, y0, x1, y1)`` tuple. gamma_filter : bool Whether to apply gamma filtering to the sample data (default: True) detector_shape : tuple[int, int] Shape of the TPX3 detector (default: (514, 514)) background_roi : ROI/tuple or a sequence of them Sample-free background region(s) — one ROI or a pooled sequence of ROIs (see ``normalize_transmission``) — for flux-proxy normalization when proton charge is unavailable. Mutually exclusive with proton-charge correction. If ``roi`` is also given the detector is cropped first, so ``background_roi`` indices are resolved in the post-crop frame. Notes ----- This function writes the normalized transmission data to disk in either HDF5 or TIFF format, depending on the file extension of `output_path`. Metadata and detector masks are included in the output. Returns ------- sc.DataArray Final normalized transmission DataArray with metadata and masks """ # Accept an ROI or a bare (x0, y0, x1, y1) tuple for every ROI argument; coerce to bounds # tuples up front so cropping and provenance see a consistent form. if roi is not None: roi = as_roi_bounds(roi) if background_roi is not None: background_roi = as_roi_bounds_list(background_roi) # Load data and convert to histogram samples = [ sc.concat( [ convert_events_to_2d_histogram(load_event_nexus(p, detector_shape=detector_shape), detector_shape) for p in run ], dim="N_image", ) for run in sample_paths ] obs = [ sc.concat( [ convert_events_to_2d_histogram(load_event_nexus(p, detector_shape=detector_shape), detector_shape) for p in run ], dim="N_image", ) for run in ob_paths ] # Combine runs if there are multiple runs sample = combine_runs(samples, metadata_keys_to_sum=[], metadata_check_match=[], normalize_by_runs=True) ob = combine_runs(obs, metadata_keys_to_sum=[], metadata_check_match=[], normalize_by_runs=True) # Apply ROI if specified if roi: sample = apply_roi(sample, roi) ob = apply_roi(ob, roi) # Average OB ob = prepare_reference(ob, dim="N_image") # Dead pixel detection sample.masks["dead_pixels"] = detect_dead_pixels(sample) # Hot pixel detection sample.masks["hot_pixels"] = detect_hot_pixels(sample) # Gamma filtering (optional) if gamma_filter: sample = apply_gamma_filter(sample) # Normalization (background_roi flux proxy when provided) transmission = normalize_transmission(sample, ob, background_roi=background_roi) # Write output metadata = { "sample_paths": [[str(p) for p in run] for run in sample_paths], "ob_paths": [[str(p) for p in run] for run in ob_paths], "gamma_filter_applied": gamma_filter, "processing_timestamp": datetime.now().isoformat(), "version": __version__, } if roi: metadata["roi_applied"] = roi if background_roi is not None: metadata["background_roi"] = ( list(background_roi[0]) if len(background_roi) == 1 else [list(b) for b in background_roi] ) if output_path.suffix.lower() in (".hdf5", ".h5"): write_hdf5( output_path, transmission, dead_pixel_mask="dead_pixels", hot_pixel_mask="hot_pixels", metadata=metadata ) elif output_path.suffix.lower() in (".tiff", ".tif"): rename_map = {} if "N_image" in transmission.dims: rename_map["N_image"] = "z" # TIFF stacks typically use 'z' for the stack dimension if rename_map: transmission = transmission.rename_dims(rename_map) daqmetadata = { "facility": "HFIR", "instrument": "MARS", "detector_type": "TPX3", "source_type": "neutron", } # Combine all masks and broadcast to the shape of the transmission data. # Mask must be same shape as the image data for scitiff. if transmission.masks: combined_mask = np.zeros_like(transmission.values, dtype=bool) for mask in transmission.masks.values(): combined_mask |= mask.values # remove other masks transmission.masks.clear() # add combined mask back in with name "scitiff-mask" transmission.masks["scitiff-mask"] = sc.array(dims=transmission.dims, values=combined_mask, dtype=bool) write_tiff_stack(output_path, transmission, metadata=metadata, daqmetadata=daqmetadata) else: raise ValueError(f"Unsupported output file format: {output_path.suffix}") logger.success("MARS TPX3 pipeline completed successfully. Output written to {}", output_path) return transmission