From abde581146ec3e56f2ce7951f321f50db19e26af Mon Sep 17 00:00:00 2001 From: Evgeniy Unzhakov Date: Mon, 5 Oct 2026 14:31:56 +0300 Subject: [PATCH] refactor: remove exp_data experimental data package Detector data handling, energy calibration, peak fitting and instrumental corrections leave the project scope (theory-only spectrum shape calculation). Removes the package, its tests, the exp pipeline integration test, the calibration notebook, and the ExpSpectrum fixtures from the shared conftest. --- exp_data/__init__.py | 42 -- exp_data/calibration.py | 416 ------------------- exp_data/corrections.py | 268 ------------- exp_data/exp_cli.py | 421 -------------------- exp_data/fitters.py | 307 -------------- exp_data/livetime.py | 219 ---------- exp_data/pipeline.py | 307 -------------- exp_data/raw_reader.py | 283 ------------- exp_data/root_io.py | 287 ------------- exp_data/spectrum.py | 211 ---------- notebooks/exp_calibration.ipynb | 321 --------------- tests/conftest.py | 44 -- tests/exp_data/conftest.py | 55 --- tests/exp_data/test_calibration.py | 138 ------- tests/exp_data/test_corrections.py | 106 ----- tests/exp_data/test_fitters.py | 114 ------ tests/exp_data/test_livetime.py | 142 ------- tests/exp_data/test_root_io.py | 84 ---- tests/exp_data/test_spectrum.py | 149 ------- tests/integration/test_exp_data_pipeline.py | 380 ------------------ 20 files changed, 4294 deletions(-) delete mode 100644 exp_data/__init__.py delete mode 100644 exp_data/calibration.py delete mode 100644 exp_data/corrections.py delete mode 100644 exp_data/exp_cli.py delete mode 100644 exp_data/fitters.py delete mode 100644 exp_data/livetime.py delete mode 100644 exp_data/pipeline.py delete mode 100644 exp_data/raw_reader.py delete mode 100644 exp_data/root_io.py delete mode 100644 exp_data/spectrum.py delete mode 100644 notebooks/exp_calibration.ipynb delete mode 100644 tests/exp_data/conftest.py delete mode 100644 tests/exp_data/test_calibration.py delete mode 100644 tests/exp_data/test_corrections.py delete mode 100644 tests/exp_data/test_fitters.py delete mode 100644 tests/exp_data/test_livetime.py delete mode 100644 tests/exp_data/test_root_io.py delete mode 100644 tests/exp_data/test_spectrum.py delete mode 100644 tests/integration/test_exp_data_pipeline.py diff --git a/exp_data/__init__.py b/exp_data/__init__.py deleted file mode 100644 index 9bdbc99..0000000 --- a/exp_data/__init__.py +++ /dev/null @@ -1,42 +0,0 @@ -""" -Experimental data handling for beta spectrometry. - -Provides tools for loading, calibrating, and pre-processing -experimental spectra from detector readouts. -""" - -from .spectrum import ExpSpectrum -from .calibration import EnergyCalibrator, CalibrationResult -from .fitters import GaussianFitter, PeakFitter -from .corrections import ( - DeadTimeCorrection, - PileUpCorrection, - BackgroundSubtractor, -) -from .livetime import LivetimeDetermination, LivetimeResult -from .pipeline import ( - SpectrumPipeline, - PipelineResult, - process_spectrum, - process_run_sequence, -) -from .root_io import write_spectrum, read_spectrum - -__all__ = [ - "ExpSpectrum", - "EnergyCalibrator", - "CalibrationResult", - "GaussianFitter", - "PeakFitter", - "DeadTimeCorrection", - "PileUpCorrection", - "BackgroundSubtractor", - "LivetimeDetermination", - "LivetimeResult", - "SpectrumPipeline", - "PipelineResult", - "process_spectrum", - "process_run_sequence", - "write_spectrum", - "read_spectrum", -] diff --git a/exp_data/calibration.py b/exp_data/calibration.py deleted file mode 100644 index a0bcc74..0000000 --- a/exp_data/calibration.py +++ /dev/null @@ -1,416 +0,0 @@ -""" -Energy calibration for beta spectrometers. - -Provides methods to calibrate detector channel numbers to physical -energy values using known spectral lines (X-rays, gamma-rays). - -Supports: -- Linear calibration: E = a + b · channel -- Quadratic calibration: E = a + b · channel + c · channel² -- Multi-point calibration with peak fitting -""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Optional - -import numpy as np -from scipy.signal import find_peaks - - -@dataclass -class CalibrationResult: - """Result of an energy calibration fit.""" - - coefficients: np.ndarray - """Calibration coefficients [a, b, c, ...] depending on polynomial order.""" - - order: int - """Polynomial order (1 = linear, 2 = quadratic).""" - - chi2_per_dof: float - """Reduced χ² of the calibration fit.""" - - n_points: int - """Number of calibration points used.""" - - covariance: np.ndarray - """Parameter covariance matrix.""" - - def __call__(self, channel: float | np.ndarray) -> np.ndarray: - """Convert channel number(s) to energy (keV).""" - return np.polyval(self.coefficients, channel) - - def invert(self) -> np.poly1d: - """Return polynomial that converts energy → channel.""" - if self.order == 1: - # coefficients from polyfit: [slope, intercept] - b, a = self.coefficients - # E = a + b*ch => ch = (E - a) / b - # np.poly1d([1/b, -a/b]) => (1/b)*E + (-a/b) = (E-a)/b ✓ - return np.poly1d([1.0 / b, -a / b]) - else: - # For quadratic, compute inverse numerically via interpolation - ch_vals = np.linspace(-10, 200, 1000) - e_vals = np.polyval(self.coefficients[::-1], ch_vals) - from scipy.interpolate import interp1d - - f = interp1d(e_vals, ch_vals, kind="linear") - - class _InversePoly: - def __init__(self, fn: Any) -> None: - self._fn = fn - - def __call__(self, x: float | np.ndarray) -> np.ndarray: - return self._fn(x) - - return _InversePoly(f) - - def resolution_at(self, energy_keV: float, sigma_channel: float = 1.0) -> float: - """Estimate energy resolution (σ) at a given energy. - - Parameters - ---------- - energy_keV : float - Energy in keV. - sigma_channel : float - Channel width in channel units (typically 1.0 for integer channels). - - Returns - ------- - float - Energy resolution σ in keV. - """ - # dE/dch = derivative of calibration polynomial at energy_keV - # coefficients are stored as [a_n, ..., a_1, a_0] - deriv_coeffs = [] - coeffs = self.coefficients - for i in range(len(coeffs) - 1): - deriv_coeffs.append(coeffs[i] * (len(coeffs) - i - 1)) - if not deriv_coeffs: - return abs(sigma_channel) - dE_dch = np.polyval(deriv_coeffs, energy_keV) - return abs(dE_dch) * sigma_channel - - -class EnergyCalibrator: - """Energy calibration for detector spectra. - - Parameters - ---------- - channel_energies : np.ndarray - Energy values (keV) per channel, shape (N_channels,). - If None, channels are assumed to be integer indices. - """ - - def __init__(self, channel_energies: Optional[np.ndarray] = None): - self._channel_energies = channel_energies - self._calibration: Optional[CalibrationResult] = None - self._calibration_points: Optional[list[tuple[float, float]]] = None - - @property - def is_calibrated(self) -> bool: - """Whether a calibration has been applied.""" - return self._calibration is not None - - def add_calibration_point( - self, - channel: int | float, - energy_keV: float, - ) -> None: - """Add a calibration point (known line energy → channel number). - - Parameters - ---------- - channel : int or float - Detector channel number. - energy_keV : float - Known energy of the calibration line (keV). - """ - if self._calibration_points is None: - self._calibration_points = [] - self._calibration_points.append((float(channel), float(energy_keV))) - - def set_calibration_points(self, points: list[tuple[float, float]]) -> None: - """Set calibration points directly. - - Parameters - ---------- - points : list of (channel, energy_keV) - Known calibration line positions. - """ - self._calibration_points = [(float(c), float(e)) for c, e in points] - - def fit_calibration( - self, - order: int = 1, - channel_energies: Optional[np.ndarray] = None, - ) -> CalibrationResult: - """Fit calibration polynomial using known points. - - Parameters - ---------- - order : int - Polynomial order (1 = linear, 2 = quadratic). - channel_energies : np.ndarray, optional - Full channel-to-energy mapping. If provided, fits to - all points; otherwise fits only to added calibration points. - - Returns - ------- - CalibrationResult - """ - if channel_energies is not None: - # Fit to full mapping - channels = np.arange(len(channel_energies), dtype=np.float64) - x = channels - y = np.asarray(channel_energies, dtype=np.float64) - elif self._calibration_points is not None: - x = np.array([p[0] for p in self._calibration_points]) - y = np.array([p[1] for p in self._calibration_points]) - else: - raise ValueError( - "No calibration data available. Use add_calibration_point() " - "or pass channel_energies." - ) - - # Fit polynomial: E(ch) = a₀ + a₁·ch + a₂·ch² + ... - try: - coeffs, cov = np.polyfit(x, y, order, cov=True) - except ValueError: - # Too few points for covariance estimation - coeffs = np.polyfit(x, y, order) - cov = np.eye(order + 1) * 1e-10 - - # Compute χ² - # np.polyfit returns descending-order coefficients [a_n, ..., a_0], - # which is exactly what np.polyval expects. - y_fit = np.polyval(coeffs, x) - # Assume uniform uncertainty of 0.1 keV for calibration points - uncertainties = np.full_like(x, 0.1) - chi2 = float(np.sum(((y - y_fit) / uncertainties) ** 2)) - n_dof = max(len(x) - order - 1, 1) - - return CalibrationResult( - coefficients=coeffs, - order=order, - chi2_per_dof=chi2 / n_dof, - n_points=len(x), - covariance=cov, - ) - - def apply_calibration( - self, - order: int = 1, - channel_energies: Optional[np.ndarray] = None, - ) -> CalibrationResult: - """Fit and store calibration for future use. - - See :meth:`fit_calibration` for parameters. - """ - self._calibration = self.fit_calibration(order, channel_energies) - return self._calibration - - def channel_to_energy(self, channel: float | np.ndarray) -> np.ndarray: - """Convert channel number to energy using current calibration. - - Parameters - ---------- - channel : float or np.ndarray - Channel number(s). - - Returns - ------- - np.ndarray - Energy in keV. - - Raises - ------ - RuntimeError - If no calibration has been applied. - """ - if self._calibration is None: - raise RuntimeError( - "No calibration applied. Call apply_calibration() first." - ) - return self._calibration(channel) - - def energy_to_channel(self, energy_keV: float | np.ndarray) -> np.ndarray: - """Convert energy to channel using current calibration. - - Parameters - ---------- - energy_keV : float or np.ndarray - Energy in keV. - - Returns - ------- - np.ndarray - Channel number(s). - - Raises - ------ - RuntimeError - If no calibration has been applied. - """ - if self._calibration is None: - raise RuntimeError( - "No calibration applied. Call apply_calibration() first." - ) - inv = self._calibration.invert() - return inv(energy_keV) - - def get_linear_calibration( - self, - low_energy_keV: float, - high_energy_keV: float, - low_channel: int, - high_channel: int, - ) -> CalibrationResult: - """Create a linear calibration from two known points. - - Parameters - ---------- - low_energy_keV : float - Energy of first calibration line (keV). - high_energy_keV : float - Energy of second calibration line (keV). - low_channel : int - Channel number of first line. - high_channel : int - Channel number of second line. - - Returns - ------- - CalibrationResult - """ - self.set_calibration_points( - [ - (float(low_channel), low_energy_keV), - (float(high_channel), high_energy_keV), - ] - ) - return self.apply_calibration(order=1) - - @classmethod - def from_am241_spectrum( - cls, - filepath: str | Path, - known_energies: Optional[list[float]] = None, - n_channels: int = 4096, - ) -> CalibrationResult: - """Create a CalibrationResult from an Am-241 calibration spectrum. - - Auto-detects peaks using scipy.signal.find_peaks(), matches to known Am-241 - lines, fits linear calibration, and returns a CalibrationResult suitable - for ExpSpectrum.apply_calibration(). - - Parameters - ---------- - filepath : str or Path - Path to the Am-241 binary calibration file (CAM1.DAT or CAM2.DAT). - known_energies : list of float, optional - Known Am-241 line energies in keV. Defaults to [13.9, 17.8, 26.3, 59.54]. - n_channels : int - Number of channels (default 4096). - - Returns - ------- - CalibrationResult - Linear calibration fit with coefficients, chi2_per_dof, and covariance. - - Raises - ------ - ValueError - If fewer than 2 peaks are matched. - """ - import struct - - filepath = Path(filepath) - - # Read raw binary - with open(filepath, "rb") as f: - f.read(128) # skip header - counts = np.array( - struct.unpack(f"{n_channels}i", f.read(n_channels * 4)), dtype=np.int64 - ) - - # Search for prominent gamma peaks across the full spectrum. - # Exclude the high-energy pulser region (last ~200 channels). - counts_search = counts.copy() - pulser_cutoff = max(n_channels - 200, 300) - counts_search[pulser_cutoff:] = 0 - - # Find peaks with moderate thresholds. - # height=100 skips noise floor; distance=5 keeps nearby peaks separate; - # width=5 constrains peaks to realistic FWHM (~5-30 channels for - # typical beta spectrometers), rejecting noise bumps and pulser tails. - peaks, _ = find_peaks(counts_search, height=100, distance=5, width=(5, 50)) - peak_energies = np.array(peaks, dtype=float) - peak_counts = counts_search[peaks] - - # Known Am-241 gamma lines (keV) - # Am-241 decays to Np-237; Np K X-rays at 13.9 keV (Kα), 17.8 keV (Kβ), - # 20.8 keV (L-series); Np gamma at 26.3 keV; Np/Am gamma at 59.54 keV. - # Source: LNHb Table de Radionucléides, IAEA nuclear data. - if known_energies is None: - known_energies = [13.9, 17.8, 26.3, 59.54] - - # Match peaks to known energies using a simpler, more robust - # approach: find the two strongest peaks that span the largest - # channel range, then assign the lowest and highest known energies - # to them. This avoids the combinatorial explosion of trying all - # energy pairs and the ambiguity of gain-based scoring. - channel_to_energy: dict[float, float] = {} - - if len(peak_energies) >= 2 and len(known_energies) >= 2: - strong_threshold = np.percentile(counts_search, 95) - - # Filter to strong peaks only. - # Use a higher absolute threshold for the low-energy end - # to avoid noise bumps being selected as calibration anchors. - # The 95th percentile is ~171, but real Am-241 peaks are >400. - absolute_threshold = max(strong_threshold, 400) - strong_peaks = [ - (peak_energies[i], peak_counts[i]) - for i in range(len(peak_energies)) - if peak_counts[i] >= absolute_threshold - ] - - if len(strong_peaks) >= 2: - # Find the pair with the largest channel separation - best_sep = 0 - best_pair = None - for i in range(len(strong_peaks)): - for j in range(i + 1, len(strong_peaks)): - ch1, _ = strong_peaks[i] - ch2, _ = strong_peaks[j] - sep = ch2 - ch1 - if sep > best_sep: - best_sep = sep - best_pair = (ch1, ch2) - - if best_pair is not None: - ch_low, ch_high = best_pair - e_min, e_max = min(known_energies), max(known_energies) - channel_to_energy[float(ch_low)] = e_min - channel_to_energy[float(ch_high)] = e_max - - if len(channel_to_energy) < 2: - raise ValueError( - f"Could not match at least 2 Am-241 peaks. " - f"Found {len(channel_to_energy)} of {len(known_energies)} expected. " - f"Check calibration file or provide known_energies." - ) - - # Build calibration points and fit using only the two strongest - # matched peaks for a clean 2-point linear fit. - sorted_points = sorted(channel_to_energy.items(), key=lambda x: x[0]) - ch_low, e_low = sorted_points[0] - ch_high, e_high = sorted_points[-1] - - calibrator = cls() - calibrator.set_calibration_points([(ch_low, e_low), (ch_high, e_high)]) - return calibrator.fit_calibration(order=1) diff --git a/exp_data/corrections.py b/exp_data/corrections.py deleted file mode 100644 index a7d35e9..0000000 --- a/exp_data/corrections.py +++ /dev/null @@ -1,268 +0,0 @@ -""" -Instrumental corrections for experimental spectra. - -Provides dead-time correction, pile-up correction, and background -subtraction for beta spectrometry data. -""" - -from __future__ import annotations - -from typing import Optional - -import numpy as np - - -class DeadTimeCorrection: - """Dead-time correction for paralyzable or non-paralyzable detectors. - - Parameters - ---------- - dead_time : float - Dead time per event in seconds. - model : str - 'non_paralyzable' (default) or 'paralyzable'. - """ - - def __init__(self, dead_time: float, model: str = "non_paralyzable"): - self.dead_time = dead_time - self.model = model - assert model in ( - "non_paralyzable", - "paralyzable", - ), f"Unknown dead-time model: {model}" - assert dead_time >= 0, "Dead time must be non-negative" - - def live_time_fraction(self, observed_rate: float) -> float: - """Calculate live time fraction given observed count rate. - - Parameters - ---------- - observed_rate : float - Observed count rate in counts/second. - - Returns - ------- - float - Fraction of time the detector is live (0 to 1). - """ - if self.model == "non_paralyzable": - return max(0.0, 1.0 - observed_rate * self.dead_time) - else: # paralyzable - return np.exp(-observed_rate * self.dead_time) - - def correct_rate(self, observed_rate: float) -> float: - """Correct observed rate for dead time. - - Parameters - ---------- - observed_rate : float - Observed count rate (counts/s). - - Returns - ------- - float - True count rate (counts/s). - """ - live_frac = self.live_time_fraction(observed_rate) - if live_frac <= 0: - return float("inf") - return observed_rate / live_frac - - def correct_counts(self, counts: np.ndarray, live_time: float) -> np.ndarray: - """Correct counts for dead time. - - Parameters - ---------- - counts : np.ndarray - Measured counts per bin. - live_time : float - Total live time in seconds. - - Returns - ------- - np.ndarray - Dead-time corrected counts. - """ - total_observed = np.sum(counts) - total_true = self.correct_rate(total_observed / live_time) * live_time - if total_observed <= 0: - return counts.copy() - return counts * (total_true / total_observed) - - -class PileUpCorrection: - """Simple pile-up correction for beta spectrometry. - - Estimates the fraction of pile-up events based on observed count rate - and dead time, and redistributes pile-up events back into the spectrum. - - This is a first-order approximation. For precise pile-up correction, - digital pulse processing or anti-pile-up electronics are recommended. - """ - - def __init__(self, dead_time: float): - self.dead_time = dead_time - - def estimate_pile_up_fraction(self, count_rate: float) -> float: - """Estimate pile-up fraction from count rate. - - For a non-paralyzable detector, pile-up fraction ≈ (count_rate × dead_time)². - - Parameters - ---------- - count_rate : float - Observed count rate in counts/second. - - Returns - ------- - float - Estimated fraction of events affected by pile-up. - """ - return (count_rate * self.dead_time) ** 2 - - def correct_spectrum( - self, - counts: np.ndarray, - live_time: float, - background: Optional[np.ndarray] = None, - ) -> np.ndarray: - """Apply first-order pile-up correction. - - Parameters - ---------- - counts : np.ndarray - Measured counts per bin. - live_time : float - Live time in seconds. - background : np.ndarray, optional - Background counts per bin. - - Returns - ------- - np.ndarray - Pile-up corrected counts. - """ - if background is not None: - net_counts = counts - background - else: - net_counts = counts.copy() - - observed_rate = np.sum(net_counts) / live_time - pileup_frac = self.estimate_pile_up_fraction(observed_rate) - - # Redistribute pile-up events proportionally to the spectrum shape - if pileup_frac > 0 and np.sum(net_counts) > 0: - correction = net_counts * pileup_frac / (1.0 - pileup_frac) - return net_counts + correction - - return net_counts - - -class BackgroundSubtractor: - """Background subtraction for beta spectra. - - Supports: - - Constant background (from sidebands) - - Polynomial background - - Region-of-interest based background estimation - """ - - @staticmethod - def constant_background( - counts: np.ndarray, - energies: np.ndarray, - low_roi: tuple[float, float], - high_roi: tuple[float, float], - ) -> np.ndarray: - """Subtract constant background estimated from sidebands. - - Parameters - ---------- - counts : np.ndarray - Measured counts per bin. - energies : np.ndarray - Energy values (keV) per bin. - low_roi : (E_min, E_max) - Energy range for low-side background estimation. - high_roi : (E_min, E_max) - Energy range for high-side background estimation. - - Returns - ------- - np.ndarray - Counts with constant background subtracted. - """ - low_mask = (energies >= low_roi[0]) & (energies <= low_roi[1]) - high_mask = (energies >= high_roi[0]) & (energies <= high_roi[1]) - - low_bg = np.mean(counts[low_mask]) if np.any(low_mask) else 0.0 - high_bg = np.mean(counts[high_mask]) if np.any(high_mask) else 0.0 - constant_bg = (low_bg + high_bg) / 2.0 - - return counts - constant_bg - - @staticmethod - def polynomial_background( - counts: np.ndarray, - energies: np.ndarray, - order: int = 2, - low_roi: tuple[float, float] = (0.0, 0.0), - high_roi: tuple[float, float] = (0.0, 0.0), - ) -> np.ndarray: - """Subtract polynomial background fitted from sidebands. - - Parameters - ---------- - counts : np.ndarray - Measured counts per bin. - energies : np.ndarray - Energy values (keV) per bin. - order : int - Polynomial order for background fit. - low_roi : (E_min, E_max) - Energy range for low-side background fitting. - high_roi : (E_min, E_max) - Energy range for high-side background fitting. - - Returns - ------- - np.ndarray - Counts with polynomial background subtracted. - """ - low_mask = (energies >= low_roi[0]) & (energies <= low_roi[1]) - high_mask = (energies >= high_roi[0]) & (energies <= high_roi[1]) - - mask = low_mask | high_mask - if not np.any(mask): - return counts - - x_fit = energies[mask] - y_fit = counts[mask] - - coeffs = np.polyfit(x_fit, y_fit, order) - bg = np.polyval(coeffs[::-1], energies) - - return counts - bg - - @staticmethod - def roi_average( - counts: np.ndarray, - roi_mask: np.ndarray, - ) -> float: - """Calculate average background from a region of interest. - - Parameters - ---------- - counts : np.ndarray - Measured counts per bin. - roi_mask : np.ndarray - Boolean mask selecting background region. - - Returns - ------- - float - Average background count per bin. - """ - if not np.any(roi_mask): - return 0.0 - return float(np.mean(counts[roi_mask])) diff --git a/exp_data/exp_cli.py b/exp_data/exp_cli.py deleted file mode 100644 index 51e2d00..0000000 --- a/exp_data/exp_cli.py +++ /dev/null @@ -1,421 +0,0 @@ -""" -Command-line interface for experimental data processing. - -Usage examples: - # Process a single spectrum - bs_exp process --raw A1.DAT --calibration CAM1.DAT --output output.root - - # Process a run sequence - bs_exp process --run-dir ./data/raw/ --prefix A --calibration CAM1.DAT - - # Determine live-time from a pulser calibration - bs_exp livetime --file CAM1.DAT --frequency 100.0 --real-time 3600.0 - - # Export calibrated spectrum to ROOT - bs_exp export --input spectrum.root --output processed.root - - # Dry run to inspect pipeline configuration - bs_exp process --raw A1.DAT --calibration CAM1.DAT --dry-run -""" - -from __future__ import annotations - -import argparse -import sys - - -def _build_parser() -> argparse.ArgumentParser: - """Build the argument parser.""" - parser = argparse.ArgumentParser( - prog="bs_exp", - description="Experimental data processing toolkit - " - "calibrate, process, and export beta decay spectra.", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=( - "Examples:\n" - " bs_exp process --raw A1.DAT --calibration CAM1.DAT --output output.root\n" - " bs_exp process --run-dir ./data/ --prefix A --calibration CAM1.DAT\n" - " bs_exp livetime --file CAM1.DAT --frequency 100.0 --real-time 3600.0\n" - " bs_exp export --input spectrum.root --output processed.root\n" - " bs_exp process --raw A1.DAT --calibration CAM1.DAT --dry-run\n" - ), - ) - - # Version - parser.add_argument( - "--version", - action="version", - version="%(prog)s 0.4.0", - ) - - # Subcommands - subparsers = parser.add_subparsers(dest="command", help="Available commands") - - # --- process subcommand --- - process_parser = subparsers.add_parser( - "process", - help="Process raw binary spectra through the full pipeline", - description=( - "Read raw binary spectra, subtract cumulative runs, " - "apply energy calibration, determine live-time, and export." - ), - ) - process_parser.add_argument( - "--raw", - type=str, - metavar="FILE", - help="Path to the raw binary spectrum file (A*.DAT).", - ) - process_parser.add_argument( - "--cumulative", - type=str, - metavar="FILE", - help="Path to the previous cumulative spectrum for subtraction.", - ) - process_parser.add_argument( - "--run-dir", - type=str, - metavar="DIR", - help="Directory containing cumulative run files (for sequence processing).", - ) - process_parser.add_argument( - "--prefix", - type=str, - default="A", - help="Filename prefix for run files (default: 'A').", - ) - process_parser.add_argument( - "--calibration", - type=str, - required=True, - metavar="FILE", - help="Path to the Am-241 calibration binary file (CAM1.DAT).", - ) - process_parser.add_argument( - "--pulser-frequency", - type=float, - default=None, - metavar="HZ", - help="Pulser injection frequency in Hz (for live-time determination).", - ) - process_parser.add_argument( - "--real-time", - type=float, - default=None, - metavar="SEC", - help="Real acquisition time in seconds (for live-time determination).", - ) - process_parser.add_argument( - "--output", - type=str, - metavar="FILE", - help="Output ROOT file path.", - ) - process_parser.add_argument( - "--run-id", - type=str, - default=None, - help="Run identifier prefix.", - ) - process_parser.add_argument( - "--source", - type=str, - default=None, - help="Source identifier (e.g., 'Tc99').", - ) - - # --- livetime subcommand --- - livetime_parser = subparsers.add_parser( - "livetime", - help="Determine live-time fraction from a pulser calibration spectrum", - description=( - "Read a pulser calibration spectrum, identify the pulser " - "peak, and compute the live-time fraction." - ), - ) - livetime_parser.add_argument( - "--file", - type=str, - required=True, - metavar="FILE", - help="Path to the pulser calibration binary file.", - ) - livetime_parser.add_argument( - "--frequency", - type=float, - required=True, - metavar="HZ", - help="Pulser injection frequency in Hz.", - ) - livetime_parser.add_argument( - "--real-time", - type=float, - required=True, - metavar="SEC", - help="Real acquisition time in seconds.", - ) - livetime_parser.add_argument( - "--output", - type=str, - metavar="FILE", - help="Output ROOT file with livetime metadata.", - ) - - # --- export subcommand --- - export_parser = subparsers.add_parser( - "export", - help="Export calibrated ExpSpectrum to ROOT format", - description="Export an ExpSpectrum (from ROOT or dict) to ROOT file format.", - ) - export_parser.add_argument( - "--input", - type=str, - required=True, - metavar="FILE", - help="Input file (ROOT or JSON dict).", - ) - export_parser.add_argument( - "--output", - type=str, - required=True, - metavar="FILE", - help="Output ROOT file path.", - ) - - # --- dry-run (global, added to all subparsers) --- - for sub in [process_parser, livetime_parser, export_parser]: - sub.add_argument( - "--dry-run", - action="store_true", - help="Validate input and display configuration without processing.", - ) - - # Logging and output (global) - log_group = parser.add_argument_group("logging and output") - log_group.add_argument( - "-v", - "--verbose", - action="count", - default=0, - help="Increase verbosity: -v=INFO, -vv=DEBUG. Repeatable.", - ) - log_group.add_argument( - "-q", - "--quiet", - action="store_true", - help="Suppress all terminal output (only errors shown).", - ) - log_group.add_argument( - "--log-file", - type=str, - metavar="PATH", - help="Write log messages to the specified file.", - ) - - return parser - - -def _run(args: argparse.Namespace) -> None: - """Execute the CLI command.""" - # Determine logging level - if args.quiet: - log_level = "WARNING" - elif args.verbose >= 2: - log_level = "DEBUG" - elif args.verbose >= 1: - log_level = "INFO" - else: - log_level = "WARNING" - - # Setup logging (minimal - reuse beta_spectrum.logging_utils if available) - try: - from beta_spectrum.logging_utils import setup_logging, get_logger - - log_config = type( - "LoggingConfig", (), {"level": log_level, "log_file": args.log_file} - )() - logger = setup_logging(log_config) - logger = get_logger("exp") - except ImportError: - import logging - - logging.basicConfig(level=getattr(logging, log_level)) - logger = logging.getLogger("exp") - - logger.info("bs_exp starting...") - - if args.command == "process": - _cmd_process(args, logger) - elif args.command == "livetime": - _cmd_livetime(args, logger) - elif args.command == "export": - _cmd_export(args, logger) - else: - parser = _build_parser() - parser.print_help() - sys.exit(1) - - -def _cmd_process(args, logger) -> None: - """Process subcommand handler.""" - from exp_data.pipeline import process_spectrum, process_run_sequence - from exp_data.root_io import write_spectrum - - # Check PyROOT/uproot availability - from exp_data.root_io import _PYROOT_AVAILABLE, _UPTOOL_AVAILABLE - - if not _PYROOT_AVAILABLE and not _UPTOOL_AVAILABLE: - logger.error("No ROOT I/O library available. Install: pip install uproot") - sys.exit(1) - - if args.dry_run: - logger.info("Dry run: pipeline configuration validated") - logger.info(" Calibration: %s", args.calibration) - logger.info(" Pulser frequency: %s Hz", args.pulser_frequency_hz or "N/A") - logger.info(" Real time: %s s", args.real_time or "N/A") - return - - if args.raw: - # Single spectrum - result = process_spectrum( - raw_path=args.raw, - calibration_path=args.calibration, - cumulative_path=args.cumulative, - run_id=args.run_id, - source=args.source, - pulser_frequency_hz=args.pulser_frequency_hz, - real_time_sec=args.real_time, - ) - - logger.info("Pipeline complete:") - logger.info( - " Validation: %s", "PASSED" if result.validation_passed else "FAILED" - ) - for msg in result.validation_messages: - logger.info(" %s", msg) - logger.info(" Calibration chi2/dof: %.4f", result.calibration.chi2_per_dof) - - if result.livetime: - logger.info( - " Live-time fraction: %.4f", result.livetime.live_time_fraction - ) - - if args.output: - write_spectrum(args.output, result.spectrum) - logger.info(" Exported to %s", args.output) - - elif args.run_dir: - # Run sequence - results = process_run_sequence( - run_dir=args.run_dir, - calibration_path=args.calibration, - prefix=args.prefix, - run_id=args.run_id, - source=args.source, - pulser_frequency_hz=args.pulser_frequency_hz, - real_time_sec=args.real_time, - ) - - logger.info("Processed %d runs", len(results)) - for i, result in enumerate(results): - logger.info( - " Run %d: %s (chi2/dof=%.4f)", - i + 1, - "PASSED" if result.validation_passed else "FAILED", - result.calibration.chi2_per_dof, - ) - - if args.output: - # Export last result - write_spectrum(args.output, results[-1].spectrum) - logger.info(" Exported last run to %s", args.output) - else: - logger.error("Specify --raw FILE or --run-dir DIR") - sys.exit(1) - - -def _cmd_livetime(args, logger) -> None: - """Livetime subcommand handler.""" - from exp_data.livetime import LivetimeDetermination - from exp_data.root_io import write_spectrum - from exp_data.spectrum import ExpSpectrum - - result = LivetimeDetermination.from_file( - filepath=args.file, - pulser_frequency_hz=args.frequency, - real_time_sec=args.real_time, - source="pulser", - ) - - logger.info("Live-time determination:") - logger.info(" Live-time fraction: %.4f", result.live_time_fraction) - logger.info(" Observed pulser rate: %.1f Hz", result.observed_pulser_rate) - logger.info( - " Pulser peak: %.2f keV (sigma=%.2f keV, chi2/dof=%.2f)", - result.pulser_peak_center_keV, - result.pulser_peak_sigma_keV, - result.pulser_peak_chi2, - ) - - if args.output: - # Create a dummy spectrum with livetime metadata - import numpy as np - - dummy = ExpSpectrum( - energies=np.array([0.0]), - counts=np.array([0.0]), - errors=np.array([0.0]), - metadata={ - "livetime": { - "fraction": result.live_time_fraction, - "pulser_frequency_hz": result.pulser_frequency_hz, - "observed_pulser_rate": result.observed_pulser_rate, - } - }, - live_time=result.real_time_sec * result.live_time_fraction, - source="pulser", - ) - write_spectrum(args.output, dummy) - logger.info(" Livetime metadata exported to %s", args.output) - - -def _cmd_export(args, logger) -> None: - """Export subcommand handler.""" - from exp_data.root_io import write_spectrum, read_spectrum - - # Read input spectrum - if str(args.input).endswith(".root"): - spectrum = read_spectrum(args.input) - elif str(args.input).endswith(".json"): - import json - from exp_data.spectrum import ExpSpectrum - - with open(args.input) as f: - data = json.load(f) - spectrum = ExpSpectrum.from_dict(data) - else: - logger.error("Unsupported input format. Use .root or .json") - sys.exit(1) - - # Write output - write_spectrum(args.output, spectrum) - logger.info("Exported to %s", args.output) - - -def main() -> None: - """Main entry point for the CLI.""" - parser = _build_parser() - args = parser.parse_args() - - try: - _run(args) - except KeyboardInterrupt: - print("\nInterrupted by user.", file=sys.stderr) - sys.exit(130) - except Exception as e: - print(f"Error: {e}", file=sys.stderr) - sys.exit(1) - - -if __name__ == "__main__": - main() diff --git a/exp_data/fitters.py b/exp_data/fitters.py deleted file mode 100644 index 051aeb1..0000000 --- a/exp_data/fitters.py +++ /dev/null @@ -1,307 +0,0 @@ -""" -Simple peak fitters for calibration source analysis. - -Provides Gaussian and Gaussian+background fitting for identifying -known X-ray and gamma-ray lines used in energy calibration. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Optional, Tuple - -import numpy as np -from scipy.optimize import least_squares - - -@dataclass -class PeakFitResult: - """Result of a peak fit.""" - - center_keV: float - """Peak center energy in keV.""" - - sigma_keV: float - """Peak width (σ) in keV.""" - - amplitude: float - """Peak amplitude (counts at peak).""" - - background: float - """Constant background level.""" - - chi2: float - """Reduced χ².""" - - n_dof: int - """Number of degrees of freedom.""" - - covariance: np.ndarray - """Parameter covariance matrix. - - Parameter order: [amplitude, center, sigma, background] - """ - - success: bool - """Whether the fit converged.""" - - def resolution(self) -> float: - """FWHM in keV.""" - return 2.35482 * self.sigma_keV - - def resolution_fraction(self) -> float: - """FWHM / center as a fraction.""" - if self.center_keV > 0: - return self.resolution() / self.center_keV - return 0.0 - - -class GaussianFitter: - """Fit a Gaussian peak on top of constant background. - - Model: f(x) = A · exp(-(x - μ)² / (2σ²)) + B - - Parameters - ---------- - energies : np.ndarray - Energy values (keV), shape (N,). - counts : np.ndarray - Measured counts, shape (N,). - uncertainties : np.ndarray, optional - Measurement uncertainties. If None, assumed to be 1.0. - """ - - def __init__( - self, - energies: np.ndarray, - counts: np.ndarray, - uncertainties: Optional[np.ndarray] = None, - ): - self.energies = np.asarray(energies, dtype=np.float64) - self.counts = np.asarray(counts, dtype=np.float64) - - if uncertainties is not None: - self.uncertainties = np.asarray(uncertainties, dtype=np.float64) - else: - self.uncertainties = np.ones_like(self.counts) - - # Filter out zero-uncertainty points - valid = self.uncertainties > 0 - self.energies = self.energies[valid] - self.counts = self.counts[valid] - self.uncertainties = self.uncertainties[valid] - - def _gaussian_background(self, params: np.ndarray, x: np.ndarray) -> np.ndarray: - """Gaussian + constant background model.""" - A, mu, sigma, B = params - return A * np.exp(-0.5 * ((x - mu) / sigma) ** 2) + B - - def _residuals(self, params: np.ndarray) -> np.ndarray: - """Weighted residuals.""" - model = self._gaussian_background(params, self.energies) - return (self.counts - model) / self.uncertainties - - def _chi2(self, params: np.ndarray) -> float: - """Compute χ².""" - residuals = self._residuals(params) - return float(np.sum(residuals**2)) - - def _compute_covariance(self, params: np.ndarray, jac: np.ndarray) -> np.ndarray: - """Estimate covariance from Jacobian.""" - jtj = jac.T @ jac - if np.linalg.cond(jtj) < 1e12: - cov = np.linalg.inv(jtj) - chi2 = self._chi2(params) - ndof = max(len(self.energies) - 4, 1) - cov = cov * chi2 / ndof - return cov - return np.diag([1e10, 1e6, 1e4, 1e6]) - - def fit( - self, - guess: Optional[Tuple[float, float, float, float]] = None, - bounds: Optional[Tuple[tuple, tuple]] = None, - ) -> PeakFitResult: - """Fit a Gaussian peak to the data. - - Parameters - ---------- - guess : (amplitude, center, sigma, background), optional - Initial parameter guess. Auto-estimated if None. - bounds : ((A_min, μ_min, σ_min, B_min), (A_max, μ_max, σ_max, B_max)), optional - Parameter bounds. - - Returns - ------- - PeakFitResult - """ - if guess is None: - guess = self._estimate_guess() - - if bounds is None: - bounds = ( - (0.0, min(self.energies), 0.1, 0.0), - (np.inf, max(self.energies), 50.0, np.inf), - ) - - result = least_squares( - self._residuals, - x0=guess, - bounds=bounds, - method="trf", - max_nfev=1000, - ) - - A, mu, sigma, B = result.x - n_dof = max(len(self.energies) - 4, 1) - chi2 = self._chi2(result.x) - - # Compute covariance from Jacobian - _, jac = result.fun, result.jac - cov = self._compute_covariance(result.x, jac) if jac is not None else np.eye(4) - - return PeakFitResult( - center_keV=mu, - sigma_keV=sigma, - amplitude=A, - background=B, - chi2=chi2 / n_dof if n_dof > 0 else chi2, - n_dof=n_dof, - covariance=cov, - success=result.success, - ) - - def _estimate_guess(self) -> Tuple[float, float, float, float]: - """Auto-estimate initial parameters from data.""" - # Find the peak - peak_idx = np.argmax(self.counts) - center = self.energies[peak_idx] - amplitude = self.counts[peak_idx] - - # Estimate sigma from FWHM - half_max = amplitude / 2.0 - left_idx = np.argmax(self.counts[:peak_idx] < half_max) - right_idx = np.argmin(self.counts[peak_idx:] < half_max) + peak_idx - - if left_idx < peak_idx and right_idx > peak_idx: - fwhm = self.energies[right_idx] - self.energies[left_idx] - sigma = fwhm / 2.35482 - else: - sigma = 1.0 # Default width - - # Background estimate from average of edges - n_edge = max(10, len(self.counts) // 20) - background = np.mean( - np.concatenate([self.counts[:n_edge], self.counts[-n_edge:]]) - ) - - return (amplitude - background, center, max(sigma, 0.1), background) - - -class PeakFitter: - """Find and fit multiple peaks in a spectrum. - - Parameters - ---------- - energies : np.ndarray - Energy values (keV). - counts : np.ndarray - Measured counts. - uncertainties : np.ndarray, optional - Measurement uncertainties. - """ - - def __init__( - self, - energies: np.ndarray, - counts: np.ndarray, - uncertainties: Optional[np.ndarray] = None, - ): - self.energies = np.asarray(energies, dtype=np.float64) - self.counts = np.asarray(counts, dtype=np.float64) - - if uncertainties is not None: - self.uncertainties = np.asarray(uncertainties, dtype=np.float64) - else: - self.uncertainties = np.sqrt(np.maximum(self.counts, 1.0)) - - def find_peaks( - self, - min_height: float = 10.0, - min_distance_keV: float = 2.0, - smoothing_window: int = 5, - ) -> np.ndarray: - """Find candidate peak positions. - - Parameters - ---------- - min_height : float - Minimum peak height above background (counts). - min_distance_keV : float - Minimum separation between peaks in keV. - smoothing_window : int - Smoothing window for peak detection (odd number). - - Returns - ------- - np.ndarray - Energy values of detected peak centers. - """ - from scipy.signal import find_peaks - - # Estimate background by smoothing - window = smoothing_window if smoothing_window % 2 == 1 else smoothing_window + 1 - if len(self.counts) > window: - from scipy.ndimage import uniform_filter1d - - background_est = uniform_filter1d(self.counts, size=window) - else: - background_est = np.full_like(self.counts, np.mean(self.counts)) - - heights = self.counts - background_est - indices, _ = find_peaks( - heights, - height=min_height, - distance=int(min_distance_keV / max(1.0, np.diff(self.energies).min())), - ) - - return self.energies[indices] - - def fit_peaks( - self, - peak_energies: Optional[np.ndarray] = None, - known_energies: Optional[np.ndarray] = None, - ) -> list[PeakFitResult]: - """Fit peaks at given or auto-detected positions. - - Parameters - ---------- - peak_energies : np.ndarray, optional - Explicit peak center energies to fit. - known_energies : np.ndarray, optional - Known calibration line energies — auto-detect peaks near these. - - Returns - ------- - list[PeakFitResult] - """ - if known_energies is not None: - # Find peaks near known energies - peak_energies = np.array([self._nearest_channel(e) for e in known_energies]) - - if peak_energies is None: - peak_energies = self.find_peaks() - - results = [] - for center in peak_energies: - fitter = GaussianFitter(self.energies, self.counts, self.uncertainties) - result = fitter.fit(guess=(10.0, center, 1.0, 0.0)) - if result.success: - results.append(result) - - return results - - def _nearest_channel(self, target_keV: float) -> float: - """Find the energy channel closest to target_keV.""" - idx = np.argmin(np.abs(self.energies - target_keV)) - return self.energies[idx] diff --git a/exp_data/livetime.py b/exp_data/livetime.py deleted file mode 100644 index c8e2170..0000000 --- a/exp_data/livetime.py +++ /dev/null @@ -1,219 +0,0 @@ -""" -Live-time determination from pulser calibration spectra. - -Uses Gaussian peak fitting to identify the pulser peak in a calibration -spectrum and computes the live-time fraction from the observed peak -count rate compared to the known injection frequency. - -Usage: - from exp_data.livetime import LivetimeDetermination - - # From an existing ExpSpectrum - determiner = LivetimeDetermination( - spectrum=pulser_spectrum, - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) - result = determiner.run() - - # Or directly from a binary file - result = LivetimeDetermination.from_file( - filepath="data/raw/CAM1.DAT", - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Optional - -import numpy as np - -from exp_data.spectrum import ExpSpectrum -from exp_data.raw_reader import read_raw_spectrum -from exp_data.fitters import GaussianFitter - - -@dataclass -class LivetimeResult: - """Result of a live-time determination.""" - - live_time_fraction: float - """Live-time fraction tau (0.0 to 1.0).""" - - real_time_sec: float - """Total real acquisition time in seconds.""" - - pulser_frequency_hz: float - """Known pulser injection frequency in Hz.""" - - observed_pulser_rate: float - """Observed pulser peak count rate in Hz.""" - - pulser_peak_center_keV: float - """Energy of the pulser peak in keV (from Gaussian fit).""" - - pulser_peak_sigma_keV: float - """Width of the pulser peak in keV.""" - - pulser_peak_chi2: float - """Reduced chi-squared of the pulser peak fit.""" - - metadata: dict = field(default_factory=dict) - """Additional metadata from the determination.""" - - -class LivetimeDetermination: - """Determine live-time fraction from a pulser calibration spectrum. - - The pulser injects known-frequency pulses into the MCA. The observed - peak count rate in the pulser channel, compared to the known injection - frequency, gives the live-time fraction: - - tau = observed_pulser_rate / expected_pulser_frequency - - Parameters - ---------- - spectrum : ExpSpectrum - Calibration spectrum with pulser input (energies in keV). - pulser_frequency_hz : float - Known pulser injection frequency in Hz. - real_time_sec : float - Total real acquisition time in seconds (from MCA metadata). - pulser_channel_guess : float, optional - Expected pulser peak channel number (for narrowing search). - If None, searches the full spectrum. - """ - - def __init__( - self, - spectrum: ExpSpectrum, - pulser_frequency_hz: float, - real_time_sec: float, - pulser_channel_guess: Optional[float] = None, - ): - self.spectrum = spectrum - self.pulser_frequency_hz = pulser_frequency_hz - self.real_time_sec = real_time_sec - self.pulser_channel_guess = pulser_channel_guess - - def run(self) -> LivetimeResult: - """Determine live-time fraction from the pulser peak. - - Returns - ------- - LivetimeResult - Live-time fraction and associated metadata. - - Raises - ------ - ValueError - If no pulser peak is found or fit fails. - """ - energies = self.spectrum.energies - counts = self.spectrum.counts - uncertainties = ( - self.spectrum.errors - if self.spectrum.errors is not None - else np.sqrt(np.maximum(counts, 0)) - ) - - # Narrow search around pulser channel guess if provided - if self.pulser_channel_guess is not None: - mask = np.abs(energies - self.pulser_channel_guess) < 5.0 - search_energies = energies[mask] - search_counts = counts[mask] - search_uncertainties = uncertainties[mask] - else: - search_energies = energies - search_counts = counts - search_uncertainties = uncertainties - - # Find the highest peak (pulser should be prominent) - peak_idx = np.argmax(search_counts) - peak_center = search_energies[peak_idx] - - # Fit Gaussian to pulser peak - fitter = GaussianFitter(search_energies, search_counts, search_uncertainties) - result = fitter.fit( - guess=(search_counts[peak_idx], peak_center, 1.0, np.mean(search_counts)) - ) - - if not result.success: - raise ValueError(f"Pulser peak fit failed: {result}") - - # Compute observed rate from peak amplitude - # Peak amplitude is in counts; divide by real_time to get rate - observed_rate = ( - result.amplitude / self.real_time_sec if self.real_time_sec > 0 else 0.0 - ) - - # Compute live-time fraction - if self.pulser_frequency_hz <= 0: - raise ValueError("Pulser frequency must be positive") - live_time_fraction = observed_rate / self.pulser_frequency_hz - - # Clamp to valid range [0, 1] - live_time_fraction = max(0.0, min(1.0, live_time_fraction)) - - return LivetimeResult( - live_time_fraction=live_time_fraction, - real_time_sec=self.real_time_sec, - pulser_frequency_hz=self.pulser_frequency_hz, - observed_pulser_rate=observed_rate, - pulser_peak_center_keV=result.center_keV, - pulser_peak_sigma_keV=result.sigma_keV, - pulser_peak_chi2=result.chi2, - metadata={ - "pulser_peak_channel": int(peak_center), - "pulser_fit_success": result.success, - }, - ) - - @staticmethod - def from_file( - filepath: str | Path, - pulser_frequency_hz: float, - real_time_sec: float, - n_channels: int = 4096, - run_id: Optional[str] = None, - source: Optional[str] = None, - ) -> LivetimeResult: - """Convenience: read spectrum from file and determine live-time. - - Parameters - ---------- - filepath : str or Path - Path to the pulser calibration binary file. - pulser_frequency_hz : float - Known pulser injection frequency in Hz. - real_time_sec : float - Total real acquisition time in seconds. - n_channels : int - Number of channels. - run_id : str, optional - Run identifier. - source : str, optional - Source identifier (e.g., 'Am241', 'pulser'). - - Returns - ------- - LivetimeResult - """ - spectrum = read_raw_spectrum( - filepath, - n_channels=n_channels, - run_id=run_id, - source=source, - live_time_sec=real_time_sec, - metadata={"operation": "livetime_calibration"}, - ) - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=pulser_frequency_hz, - real_time_sec=real_time_sec, - ) - return determiner.run() diff --git a/exp_data/pipeline.py b/exp_data/pipeline.py deleted file mode 100644 index a5f90a9..0000000 --- a/exp_data/pipeline.py +++ /dev/null @@ -1,307 +0,0 @@ -""" -Experimental data processing pipeline. - -Orchestrates the workflow: read -> subtract -> calibrate -> livetime -> validate -> export. - -Usage: - from exp_data.pipeline import process_spectrum, process_run_sequence - - # Single spectrum - result = process_spectrum( - raw_path="data/raw/A1.DAT", - calibration_path="data/raw/CAM1.DAT", - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) - - # Run sequence (cumulative subtraction) - results = process_run_sequence( - run_dir="data/raw/", - prefix="A", - calibration_path="data/raw/CAM1.DAT", - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Optional - -import numpy as np - -from exp_data.spectrum import ExpSpectrum -from exp_data.raw_reader import read_raw_spectrum, subtract_spectra, get_run_files -from exp_data.calibration import EnergyCalibrator, CalibrationResult -from exp_data.livetime import LivetimeDetermination, LivetimeResult - - -@dataclass -class PipelineResult: - """Result of a pipeline execution.""" - - spectrum: ExpSpectrum - """Calibrated, corrected ExpSpectrum.""" - - calibration: CalibrationResult - """Energy calibration result.""" - - livetime: Optional[LivetimeResult] - """Live-time determination result (None if not computed).""" - - validation_passed: bool - """Whether all quality checks passed.""" - - validation_messages: list[str] = field(default_factory=list) - """Validation messages (warnings/errors).""" - - -class SpectrumPipeline: - """Orchestrate the full experimental data processing pipeline. - - Workflow: - 1. Read raw binary spectrum - 2. (Optional) Subtract from previous cumulative spectrum - 3. Apply energy calibration from Am-241 - 4. Determine live-time from pulser (if available) - 5. Validate quality - 6. Export (optional) - """ - - def __init__( - self, - calibration_path: str | Path, - pulser_frequency_hz: Optional[float] = None, - real_time_sec: Optional[float] = None, - ): - self.calibration_path = Path(calibration_path) - self.pulser_frequency_hz = pulser_frequency_hz - self.real_time_sec = real_time_sec - - def process( - self, - raw_path: str | Path, - cumulative_path: Optional[str | Path] = None, - run_id: Optional[str] = None, - source: Optional[str] = None, - ) -> PipelineResult: - """Process a single spectrum through the full pipeline. - - Parameters - ---------- - raw_path : str or Path - Path to the raw binary spectrum file. - cumulative_path : str or Path, optional - Path to the previous cumulative spectrum for subtraction. - run_id : str, optional - Run identifier. - source : str, optional - Source identifier. - - Returns - ------- - PipelineResult - """ - # Step 1: Read raw spectrum - spectrum = read_raw_spectrum( - raw_path, - run_id=run_id, - source=source, - metadata={"operation": "read"}, - ) - - # Step 2: Subtract from cumulative (if provided) - if cumulative_path is not None: - cumulative = read_raw_spectrum(cumulative_path) - spectrum = subtract_spectra(spectrum, cumulative, run_id=run_id) - - # Step 3: Apply energy calibration - calibration = EnergyCalibrator.from_am241_spectrum(self.calibration_path) - spectrum = spectrum.apply_calibration(calibration) - - # Step 4: Determine live-time (if pulser data available) - livetime = None - if self.pulser_frequency_hz is not None and self.real_time_sec is not None: - livetime = LivetimeDetermination.from_file( - self.calibration_path, - pulser_frequency_hz=self.pulser_frequency_hz, - real_time_sec=self.real_time_sec, - run_id=f"{run_id}_livetime" if run_id else "livetime", - source="pulser", - ) - # Store live-time fraction in metadata - spectrum.metadata["live_time"] = { - "fraction": livetime.live_time_fraction, - "real_time_sec": livetime.real_time_sec, - "pulser_frequency_hz": livetime.pulser_frequency_hz, - "observed_pulser_rate": livetime.observed_pulser_rate, - "method": "pulser", - } - spectrum.live_time = livetime.real_time_sec * livetime.live_time_fraction - - # Step 5: Validate quality - messages = self._validate(spectrum, calibration, livetime) - passed = len([m for m in messages if m.startswith("ERROR:")]) == 0 - - return PipelineResult( - spectrum=spectrum, - calibration=calibration, - livetime=livetime, - validation_passed=passed, - validation_messages=messages, - ) - - def process_sequence( - self, - run_dir: str | Path, - prefix: str = "A", - first_file: Optional[str | Path] = None, - run_id: Optional[str] = None, - source: Optional[str] = None, - ) -> list[PipelineResult]: - """Process a sequence of cumulative binary files. - - Reads all files matching prefix in run_dir, subtracts consecutive - pairs to extract per-run spectra. - - Parameters - ---------- - run_dir : str or Path - Directory containing run files. - prefix : str - Filename prefix (default 'A'). - first_file : str or Path, optional - Path to the first cumulative file (baseline for subtraction). - If None, the first file in the sorted list is used as baseline. - run_id : str, optional - Base run identifier (appended with run number). - source : str, optional - Source identifier. - - Returns - ------- - list of PipelineResult - One result per extracted run. - """ - files = get_run_files(run_dir, prefix=prefix) - - if len(files) < 2: - raise ValueError( - f"Need at least 2 cumulative files for subtraction, " - f"found {len(files)} in {run_dir}" - ) - - results = [] - for i, filepath in enumerate(files[1:], start=1): - result = self.process( - raw_path=filepath, - cumulative_path=files[0], - run_id=f"{run_id}_{i}" if run_id else None, - source=source, - ) - results.append(result) - - return results - - def _validate( - self, - spectrum: ExpSpectrum, - calibration: CalibrationResult, - livetime: Optional[LivetimeResult], - ) -> list[str]: - """Validate spectrum quality. - - Checks: - - Calibration fit quality (chi2/dof < 3) - - Minimum total counts (> 100) - - Live-time fraction in valid range (if available) - - Energy range sanity (energies increasing, positive) - """ - messages = [] - - # Calibration quality - if calibration.chi2_per_dof > 3.0: - messages.append( - f"ERROR: Calibration fit quality poor (chi2/dof = {calibration.chi2_per_dof:.2f} > 3.0)" - ) - elif calibration.chi2_per_dof > 1.5: - messages.append( - f"WARNING: Calibration fit quality marginal (chi2/dof = {calibration.chi2_per_dof:.2f})" - ) - - # Minimum counts - total = spectrum.total_counts - if total < 100: - messages.append( - f"ERROR: Insufficient statistics (total counts = {total:.0f} < 100)" - ) - - # Live-time fraction - if livetime is not None: - if livetime.live_time_fraction < 0.5: - messages.append( - f"WARNING: Low live-time fraction (tau = {livetime.live_time_fraction:.3f} < 0.5)" - ) - elif livetime.live_time_fraction <= 0.0: - messages.append( - f"ERROR: Invalid live-time fraction (tau = {livetime.live_time_fraction:.3f} <= 0)" - ) - - # Energy range sanity - if spectrum.energies[0] < 0: - messages.append("ERROR: Negative minimum energy") - if not np.all(np.diff(spectrum.energies) >= 0): - messages.append("ERROR: Energies not monotonically increasing") - - return messages - - -# Convenience functions - - -def process_spectrum( - raw_path: str | Path, - calibration_path: str | Path, - run_id: Optional[str] = None, - source: Optional[str] = None, - cumulative_path: Optional[str | Path] = None, - pulser_frequency_hz: Optional[float] = None, - real_time_sec: Optional[float] = None, -) -> PipelineResult: - """Convenience: process a single spectrum through the full pipeline. - - See SpectrumPipeline.process() for details. - """ - pipeline = SpectrumPipeline( - calibration_path=calibration_path, - pulser_frequency_hz=pulser_frequency_hz, - real_time_sec=real_time_sec, - ) - return pipeline.process( - raw_path, cumulative_path=cumulative_path, run_id=run_id, source=source - ) - - -def process_run_sequence( - run_dir: str | Path, - calibration_path: str | Path, - prefix: str = "A", - run_id: Optional[str] = None, - source: Optional[str] = None, - pulser_frequency_hz: Optional[float] = None, - real_time_sec: Optional[float] = None, -) -> list[PipelineResult]: - """Convenience: process a run sequence through the full pipeline. - - See SpectrumPipeline.process_sequence() for details. - """ - pipeline = SpectrumPipeline( - calibration_path=calibration_path, - pulser_frequency_hz=pulser_frequency_hz, - real_time_sec=real_time_sec, - ) - return pipeline.process_sequence( - run_dir, prefix=prefix, run_id=run_id, source=source - ) diff --git a/exp_data/raw_reader.py b/exp_data/raw_reader.py deleted file mode 100644 index 9acd249..0000000 --- a/exp_data/raw_reader.py +++ /dev/null @@ -1,283 +0,0 @@ -""" -Raw binary spectrum reader for 4096-channel MCA data. - -Reads binary files with 128-byte header + 4096 channels of int32 counts. -Used for both Tc-99 beta spectra and Am-241 calibration sources. -""" - -from __future__ import annotations - -import struct -from pathlib import Path -from typing import Optional - -import numpy as np -from scipy.signal import find_peaks - -from exp_data.spectrum import ExpSpectrum - -# Binary format constants -HEADER_SIZE = 128 # bytes -CHANNEL_SIZE = 4 # int32 -DEFAULT_N_CHANNELS = 4096 -# 16512 = 128 + 4096 * 4 -FILE_SIZE = HEADER_SIZE + DEFAULT_N_CHANNELS * CHANNEL_SIZE - - -def read_raw_spectrum( - filepath: str | Path, - n_channels: int = DEFAULT_N_CHANNELS, - energy_keV: Optional[np.ndarray] = None, - run_id: Optional[str] = None, - source: Optional[str] = None, - live_time_sec: Optional[float] = None, - metadata: Optional[dict] = None, -) -> ExpSpectrum: - """Read a raw binary spectrum file. - - Parameters - ---------- - filepath - Path to the binary spectrum file. - n_channels - Number of energy channels (default 4096). - energy_keV - Energy calibration array (n_channels,). If None, uses channel indices. - run_id - Run identifier string. - source - Source identifier (e.g. 'Tc99', 'Am241'). - live_time_sec - Live acquisition time in seconds. - metadata - Additional metadata dict. - - Returns - ------- - ExpSpectrum - Parsed spectrum with metadata. - """ - filepath = Path(filepath) - if not filepath.exists(): - raise FileNotFoundError(f"Spectrum file not found: {filepath}") - - file_size = filepath.stat().st_size - expected = HEADER_SIZE + n_channels * CHANNEL_SIZE - if file_size != expected: - raise ValueError( - f"File size {file_size} != expected {expected} " - f"(header={HEADER_SIZE} + {n_channels}ch * {CHANNEL_SIZE}B)" - ) - - with open(filepath, "rb") as f: - # Skip header - f.seek(HEADER_SIZE) - counts = np.fromfile(f, dtype=np.int32, count=n_channels) - - if energy_keV is None: - energy_keV = np.arange(n_channels, dtype=np.float64) - - return ExpSpectrum( - energies=energy_keV, - counts=counts.astype(np.float64), - metadata=metadata or {}, - run_id=run_id, - source=source, - live_time=live_time_sec, - ) - - -def subtract_spectra( - cumulative_a: ExpSpectrum, - cumulative_b: ExpSpectrum, - run_id: Optional[str] = None, -) -> ExpSpectrum: - """Subtract two cumulative spectra to get the incremental run. - - Both spectra must have the same number of channels. - Result = A - B (where A is the later cumulative measurement). - - Parameters - ---------- - cumulative_a - Later cumulative spectrum (larger time). - cumulative_b - Earlier cumulative spectrum (smaller time). - run_id - Run ID for the difference spectrum. - - Returns - ------- - ExpSpectrum - The difference spectrum (A - B). - """ - assert len(cumulative_a.counts) == len(cumulative_b.counts), ( - f"Spectra must have same channel count: " - f"{len(cumulative_a.counts)} vs {len(cumulative_b.counts)}" - ) - - diff_counts = cumulative_a.counts - cumulative_b.counts - # Propagate errors: σ_diff = sqrt(σ_a² + σ_b²) - diff_errors = np.sqrt(cumulative_a.errors**2 + cumulative_b.errors**2) - - return ExpSpectrum( - energies=cumulative_a.energies.copy(), - counts=diff_counts, - errors=diff_errors, - metadata={ - **cumulative_a.metadata, - "operation": "subtract", - "parent_a": cumulative_a.run_id, - "parent_b": cumulative_b.run_id, - }, - run_id=run_id, - source=cumulative_a.source, - ) - - -def calibrate_from_am241( - filepath: str | Path, - known_peaks: dict[float, int] | None = None, - n_channels: int = DEFAULT_N_CHANNELS, -) -> tuple[float, float]: - """Calibrate energy scale using Am-241 calibration source. - - Matches known Am-241 gamma lines to detected peaks in the - low-energy region (ch < 300) and fits a linear calibration. - - Parameters - ---------- - filepath - Path to the Am-241 binary file (CAM1.DAT or CAM2.DAT). - known_peaks - Mapping of energy_keV -> expected channel. If None, auto-detects - Am-241 peaks at 13.9, 17.8, 26.3 keV. - n_channels - Number of channels. - - Returns - ------- - (offset_keV, gain_keV_per_ch) - Linear calibration: E = offset + gain * channel - """ - filepath = Path(filepath) - with open(filepath, "rb") as f: - f.read(128) # skip header - counts = np.array( - struct.unpack(f"{n_channels}i", f.read(n_channels * 4)), dtype=np.int64 - ) - - # Search for peaks in low-energy region only - counts_low = counts.copy() - counts_low[300:] = 0 - peaks, props = find_peaks(counts_low, height=10, distance=10) - peak_energies = np.array(peaks, dtype=float) - _peak_heights = props["peak_heights"] # noqa: F841 - - # Known Am-241 gamma lines (keV) - am241_lines = {13.9: None, 17.8: None, 26.3: None} - if known_peaks: - am241_lines = known_peaks - else: - # Match lines to peaks by proximity - for E in am241_lines: - best = peak_energies[np.argmin(np.abs(peak_energies - E / 0.1))] - am241_lines[E] = int(round(best)) - - # Use all matched peaks for calibration - energies = np.array(list(am241_lines.keys())) - channels = np.array(list(am241_lines.values())) - - # Linear fit - coeffs = np.polyfit(channels, energies, 1) - gain = coeffs[0] - offset = coeffs[1] - - return float(offset), float(gain) - - -def read_am241_calibration( - filepath: str | Path, - channel_energy_map: dict[int, float] | None = None, - n_channels: int = DEFAULT_N_CHANNELS, -) -> ExpSpectrum: - """Read Am-241 calibration spectrum. - - Parameters - ---------- - filepath - Path to the binary calibration file (CAM1.DAT or CAM2.DAT). - channel_energy_map - Mapping of channel index -> energy_keV for known peaks. - If None, assumes linear calibration from channel indices. - n_channels - Number of channels. - - Returns - ------- - ExpSpectrum - Calibration spectrum. - """ - filepath = Path(filepath) - - # Am-241 known peaks (keV): 13.9, 17.8, 26.3, 59.54 - # The spectrum typically has a few prominent peaks - with open(filepath, "rb") as f: - _header = f.read(128) - counts_raw = struct.unpack(f"{n_channels}i", f.read(n_channels * 4)) - - counts = np.array(counts_raw, dtype=np.float64) - energy_keV = np.arange(n_channels, dtype=np.float64) - - metadata = { - "source": "Am241", - "calibration_type": "energy", - "known_peaks_keV": [13.9, 17.8, 26.3, 59.54], - } - - return ExpSpectrum( - energies=energy_keV, - counts=counts, - metadata=metadata, - source="Am241", - ) - - -def get_run_files( - directory: str | Path, - prefix: str = "A", - sort_key: str = "natural", -) -> list[Path]: - """Get sorted list of run files in directory. - - Parameters - ---------- - directory - Directory containing run files. - prefix - Filename prefix (default 'A'). - sort_key - Sort method: 'natural' (A1, A2, ..., A10, A11) or 'lexicographic' (A1, A10, A100). - - Returns - ------- - list[Path] - Sorted file paths. - """ - directory = Path(directory) - runs = [f for f in directory.iterdir() if f.is_file() and f.name.startswith(prefix)] - - if sort_key == "natural": - runs.sort(key=lambda p: _natural_sort_key(p.name)) - else: - runs.sort(key=lambda p: p.name) - - return runs - - -def _natural_sort_key(s: str) -> list: - """Sort key for natural ordering: A1, A2, ..., A10, A100.""" - import re - - parts = re.split(r"(\d+)", s) - return [int(p) if p.isdigit() else p.lower() for p in parts] diff --git a/exp_data/root_io.py b/exp_data/root_io.py deleted file mode 100644 index 67c1122..0000000 --- a/exp_data/root_io.py +++ /dev/null @@ -1,287 +0,0 @@ -""" -ROOT file I/O for experimental spectra. - -Supports PyROOT (primary) and uproot (fallback) for reading and writing -ExpSpectrum objects to ROOT TTree/TBranch format. - -Usage: - from exp_data.root_io import write_spectrum, read_spectrum - - # Write - write_spectrum("output.root", spectrum) - - # Read - spectrum = read_spectrum("output.root") -""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import TYPE_CHECKING - -import numpy as np - -if TYPE_CHECKING: - from exp_data.spectrum import ExpSpectrum - - -# Try PyROOT first, fall back to uproot -_PYROOT_AVAILABLE = False -_UPTOOL_AVAILABLE = False - -try: - import ROOT # noqa: F401 - - _PYROOT_AVAILABLE = True -except ImportError: - try: - import uproot # type: ignore[import-not-found] # noqa: F401 - - _UPTOOL_AVAILABLE = True - except ImportError: - raise ImportError( - "Neither PyROOT nor uproot is available. " - "Install one of: pip install beta-spectrum[root] (PyROOT) " - "or pip install uproot (pure Python fallback)." - ) - - -def write_spectrum( - filepath: str | Path, - spectrum: "ExpSpectrum", - tree_name: str = "spectrum", -) -> None: - """Write an ExpSpectrum to a ROOT file. - - Creates a TTree with TBranches for energies, counts, errors, - scalar fields (dead_time, live_time, source, run_id, date), - and metadata as a JSON string. - - Parameters - ---------- - filepath : str or Path - Output ROOT file path. - spectrum : ExpSpectrum - Spectrum to write. - tree_name : str - Name of the TTree in the ROOT file. - """ - filepath = Path(filepath) - - if _PYROOT_AVAILABLE: - _write_pyroot(filepath, spectrum, tree_name) - elif _UPTOOL_AVAILABLE: - _write_uproot(filepath, spectrum, tree_name) - else: - raise RuntimeError("No ROOT I/O library available") - - -def read_spectrum( - filepath: str | Path, - tree_name: str = "spectrum", -) -> "ExpSpectrum": - """Read an ExpSpectrum from a ROOT file. - - Parameters - ---------- - filepath : str or Path - Input ROOT file path. - tree_name : str - Name of the TTree in the ROOT file. - - Returns - ------- - ExpSpectrum - Reconstructed spectrum with all metadata. - """ - filepath = Path(filepath) - - if _PYROOT_AVAILABLE: - return _read_pyroot(filepath, tree_name) - elif _UPTOOL_AVAILABLE: - return _read_uproot(filepath, tree_name) - else: - raise RuntimeError("No ROOT I/O library available") - - -# --- PyROOT implementation --- - - -def _write_pyroot( - filepath: Path, - spectrum: "ExpSpectrum", - tree_name: str, -) -> None: - """Write ExpSpectrum using PyROOT TTree/TBranch.""" - import ROOT - from array import array - - with ROOT.TFile(str(filepath), "RECREATE"): - tree = ROOT.TTree(tree_name, "Experimental beta spectrum") - - # Data arrays - energies_arr = np.array(spectrum.energies, dtype=np.float64) - counts_arr = np.array(spectrum.counts, dtype=np.float64) - errors_arr = np.array(spectrum.errors, dtype=np.float64) - - # TBranches for arrays - n = len(energies_arr) - tree.Branch("energies", energies_arr, f"energies[{n}]/D") - tree.Branch("counts", counts_arr, f"counts[{n}]/D") - tree.Branch("errors", errors_arr, f"errors[{n}]/D") - - # Scalar TBranches (using array of length 1) - dead_time_val = array("d", [spectrum.dead_time]) - live_time_val = array( - "d", [spectrum.live_time if spectrum.live_time is not None else 0.0] - ) - - tree.Branch("dead_time", dead_time_val, "dead_time/D") - tree.Branch("live_time", live_time_val, "live_time/D") - - # String fields - source_str = ROOT.TObjString(spectrum.source or "") - run_id_str = ROOT.TObjString(spectrum.run_id or "") - date_str = ROOT.TObjString(spectrum.date or "") - - tree.Branch("source", source_str) - tree.Branch("run_id", run_id_str) - tree.Branch("date", date_str) - - # Metadata as JSON - metadata_json = json.dumps(spectrum.metadata) - metadata_str = ROOT.TObjString(metadata_json) - tree.Branch("metadata_json", metadata_str) - - tree.Fill() - tree.Write() - - -def _read_pyroot( - filepath: Path, - tree_name: str, -) -> "ExpSpectrum": - """Read ExpSpectrum from ROOT file using PyROOT.""" - import ROOT - import json - - f = ROOT.TFile(str(filepath), "READ") - tree = f.Get(tree_name) - - # Read first (and only) entry - tree.GetEntry(0) - - energies = np.array(tree.energies, dtype=np.float64) - counts = np.array(tree.counts, dtype=np.float64) - errors = np.array(tree.errors, dtype=np.float64) - - dead_time = float(tree.dead_time) - live_time = float(tree.live_time) if tree.live_time > 0 else None - source_raw = tree.source - source = str(source_raw) if source_raw and str(source_raw) else None - run_id_raw = tree.run_id - run_id = str(run_id_raw) if run_id_raw and str(run_id_raw) else None - date_raw = tree.date - date = str(date_raw) if date_raw and str(date_raw) else None - - # Parse metadata JSON - metadata_raw = tree.metadata_json - # Handle various ROOT string types (TObjString, TString, str, bytes) - if hasattr(metadata_raw, "__str__"): - try: - metadata_str = str(metadata_raw) - except Exception: - metadata_str = "" - elif hasattr(metadata_raw, "GetString"): - metadata_str = str(metadata_raw.GetString()) - else: - metadata_str = "" - metadata = json.loads(metadata_str) if metadata_str else {} - - from exp_data.spectrum import ExpSpectrum - - return ExpSpectrum( - energies=energies, - counts=counts, - errors=errors, - metadata=metadata, - dead_time=dead_time, - live_time=live_time, - source=source, - run_id=run_id, - date=date, - ) - - -# --- uproot implementation --- - - -def _write_uproot( - filepath: Path, - spectrum: "ExpSpectrum", - tree_name: str, -) -> None: - """Write ExpSpectrum using uproot (pure Python).""" - import uproot - - with uproot.recreate(str(filepath)) as f: - f[tree_name] = { - "energies": np.array(spectrum.energies, dtype=">f8"), - "counts": np.array(spectrum.counts, dtype=">f8"), - "errors": np.array(spectrum.errors, dtype=">f8"), - "dead_time": np.array([spectrum.dead_time], dtype=">f8"), - "live_time": np.array( - [spectrum.live_time if spectrum.live_time is not None else 0.0], - dtype=">f8", - ), - "source": np.array([spectrum.source or ""], dtype="S100"), - "run_id": np.array([spectrum.run_id or ""], dtype="S100"), - "date": np.array([spectrum.date or ""], dtype="S100"), - "metadata_json": np.array([json.dumps(spectrum.metadata)], dtype="S500"), - } - - -def _read_uproot( - filepath: Path, - tree_name: str, -) -> "ExpSpectrum": - """Read ExpSpectrum from ROOT file using uproot.""" - import uproot - import json - - with uproot.open(str(filepath)) as f: - tree = f[tree_name] - data = tree.arrays(library="np") - - energies = np.array(data["energies"], dtype=np.float64) - counts = np.array(data["counts"], dtype=np.float64) - errors = np.array(data["errors"], dtype=np.float64) - - dead_time = float(data["dead_time"][0]) - live_time_raw = float(data["live_time"][0]) - live_time = live_time_raw if live_time_raw > 0 else None - - source_raw = data["source"][0] - source = source_raw.decode() if isinstance(source_raw, bytes) else source_raw - run_id_raw = data["run_id"][0] - run_id = run_id_raw.decode() if isinstance(run_id_raw, bytes) else run_id_raw - date_raw = data["date"][0] - date = date_raw.decode() if isinstance(date_raw, bytes) else date_raw - - metadata_str = data["metadata_json"][0] - if isinstance(metadata_str, bytes): - metadata_str = metadata_str.decode() - metadata = json.loads(metadata_str) if metadata_str else {} - - return ExpSpectrum( - energies=energies, - counts=counts, - errors=errors, - metadata=metadata, - dead_time=dead_time, - live_time=live_time, - source=source if source else None, - run_id=run_id if run_id else None, - date=date if date else None, - ) diff --git a/exp_data/spectrum.py b/exp_data/spectrum.py deleted file mode 100644 index 44f4557..0000000 --- a/exp_data/spectrum.py +++ /dev/null @@ -1,211 +0,0 @@ -""" -Experimental spectrum container. - -Stores measured detector spectra with metadata, uncertainties, -and calibration information. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Optional - -if TYPE_CHECKING: - from exp_data.calibration import CalibrationResult - -import numpy as np - - -@dataclass -class ExpSpectrum: - """Container for an experimental beta spectrum measurement. - - Parameters - ---------- - energies : np.ndarray - Energy calibration values for each channel (keV), shape (N_channels,). - counts : np.ndarray - Raw counts per channel, shape (N_channels,). - errors : np.ndarray, optional - Statistical uncertainties per channel. Defaults to sqrt(counts). - metadata : dict, optional - Arbitrary metadata (source, run ID, date, etc.). - dead_time : float, optional - Dead time fraction (0.0 to 1.0). - live_time : float, optional - Live acquisition time in seconds. - source : str, optional - Radioactive source identifier. - run_id : str, optional - Unique run identifier. - date : str, optional - Measurement date (ISO format). - """ - - energies: np.ndarray - counts: np.ndarray - errors: Optional[np.ndarray] = None - metadata: dict = field(default_factory=dict) - dead_time: float = 0.0 - live_time: Optional[float] = None - source: Optional[str] = None - run_id: Optional[str] = None - date: Optional[str] = None - - def __post_init__(self) -> None: - """Validate and compute default errors.""" - self.energies = np.asarray(self.energies, dtype=np.float64) - self.counts = np.asarray(self.counts, dtype=np.float64) - assert self.energies.shape == self.counts.shape, ( - f"Energies ({self.energies.shape}) and counts ({self.counts.shape}) " - "must have the same shape" - ) - assert np.all(self.energies >= 0), "Energies must be non-negative" - - # Difference spectra (cumulative subtraction) can have negative counts - # in channels where the true count is zero and Poisson noise causes - # the earlier cumulative to slightly exceed the later one. - # This is normal and expected — errors propagate correctly. - - if self.errors is None: - # Poisson statistics: σ = √N for positive counts - pos_mask = self.counts >= 0 - errors = np.zeros_like(self.counts) - errors[pos_mask] = np.sqrt(self.counts[pos_mask]) - neg_mask = ~pos_mask - if neg_mask.any(): - errors[neg_mask] = np.sqrt(np.abs(self.counts[neg_mask])) - self.errors = errors - else: - self.errors = np.asarray(self.errors, dtype=np.float64) - assert self.errors.shape == self.counts.shape - - @property - def n_channels(self) -> int: - """Number of energy channels.""" - return len(self.energies) - - @property - def energy_range(self) -> tuple[float, float]: - """(min_energy_keV, max_energy_keV).""" - return (float(self.energies[0]), float(self.energies[-1])) - - @property - def total_counts(self) -> float: - """Total counts in the spectrum.""" - return float(np.sum(self.counts)) - - @property - def live_time_corrected_rate(self) -> Optional[float]: - """Count rate corrected for dead time, in counts per second.""" - if self.live_time is None or self.live_time <= 0: - return None - live_seconds = self.live_time * (1.0 - self.dead_time) - return self.total_counts / live_seconds - - def normalize(self) -> np.ndarray: - """Return spectrum normalized to unit area.""" - integral = np.trapezoid(self.counts, self.energies) - if integral > 0: - return self.counts / integral - return self.counts.copy() - - def apply_dead_time_correction(self) -> ExpSpectrum: - """Return a new spectrum corrected for dead time. - - Correction: N_true = N_measured / (1 - τ), where τ is dead time fraction. - """ - if self.dead_time <= 0 or self.dead_time >= 1.0: - return ExpSpectrum( - energies=self.energies.copy(), - counts=self.counts.copy(), - errors=self.errors.copy(), - metadata=self.metadata.copy(), - dead_time=self.dead_time, - live_time=self.live_time, - source=self.source, - run_id=self.run_id, - date=self.date, - ) - - corrected_counts = self.counts / (1.0 - self.dead_time) - corrected_errors = self.errors / (1.0 - self.dead_time) - - return ExpSpectrum( - energies=self.energies.copy(), - counts=corrected_counts, - errors=corrected_errors, - metadata=self.metadata.copy(), - dead_time=0.0, # Corrected - live_time=self.live_time, - source=self.source, - run_id=self.run_id, - date=self.date, - ) - - def apply_calibration(self, calibration: "CalibrationResult") -> ExpSpectrum: - """Return a new spectrum with energies recalibrated via the given CalibrationResult. - - The calibration polynomial is applied to the current energy axis: - calibrated_energies = calibration(channel_indices). - - Parameters - ---------- - calibration : CalibrationResult - A CalibrationResult from EnergyCalibrator (e.g., from Am-241 peak detection). - - Returns - ------- - ExpSpectrum - New ExpSpectrum with recalibrated energies, all other fields preserved. - """ - - calibrated_energies = calibration(self.energies) - return ExpSpectrum( - energies=calibrated_energies, - counts=self.counts.copy(), - errors=self.errors.copy(), - metadata={ - **self.metadata, - "calibration": { - "type": "linear" if calibration.order == 1 else "quadratic", - "coefficients": calibration.coefficients.tolist(), - "chi2_per_dof": calibration.chi2_per_dof, - "n_points": calibration.n_points, - }, - }, - dead_time=self.dead_time, - live_time=self.live_time, - source=self.source, - run_id=self.run_id, - date=self.date, - ) - - def to_dict(self) -> dict: - """Serialize to dictionary.""" - return { - "energies": self.energies.tolist(), - "counts": self.counts.tolist(), - "errors": self.errors.tolist(), - "metadata": self.metadata, - "dead_time": self.dead_time, - "live_time": self.live_time, - "source": self.source, - "run_id": self.run_id, - "date": self.date, - } - - @classmethod - def from_dict(cls, data: dict) -> ExpSpectrum: - """Deserialize from dictionary.""" - return cls( - energies=np.array(data["energies"], dtype=np.float64), - counts=np.array(data["counts"], dtype=np.float64), - errors=np.array(data["errors"], dtype=np.float64), - metadata=data.get("metadata", {}), - dead_time=data.get("dead_time", 0.0), - live_time=data.get("live_time"), - source=data.get("source"), - run_id=data.get("run_id"), - date=data.get("date"), - ) diff --git a/notebooks/exp_calibration.ipynb b/notebooks/exp_calibration.ipynb deleted file mode 100644 index a87c65a..0000000 --- a/notebooks/exp_calibration.ipynb +++ /dev/null @@ -1,321 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Experimental Data Calibration Pipeline\n", - "\n", - "Interactive workflow for processing raw CAMAC MCA binary spectra:\n", - "read -> subtract -> calibrate -> livetime -> validate -> export\n", - "\n", - "## Prerequisites\n", - "- Raw binary spectrum files (A*.DAT format)\n", - "- Am-241 calibration file (CAM1.DAT)\n", - "- Pulser frequency and real time (from calibration run metadata)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Use non-interactive backend only when no display is available (e.g. jupyter execute)\n", - "import matplotlib as mpl\n", - "try:\n", - " import matplotlib.pyplot as plt\n", - " # Check if we're in a headless environment\n", - " import os\n", - " if not os.environ.get(\"DISPLAY\") and not os.environ.get(\"WAYLAND_DISPLAY\"):\n", - " mpl.use(\"Agg\")\n", - "except ImportError:\n", - " mpl.use(\"Agg\")\n", - " import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "\n", - "from exp_data.pipeline import process_spectrum, process_run_sequence\n", - "from exp_data.root_io import write_spectrum, read_spectrum\n", - "from exp_data.livetime import LivetimeDetermination\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Configuration" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Paths\n", - "RAW_FILE = \"../data/raw/A1\"\n", - "CALIBRATION_FILE = \"../data/raw/CAM1.DAT\"\n", - "RUN_DIR = \"../data/raw/\"\n", - "\n", - "# Pulser parameters (from calibration run metadata)\n", - "PULSER_FREQUENCY_HZ = 100.0\n", - "REAL_TIME_SEC = 3600.0" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Raw Spectrum Visualization\n", - "Visualize raw spectra before calibration to identify peaks and verify calibration source." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Visualize raw spectra before calibration\n", - "# This helps identify what calibration source is actually present\n", - "\n", - "# Detect header size from file\n", - "def detect_header_size(filepath):\n", - " \"\"\"Detect MCA header size from the file size.\"\"\"\n", - " # Simple heuristic: if file is exactly 32 + 4096*4 = 16412, header=32\n", - " # If file is 128 + 4096*4 = 16512, header=128\n", - " import os\n", - " size = os.path.getsize(filepath)\n", - " expected_32 = 32 + 4096 * 4\n", - " expected_128 = 128 + 4096 * 4\n", - " if size == expected_32:\n", - " return 32\n", - " elif size == expected_128:\n", - " return 128\n", - " else:\n", - " return 32 # default\n", - "\n", - "# --- Raw run spectrum ---\n", - "header_size = detect_header_size(RAW_FILE)\n", - "with open(RAW_FILE, 'rb') as f:\n", - " f.read(header_size)\n", - " raw_counts = np.fromfile(f, dtype=np.int32, count=4096)\n", - "\n", - "fig, axes = plt.subplots(2, 2, figsize=(14, 8))\n", - "\n", - "ax = axes[0, 0]\n", - "ax.plot(raw_counts, 'b-', linewidth=0.5)\n", - "ax.set_title('Raw Run Spectrum (counts vs channel)')\n", - "ax.set_xlabel('Channel')\n", - "ax.set_ylabel('Counts')\n", - "ax.set_yscale('log')\n", - "ax.set_xlim(0, 600)\n", - "\n", - "ax = axes[0, 1]\n", - "ax.plot(raw_counts, 'b-', linewidth=0.5)\n", - "ax.set_title('Raw Run Spectrum (full range)')\n", - "ax.set_xlabel('Channel')\n", - "ax.set_ylabel('Counts')\n", - "ax.set_yscale('log')\n", - "\n", - "# --- Calibration spectrum ---\n", - "with open(CALIBRATION_FILE, 'rb') as f:\n", - " f.read(128) # CAM files have 128-byte header\n", - " cal_counts = np.fromfile(f, dtype=np.int32, count=4096)\n", - "\n", - "ax = axes[1, 0]\n", - "ax.plot(cal_counts, 'r-', linewidth=0.5)\n", - "ax.set_title('Calibration Spectrum (counts vs channel)')\n", - "ax.set_xlabel('Channel')\n", - "ax.set_ylabel('Counts')\n", - "ax.set_xlim(0, 600)\n", - "ax.set_yscale('log')\n", - "\n", - "ax = axes[1, 1]\n", - "ax.plot(cal_counts, 'r-', linewidth=0.5)\n", - "ax.set_title('Calibration Spectrum (full range)')\n", - "ax.set_xlabel('Channel')\n", - "ax.set_ylabel('Counts')\n", - "ax.set_yscale('log')\n", - "\n", - "plt.tight_layout()\n", - "plt.show()\n", - "\n", - "# Print some stats\n", - "print(f'Raw run: {len(raw_counts)} channels, max={raw_counts.max()} at ch {np.argmax(raw_counts)}')\n", - "print(f'Calibration: {len(cal_counts)} channels, max={cal_counts.max()} at ch {np.argmax(cal_counts)}')\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Process Single Spectrum" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "result = process_spectrum(\n", - " raw_path=RAW_FILE,\n", - " calibration_path=CALIBRATION_FILE,\n", - " pulser_frequency_hz=PULSER_FREQUENCY_HZ,\n", - " real_time_sec=REAL_TIME_SEC,\n", - ")\n", - "\n", - "print(f\"Validation: {'PASSED' if result.validation_passed else 'FAILED'}\")\n", - "print(f\"Calibration chi2/dof: {result.calibration.chi2_per_dof:.4f}\")\n", - "for msg in result.validation_messages:\n", - " print(f\" {msg}\")\n", - "if result.livetime:\n", - " print(f\"Live-time fraction: {result.livetime.live_time_fraction:.4f}\")\n", - " print(f\"Pulser peak: {result.livetime.pulser_peak_center_keV:.1f} keV\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Plot Calibrated Spectrum" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "plt.figure(figsize=(10, 6))\n", - "plt.plot(result.spectrum.energies, result.spectrum.counts, 'b-', label='Counts')\n", - "plt.xlabel('Energy (keV)')\n", - "plt.ylabel('Counts')\n", - "plt.yscale('log')\n", - "plt.title('Calibrated Beta Spectrum')\n", - "plt.legend()\n", - "plt.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Process Run Sequence" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "results = process_run_sequence(\n", - " run_dir=RUN_DIR,\n", - " calibration_path=CALIBRATION_FILE,\n", - " prefix=\"A\",\n", - " pulser_frequency_hz=PULSER_FREQUENCY_HZ,\n", - " real_time_sec=REAL_TIME_SEC,\n", - ")\n", - "\n", - "print(f\"Processed {len(results)} runs\")\n", - "for i, r in enumerate(results):\n", - " print(f\" Run {i+1}: {'PASSED' if r.validation_passed else 'FAILED'} \"\n", - " f\"(chi2/dof={r.calibration.chi2_per_dof:.4f})\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Export to ROOT" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "write_spectrum(\"output.root\", result.spectrum)\n", - "print(\"Exported to output.root\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Read Back from ROOT" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "read_back = read_spectrum(\"output.root\")\n", - "print(f\"Read {len(read_back.energies)} channels\")\n", - "print(f\"Source: {read_back.source}\")\n", - "print(f\"Run ID: {read_back.run_id}\")\n", - "print(f\"Calibration: {read_back.metadata.get('calibration', {})}\")\n", - "\n", - "# Verify round-trip\n", - "assert np.allclose(read_back.energies, result.spectrum.energies)\n", - "assert np.allclose(read_back.counts, result.spectrum.counts)\n", - "print(\"Round-trip verified\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Livetime Determination" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Direct livetime determination from calibration file\n", - "livetime_result = LivetimeDetermination.from_file(\n", - " filepath=CALIBRATION_FILE,\n", - " pulser_frequency_hz=PULSER_FREQUENCY_HZ,\n", - " real_time_sec=REAL_TIME_SEC,\n", - " source=\"pulser\",\n", - ")\n", - "\n", - "print(f\"Live-time fraction: {livetime_result.live_time_fraction:.4f}\")\n", - "print(f\"Observed pulser rate: {livetime_result.observed_pulser_rate:.1f} Hz\")\n", - "print(f\"Pulser peak: {livetime_result.pulser_peak_center_keV:.1f} keV \"\n", - " f\"(sigma={livetime_result.pulser_peak_sigma_keV:.2f} keV, \"\n", - " f\"chi2/dof={livetime_result.pulser_peak_chi2:.2f})\")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.14.4" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/tests/conftest.py b/tests/conftest.py index 3e03fcf..3bbb184 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,9 +3,6 @@ import numpy as np import pytest -from exp_data.spectrum import ExpSpectrum -from exp_data.calibration import CalibrationResult - @pytest.fixture() def W_low(): @@ -38,47 +35,6 @@ def large_W(): return np.linspace(1.002, 4.9, 1000) -@pytest.fixture() -def synthetic_spectrum(): - """Create a synthetic ExpSpectrum (4096 channels).""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 1000 + np.random.poisson(5, n_channels) - return ExpSpectrum( - energies=energies, - counts=counts, - source="Tc99", - run_id="synthetic", - ) -@pytest.fixture() -def calibration_result(): - """Create a synthetic CalibrationResult (linear, E = 0.4 + 0.1*ch).""" - coeffs = np.array([0.1, 0.4]) # polyfit order: [slope, intercept] - cov = np.eye(2) * 1e-6 - return CalibrationResult( - coefficients=coeffs, - order=1, - chi2_per_dof=0.5, - n_points=2, - covariance=cov, - ) - -@pytest.fixture() -def full_spectrum(): - """Create a full ExpSpectrum with all fields populated.""" - return ExpSpectrum( - energies=np.arange(4096, dtype=np.float64), - counts=np.exp(-np.arange(4096) / 500) * 1000 + np.random.poisson(5, 4096), - source="Tc99", - run_id="run_042", - date="2025-01-15", - dead_time=0.05, - live_time=3420.0, - metadata={ - "calibration": {"type": "linear", "coefficients": [0.1, 0.4]}, - "livetime": {"fraction": 0.95, "method": "pulser"}, - }, - ) diff --git a/tests/exp_data/conftest.py b/tests/exp_data/conftest.py deleted file mode 100644 index 521f070..0000000 --- a/tests/exp_data/conftest.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Shared fixtures for exp_data tests.""" - -from __future__ import annotations - -import numpy as np -import pytest - -from exp_data.spectrum import ExpSpectrum -from exp_data.calibration import CalibrationResult - - -@pytest.fixture -def synthetic_spectrum(): - """Create a synthetic ExpSpectrum (4096 channels).""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 1000 + np.random.poisson(5, n_channels) - return ExpSpectrum( - energies=energies, - counts=counts, - source="Tc99", - run_id="synthetic", - ) - - -@pytest.fixture -def calibration_result(): - """Create a synthetic CalibrationResult (linear, E = 0.4 + 0.1*ch).""" - coeffs = np.array([0.1, 0.4]) # polyfit order: [slope, intercept] - cov = np.eye(2) * 1e-6 - return CalibrationResult( - coefficients=coeffs, - order=1, - chi2_per_dof=0.5, - n_points=2, - covariance=cov, - ) - - -@pytest.fixture -def full_spectrum(): - """Create a full ExpSpectrum with all fields populated.""" - return ExpSpectrum( - energies=np.arange(4096, dtype=np.float64), - counts=np.exp(-np.arange(4096) / 500) * 1000 + np.random.poisson(5, 4096), - source="Tc99", - run_id="run_042", - date="2025-01-15", - dead_time=0.05, - live_time=3420.0, - metadata={ - "calibration": {"type": "linear", "coefficients": [0.1, 0.4]}, - "livetime": {"fraction": 0.95, "method": "pulser"}, - }, - ) diff --git a/tests/exp_data/test_calibration.py b/tests/exp_data/test_calibration.py deleted file mode 100644 index 5cc4d87..0000000 --- a/tests/exp_data/test_calibration.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Tests for exp_data.calibration.""" - -import numpy as np -import pytest - -from exp_data.calibration import CalibrationResult, EnergyCalibrator - - -class TestCalibrationResult: - """Test CalibrationResult.""" - - def test_linear_call(self): - # polyfit returns [slope, intercept] = [0.5, 5.0] - # polyval([0.5, 5.0], ch) = 0.5*ch + 5.0 - result = CalibrationResult( - coefficients=np.array([0.5, 5.0]), # E = 0.5*ch + 5.0 - order=1, - chi2_per_dof=0.01, - n_points=2, - covariance=np.eye(2), - ) - np.testing.assert_allclose(result(0), 5.0) - np.testing.assert_allclose(result(10), 10.0) - - def test_quadratic_call(self): - # polyfit returns [a2, a1, a0] = [0.01, 0.1, 1.0] - # polyval([0.01, 0.1, 1.0], ch) = 0.01*ch² + 0.1*ch + 1.0 - result = CalibrationResult( - coefficients=np.array([0.01, 0.1, 1.0]), - order=2, - chi2_per_dof=0.01, - n_points=3, - covariance=np.eye(3), - ) - np.testing.assert_allclose(result(0), 1.0, rtol=1e-10) - # ch=10: 0.01*100 + 0.1*10 + 1.0 = 1 + 1 + 1 = 3 - np.testing.assert_allclose(result(10), 3.0) - - def test_invert(self): - result = CalibrationResult( - coefficients=np.array([0.5, 5.0]), # E = 0.5*ch + 5.0 - order=1, - chi2_per_dof=0.01, - n_points=2, - covariance=np.eye(2), - ) - inv = result.invert() - # E=10 => ch = (10-5)/0.5 = 10 - np.testing.assert_allclose(inv(10.0), 10.0, rtol=0.01) - - def test_resolution_at(self): - result = CalibrationResult( - coefficients=np.array([1.0, 0.0]), # E = 1.0*ch + 0 - order=1, - chi2_per_dof=0.01, - n_points=2, - covariance=np.eye(2), - ) - # dE/dch = 1.0, resolution = 1.0 * 1.0 = 1.0 keV - res = result.resolution_at(50.0, sigma_channel=1.0) - np.testing.assert_allclose(res, 1.0) - - def test_resolution_at_quadratic(self): - result = CalibrationResult( - coefficients=np.array([0.02, 1.0, 0.0]), # E = 0.02*ch² + 1.0*ch - order=2, - chi2_per_dof=0.01, - n_points=3, - covariance=np.eye(3), - ) - # dE/dch = 2*0.02*ch + 1.0 = 0.04*ch + 1.0 - # at ch=50: dE/dch = 0.04*50 + 1.0 = 3.0 - expected = 0.04 * 50.0 + 1.0 - res = result.resolution_at(50.0, sigma_channel=1.0) - np.testing.assert_allclose(res, expected) - - -class TestEnergyCalibrator: - """Test EnergyCalibrator.""" - - def test_not_calibrated_initially(self): - cal = EnergyCalibrator() - assert not cal.is_calibrated - - def test_get_linear_calibration(self): - cal = EnergyCalibrator() - result = cal.get_linear_calibration( - low_energy_keV=0.0, - high_energy_keV=100.0, - low_channel=0, - high_channel=100, - ) - assert cal.is_calibrated - assert result.order == 1 - np.testing.assert_allclose(result(50), 50.0) - - def test_apply_calibration_with_points(self): - cal = EnergyCalibrator() - cal.add_calibration_point(0, 0.0) - cal.add_calibration_point(50, 50.0) - cal.add_calibration_point(100, 100.0) - - result = cal.fit_calibration(order=1) - assert result.order == 1 - # Perfect fit on collinear points - np.testing.assert_allclose(result(50), 50.0, rtol=1e-6) - - def test_channel_to_energy_raises_without_calibration(self): - cal = EnergyCalibrator() - with pytest.raises(RuntimeError): - cal.channel_to_energy(50) - - def test_energy_to_channel_raises_without_calibration(self): - cal = EnergyCalibrator() - with pytest.raises(RuntimeError): - cal.energy_to_channel(50.0) - - def test_full_channel_mapping(self): - """Calibrate using full channel-to-energy mapping.""" - # Simulate: E = 0.5 + 0.1*ch + 0.001*ch² - n_channels = 200 - channels = np.arange(n_channels, dtype=np.float64) - energies = 0.5 + 0.1 * channels + 0.001 * channels**2 - - cal = EnergyCalibrator() - result = cal.fit_calibration(order=2, channel_energies=energies) - - assert result.order == 2 - np.testing.assert_allclose(result(channels), energies, rtol=1e-10) - - def test_quadratic_calibration(self): - cal = EnergyCalibrator() - cal.add_calibration_point(0, 0.0) - cal.add_calibration_point(50, 52.0) - cal.add_calibration_point(100, 108.0) # slight quadratic deviation - - result = cal.fit_calibration(order=2) - assert result.order == 2 diff --git a/tests/exp_data/test_corrections.py b/tests/exp_data/test_corrections.py deleted file mode 100644 index e38d0b2..0000000 --- a/tests/exp_data/test_corrections.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Tests for exp_data.corrections.""" - -import numpy as np - -from exp_data.corrections import ( - DeadTimeCorrection, - PileUpCorrection, - BackgroundSubtractor, -) - - -class TestDeadTimeCorrection: - """Test dead-time correction.""" - - def test_non_paralyzable_live_time(self): - dt = DeadTimeCorrection(dead_time=0.001, model="non_paralyzable") - # At 100 cps: live = 1 - 100*0.001 = 0.9 - assert np.isclose(dt.live_time_fraction(100.0), 0.9) - - def test_non_paralyzable_correct_rate(self): - dt = DeadTimeCorrection(dead_time=0.01, model="non_paralyzable") - # True rate = observed / (1 - obs*dt) = 90 / (1 - 90*0.01) = 90 / 0.1 = 900 - assert np.isclose(dt.correct_rate(90.0), 900.0, rtol=0.01) - - def test_paralyzable_live_time(self): - dt = DeadTimeCorrection(dead_time=0.001, model="paralyzable") - # At low rate: exp(-0.1) ≈ 0.905 - frac = dt.live_time_fraction(100.0) - assert np.isclose(frac, np.exp(-0.1), rtol=1e-5) - - def test_correct_counts(self): - dt = DeadTimeCorrection(dead_time=0.01, model="non_paralyzable") - counts = np.array([100.0, 200.0, 300.0]) - live_time = 10.0 - corrected = dt.correct_counts(counts, live_time) - # Total observed rate = 600/10 = 60 cps - # True rate = 60 / (1 - 60*0.01) = 60 / 0.94 = 63.83 cps - # Scale factor = 63.83/60 = 1.0638 - scale = 1.0 / (1 - 60.0 * 0.01) - expected = counts * scale - np.testing.assert_allclose(corrected, expected) - - -class TestPileUpCorrection: - """Test pile-up correction.""" - - def test_pile_up_fraction(self): - pu = PileUpCorrection(dead_time=0.001) - # At 100 cps: (0.1)² = 0.01 - frac = pu.estimate_pile_up_fraction(100.0) - assert np.isclose(frac, 0.01) - - def test_correct_spectrum(self): - pu = PileUpCorrection(dead_time=0.001) - counts = np.array([100.0, 200.0, 300.0]) - live_time = 10.0 - corrected = pu.correct_spectrum(counts, live_time) - # Pile-up correction adds a small fraction back - assert np.all(corrected >= counts) - - -class TestBackgroundSubtractor: - """Test background subtraction.""" - - def test_constant_background(self): - energies = np.linspace(0, 100, 200) - counts = 50 * np.exp(-0.5 * ((energies - 50) / 5) ** 2) + 10.0 # peak + bg - - corrected = BackgroundSubtractor.constant_background( - counts, - energies, - low_roi=(0.0, 10.0), - high_roi=(90.0, 100.0), - ) - - # Peak should still be visible, background removed - assert np.max(corrected) > 0 - # Background regions should be near zero - assert np.mean(corrected[energies <= 10]) < 5 - assert np.mean(corrected[energies >= 90]) < 5 - - def test_polynomial_background(self): - energies = np.linspace(0, 100, 200) - counts = 50 * np.exp(-0.5 * ((energies - 50) / 5) ** 2) + 10.0 + 0.05 * energies - - corrected = BackgroundSubtractor.polynomial_background( - counts, - energies, - order=1, - low_roi=(0.0, 10.0), - high_roi=(90.0, 100.0), - ) - - assert np.max(corrected) > 0 - - def test_roi_average(self): - counts = np.array([10.0, 12.0, 50.0, 10.0, 11.0]) - mask = np.array([True, True, False, True, True]) - avg = BackgroundSubtractor.roi_average(counts, mask) - np.testing.assert_allclose(avg, 10.75) - - def test_roi_average_empty(self): - counts = np.array([10.0, 20.0]) - mask = np.array([False, False]) - avg = BackgroundSubtractor.roi_average(counts, mask) - assert avg == 0.0 diff --git a/tests/exp_data/test_fitters.py b/tests/exp_data/test_fitters.py deleted file mode 100644 index e9bbee1..0000000 --- a/tests/exp_data/test_fitters.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Tests for exp_data.fitters.""" - -import numpy as np - -from exp_data.fitters import GaussianFitter, PeakFitter - - -class TestGaussianFitter: - """Test Gaussian peak fitting.""" - - def _make_synthetic_peak( - self, - center=50.0, - sigma=2.0, - amplitude=100.0, - background=10.0, - n_points=500, - energy_range=(0, 100), - ): - """Create synthetic Gaussian peak with noise.""" - energies = np.linspace(energy_range[0], energy_range[1], n_points) - counts = ( - amplitude * np.exp(-0.5 * ((energies - center) / sigma) ** 2) + background - ) - counts += np.random.RandomState(42).normal( - 0, np.sqrt(np.maximum(counts, 1)), n_points - ) - return energies, counts - - def test_fit_returns_valid_result(self): - energies, counts = self._make_synthetic_peak() - fitter = GaussianFitter(energies, counts) - result = fitter.fit() - - assert result.success - assert result.chi2 > 0 - assert result.n_dof > 0 - assert result.covariance.shape == (4, 4) - - def test_fit_recovers_parameters(self): - center, sigma, amplitude, background = 50.0, 2.0, 100.0, 10.0 - energies, counts = self._make_synthetic_peak( - center=center, sigma=sigma, amplitude=amplitude, background=background - ) - fitter = GaussianFitter(energies, counts) - result = fitter.fit() - - np.testing.assert_allclose(result.center_keV, center, rtol=0.05) - np.testing.assert_allclose(result.sigma_keV, sigma, rtol=0.1) - np.testing.assert_allclose(result.amplitude, amplitude, rtol=0.15) - - def test_resolution_property(self): - energies, counts = self._make_synthetic_peak(sigma=2.0) - fitter = GaussianFitter(energies, counts) - result = fitter.fit() - - expected_fwhm = 2.35482 * 2.0 - np.testing.assert_allclose(result.resolution(), expected_fwhm, rtol=0.05) - - def test_resolution_fraction(self): - energies, counts = self._make_synthetic_peak(center=50.0, sigma=2.0) - fitter = GaussianFitter(energies, counts) - result = fitter.fit() - - expected_frac = (2.35482 * 2.0) / 50.0 - np.testing.assert_allclose( - result.resolution_fraction(), expected_frac, rtol=0.05 - ) - - def test_auto_guess(self): - energies, counts = self._make_synthetic_peak() - fitter = GaussianFitter(energies, counts) - result = fitter.fit() - assert result.success - - -class TestPeakFitter: - """Test multi-peak detection and fitting.""" - - def _make_multi_peak(self): - """Create spectrum with two peaks.""" - energies = np.linspace(0, 100, 1000) - counts = np.zeros_like(energies) - counts += 50 * np.exp(-0.5 * ((energies - 20) / 2) ** 2) - counts += 80 * np.exp(-0.5 * ((energies - 70) / 3) ** 2) - counts += 5 # constant background - counts += np.random.RandomState(42).normal( - 0, np.sqrt(np.maximum(counts, 1)), len(energies) - ) - return energies, counts - - def test_find_peaks(self): - energies, counts = self._make_multi_peak() - fitter = PeakFitter(energies, counts) - peaks = fitter.find_peaks(min_height=10.0, min_distance_keV=10.0) - - assert len(peaks) >= 1 # At least one peak should be found - - def test_fit_peaks_with_known_energies(self): - energies, counts = self._make_multi_peak() - fitter = PeakFitter(energies, counts) - results = fitter.fit_peaks(known_energies=np.array([20.0, 70.0])) - - assert len(results) >= 1 - for r in results: - assert r.success - assert r.n_dof > 0 - - def test_fit_peaks_auto_detect(self): - energies, counts = self._make_multi_peak() - fitter = PeakFitter(energies, counts) - results = fitter.fit_peaks() - - assert len(results) >= 1 diff --git a/tests/exp_data/test_livetime.py b/tests/exp_data/test_livetime.py deleted file mode 100644 index 13c1cde..0000000 --- a/tests/exp_data/test_livetime.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Unit tests for the LivetimeDetermination class.""" - -from __future__ import annotations - -import numpy as np -import pytest - -from exp_data.spectrum import ExpSpectrum -from exp_data.livetime import LivetimeDetermination, LivetimeResult - - -class TestLivetimeResult: - def test_basic(self): - result = LivetimeResult( - live_time_fraction=0.95, - real_time_sec=3600.0, - pulser_frequency_hz=100.0, - observed_pulser_rate=95.0, - pulser_peak_center_keV=10.0, - pulser_peak_sigma_keV=0.5, - pulser_peak_chi2=1.2, - ) - assert result.live_time_fraction == 0.95 - assert result.metadata == {} - - def test_metadata_populated(self): - result = LivetimeResult( - live_time_fraction=0.95, - real_time_sec=3600.0, - pulser_frequency_hz=100.0, - observed_pulser_rate=95.0, - pulser_peak_center_keV=10.0, - pulser_peak_sigma_keV=0.5, - pulser_peak_chi2=1.2, - metadata={"pulser_peak_channel": 100, "pulser_fit_success": True}, - ) - assert result.metadata["pulser_peak_channel"] == 100 - - -class TestLivetimeDetermination: - def test_from_file_requires_positive_frequency(self): - """Test that zero frequency raises ValueError.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.ones(n_channels) * 10 - spectrum = ExpSpectrum(energies=energies, counts=counts) - - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=0.0, - real_time_sec=3600.0, - ) - with pytest.raises(ValueError, match="positive"): - determiner.run() - - def test_live_time_clamped_to_1(self): - """Test that live-time fraction > 1.0 is clamped to 1.0.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 10 # background - # Add prominent pulser peak at channel 100 - pulser_channel = 100 - pulser_peak = 500 * np.exp(-0.5 * ((energies - pulser_channel) / 2.0) ** 2) - counts += pulser_peak.astype(float) - spectrum = ExpSpectrum(energies=energies, counts=counts) - - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=10.0, # low frequency -> observed rate > frequency - real_time_sec=0.01, # very short time -> high observed rate - ) - result = determiner.run() - assert result.live_time_fraction == 1.0 - - def test_live_time_clamped_to_0(self): - """Test that live-time fraction < 0 is clamped to 0.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 10 # background only, no pulser peak - spectrum = ExpSpectrum(energies=energies, counts=counts) - - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) - result = determiner.run() - # With no pulser peak, the fitter may find a tiny spurious peak - # but the observed rate will be negligible -> live_time ~ 0 - assert result.live_time_fraction >= 0.0 - assert result.live_time_fraction <= 0.01 - - def test_pulser_peak_detection(self): - """Test that LivetimeDetermination correctly identifies a pulser peak.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 10 # background - # Add pulser peak at channel 100 (10 keV) - pulser_channel = 100 - pulser_peak = 500 * np.exp(-0.5 * ((energies - pulser_channel) / 2.0) ** 2) - counts += pulser_peak.astype(float) - - spectrum = ExpSpectrum( - energies=energies, - counts=counts, - source="pulser", - ) - - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) - result = determiner.run() - - assert result.live_time_fraction > 0 - assert result.live_time_fraction <= 1.0 - assert result.pulser_peak_center_keV == pytest.approx(pulser_channel, abs=5.0) - - def test_result_contains_pulser_info(self): - """Test that result contains pulser peak info.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 10 - pulser_channel = 100 - pulser_peak = 500 * np.exp(-0.5 * ((energies - pulser_channel) / 2.0) ** 2) - counts += pulser_peak.astype(float) - - spectrum = ExpSpectrum(energies=energies, counts=counts, source="pulser") - - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) - result = determiner.run() - - assert result.pulser_peak_center_keV > 0 - assert result.pulser_peak_sigma_keV > 0 - assert result.pulser_peak_chi2 > 0 - assert result.pulser_frequency_hz == 100.0 - assert result.real_time_sec == 3600.0 diff --git a/tests/exp_data/test_root_io.py b/tests/exp_data/test_root_io.py deleted file mode 100644 index 6564938..0000000 --- a/tests/exp_data/test_root_io.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Unit tests for ROOT I/O (skip if no ROOT library available).""" - -from __future__ import annotations - - -import numpy as np -import pytest - -from exp_data.spectrum import ExpSpectrum - - -def _has_root_io(): - """Check if any ROOT I/O library is available.""" - try: - import ROOT # noqa: F401 - - return True - except ImportError: - pass - try: - import uproot # noqa: F401 - - return True - except ImportError: - return False - - -@pytest.fixture -def sample_spectrum(): - """Create a sample ExpSpectrum for testing.""" - return ExpSpectrum( - energies=np.arange(100, dtype=np.float64), - counts=np.random.poisson(10, 100).astype(np.float64), - source="Tc99", - run_id="test_run", - dead_time=0.05, - live_time=3420.0, - date="2025-01-15", - metadata={ - "calibration": {"type": "linear", "coefficients": [0.1, 0.4]}, - "livetime": {"fraction": 0.95}, - }, - ) - - -@pytest.mark.skipif( - not _has_root_io(), - reason="Neither PyROOT nor uproot available", -) -class TestRootIOWriteRead: - def test_round_trip(self, sample_spectrum, tmp_path): - """Test write and read round-trip.""" - from exp_data.root_io import write_spectrum, read_spectrum - - root_path = tmp_path / "test.root" - write_spectrum(str(root_path), sample_spectrum) - read_back = read_spectrum(str(root_path)) - - np.testing.assert_array_equal(read_back.energies, sample_spectrum.energies) - np.testing.assert_array_equal(read_back.counts, sample_spectrum.counts) - - def test_scalar_fields(self, sample_spectrum, tmp_path): - """Test that scalar fields survive round-trip.""" - from exp_data.root_io import write_spectrum, read_spectrum - - root_path = tmp_path / "test.root" - write_spectrum(str(root_path), sample_spectrum) - read_back = read_spectrum(str(root_path)) - - assert read_back.dead_time == sample_spectrum.dead_time - assert read_back.source == sample_spectrum.source - assert read_back.run_id == sample_spectrum.run_id - - def test_metadata(self, sample_spectrum, tmp_path): - """Test that metadata survives round-trip.""" - from exp_data.root_io import write_spectrum, read_spectrum - - root_path = tmp_path / "test.root" - write_spectrum(str(root_path), sample_spectrum) - read_back = read_spectrum(str(root_path)) - - assert "calibration" in read_back.metadata - assert read_back.metadata["calibration"]["type"] == "linear" - assert read_back.metadata["livetime"]["fraction"] == 0.95 diff --git a/tests/exp_data/test_spectrum.py b/tests/exp_data/test_spectrum.py deleted file mode 100644 index 0423258..0000000 --- a/tests/exp_data/test_spectrum.py +++ /dev/null @@ -1,149 +0,0 @@ -"""Tests for exp_data.spectrum.""" - -import numpy as np -import pytest - -from exp_data.spectrum import ExpSpectrum - - -class TestExpSpectrumBasic: - """Basic ExpSpectrum construction and properties.""" - - def test_construct_with_errors(self): - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0, 2.0]), - counts=np.array([10.0, 20.0, 30.0]), - errors=np.array([3.16, 4.47, 5.48]), - ) - assert spectrum.n_channels == 3 - np.testing.assert_array_equal(spectrum.energies, [0.0, 1.0, 2.0]) - - def test_automatic_poisson_errors(self): - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0, 2.0]), - counts=np.array([10.0, 20.0, 30.0]), - ) - expected = np.sqrt([10.0, 20.0, 30.0]) - np.testing.assert_allclose(spectrum.errors, expected) - - def test_zero_counts_get_zero_error(self): - """Zero counts produce zero error (not sqrt(1)).""" - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0]), - counts=np.array([0.0, 100.0]), - ) - assert spectrum.errors[0] == 0.0 # sqrt(0) = 0 - - def test_shape_mismatch_raises(self): - with pytest.raises(AssertionError): - ExpSpectrum( - energies=np.array([0.0, 1.0]), - counts=np.array([10.0, 20.0, 30.0]), - ) - - def test_negative_energies_raises(self): - with pytest.raises(AssertionError): - ExpSpectrum( - energies=np.array([-1.0, 0.0]), - counts=np.array([10.0, 20.0]), - ) - - def test_negative_counts_allowed_for_difference_spectra(self): - """Negative counts are allowed — they arise in cumulative subtraction. - - Difference spectra can have negative counts in channels where the true - count is zero and Poisson noise causes the earlier cumulative to slightly - exceed the later one. - """ - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0]), - counts=np.array([10.0, -5.0]), - ) - # Should not raise; errors use abs(counts) for negative values - assert spectrum.errors[1] == np.sqrt(5.0) - - def test_total_counts(self): - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0, 2.0]), - counts=np.array([10.0, 20.0, 30.0]), - ) - assert spectrum.total_counts == 60.0 - - def test_energy_range(self): - spectrum = ExpSpectrum( - energies=np.array([0.5, 1.5, 2.5]), - counts=np.array([10.0, 20.0, 30.0]), - ) - assert spectrum.energy_range == (0.5, 2.5) - - def test_normalize(self): - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0, 2.0]), - counts=np.array([10.0, 20.0, 30.0]), - ) - norm = spectrum.normalize() - integral = np.trapezoid(norm, spectrum.energies) - assert np.isclose(integral, 1.0, rtol=1e-10) - - -class TestExpSpectrumDeadTime: - """Dead-time correction on ExpSpectrum.""" - - def test_no_correction_when_dead_time_zero(self): - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0]), - counts=np.array([10.0, 20.0]), - ) - corrected = spectrum.apply_dead_time_correction() - np.testing.assert_array_equal(corrected.counts, spectrum.counts) - - def test_correction_applied(self): - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0]), - counts=np.array([10.0, 20.0]), - dead_time=0.1, - ) - corrected = spectrum.apply_dead_time_correction() - expected = spectrum.counts / 0.9 - np.testing.assert_allclose(corrected.counts, expected, rtol=1e-10) - assert corrected.dead_time == 0.0 - - def test_live_time_corrected_rate(self): - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0]), - counts=np.array([10.0, 20.0]), - dead_time=0.1, - live_time=10.0, - ) - rate = spectrum.live_time_corrected_rate - assert rate is not None - # Total counts / (live_time * (1 - dead_time)) - expected = 30.0 / (10.0 * 0.9) - assert np.isclose(rate, expected) - - -class TestExpSpectrumSerialization: - """Dict round-trip for ExpSpectrum.""" - - def test_to_from_dict(self): - spectrum = ExpSpectrum( - energies=np.array([0.0, 1.0, 2.0]), - counts=np.array([10.0, 20.0, 30.0]), - dead_time=0.05, - live_time=100.0, - source="Tc99", - run_id="run_042", - date="2025-01-15", - metadata={"detector": "MMC", "temperature": "25 mK"}, - ) - - recovered = ExpSpectrum.from_dict(spectrum.to_dict()) - - np.testing.assert_allclose(recovered.energies, spectrum.energies) - np.testing.assert_allclose(recovered.counts, spectrum.counts) - assert recovered.dead_time == spectrum.dead_time - assert recovered.live_time == spectrum.live_time - assert recovered.source == spectrum.source - assert recovered.run_id == spectrum.run_id - assert recovered.date == spectrum.date - assert recovered.metadata == spectrum.metadata diff --git a/tests/integration/test_exp_data_pipeline.py b/tests/integration/test_exp_data_pipeline.py deleted file mode 100644 index b5650f6..0000000 --- a/tests/integration/test_exp_data_pipeline.py +++ /dev/null @@ -1,380 +0,0 @@ -"""Integration tests for the experimental data processing pipeline.""" - -from __future__ import annotations - -import numpy as np -import pytest - -from exp_data.spectrum import ExpSpectrum -from exp_data.calibration import CalibrationResult -from exp_data.pipeline import SpectrumPipeline -from exp_data.livetime import LivetimeDetermination - - -def _has_root_io(): - """Check if any ROOT I/O library is available.""" - try: - import ROOT # noqa: F401 - - return True - except ImportError: - pass - try: - import uproot # noqa: F401 - - return True - except ImportError: - return False - - -# --- ExpSpectrum.apply_calibration tests --- - - -class TestApplyCalibration: - def test_returns_new_spectrum(self, synthetic_spectrum, calibration_result): - """Test that apply_calibration returns a new ExpSpectrum.""" - calibrated = synthetic_spectrum.apply_calibration(calibration_result) - assert calibrated is not synthetic_spectrum - assert calibrated.n_channels == synthetic_spectrum.n_channels - - def test_energies_recalibrated(self, synthetic_spectrum, calibration_result): - """Test that energies are recalibrated via the calibration polynomial.""" - calibrated = synthetic_spectrum.apply_calibration(calibration_result) - expected = calibration_result(synthetic_spectrum.energies) - np.testing.assert_allclose(calibrated.energies, expected) - - def test_counts_preserved(self, synthetic_spectrum, calibration_result): - """Test that counts are preserved (copied, not modified).""" - calibrated = synthetic_spectrum.apply_calibration(calibration_result) - np.testing.assert_array_equal(calibrated.counts, synthetic_spectrum.counts) - - def test_errors_preserved(self, synthetic_spectrum, calibration_result): - """Test that errors are preserved.""" - calibrated = synthetic_spectrum.apply_calibration(calibration_result) - np.testing.assert_array_equal(calibrated.errors, synthetic_spectrum.errors) - - def test_metadata_updated_with_calibration( - self, synthetic_spectrum, calibration_result - ): - """Test that metadata includes calibration info.""" - calibrated = synthetic_spectrum.apply_calibration(calibration_result) - assert "calibration" in calibrated.metadata - assert calibrated.metadata["calibration"]["type"] == "linear" - assert ( - calibrated.metadata["calibration"]["chi2_per_dof"] - == calibration_result.chi2_per_dof - ) - - def test_scalar_fields_preserved(self, synthetic_spectrum, calibration_result): - """Test that all scalar fields are preserved.""" - spectrum = ExpSpectrum( - energies=synthetic_spectrum.energies, - counts=synthetic_spectrum.counts, - dead_time=0.05, - live_time=3420.0, - source="Tc99", - run_id="run_042", - date="2025-01-15", - ) - calibrated = spectrum.apply_calibration(calibration_result) - assert calibrated.dead_time == 0.05 - assert calibrated.live_time == 3420.0 - assert calibrated.source == "Tc99" - assert calibrated.run_id == "run_042" - assert calibrated.date == "2025-01-15" - - def test_metadata_merged_not_replaced(self, synthetic_spectrum, calibration_result): - """Test that existing metadata is merged, not replaced.""" - spectrum = ExpSpectrum( - energies=synthetic_spectrum.energies, - counts=synthetic_spectrum.counts, - metadata={"existing_key": "existing_value"}, - ) - calibrated = spectrum.apply_calibration(calibration_result) - assert "existing_key" in calibrated.metadata - assert calibrated.metadata["existing_key"] == "existing_value" - assert "calibration" in calibrated.metadata - - -# --- LivetimeDetermination tests --- - - -class TestLivetimeDetermination: - def test_pulser_peak_detection(self): - """Test that LivetimeDetermination correctly identifies a pulser peak.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 10 # background - # Add pulser peak at channel 100 (10 keV) - pulser_channel = 100 - pulser_peak = 500 * np.exp(-0.5 * ((energies - pulser_channel) / 2.0) ** 2) - counts += pulser_peak.astype(float) - - spectrum = ExpSpectrum( - energies=energies, - counts=counts, - source="pulser", - ) - - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) - result = determiner.run() - - assert result.live_time_fraction > 0 - assert result.live_time_fraction <= 1.0 - assert result.pulser_peak_center_keV == pytest.approx(pulser_channel, abs=5.0) - - def test_invalid_pulser_frequency(self): - """Test that zero frequency raises ValueError.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.ones(n_channels) * 10 - - spectrum = ExpSpectrum(energies=energies, counts=counts) - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=0.0, - real_time_sec=3600.0, - ) - with pytest.raises(ValueError, match="positive"): - determiner.run() - - def test_live_time_clamped_to_1(self): - """Test that live-time fraction > 1.0 is clamped to 1.0.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 10 # background - # Add prominent pulser peak at channel 100 - pulser_channel = 100 - pulser_peak = 500 * np.exp(-0.5 * ((energies - pulser_channel) / 2.0) ** 2) - counts += pulser_peak.astype(float) - spectrum = ExpSpectrum(energies=energies, counts=counts) - - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=10.0, # low frequency -> observed rate > frequency - real_time_sec=0.01, # very short time -> high observed rate - ) - result = determiner.run() - assert result.live_time_fraction == 1.0 - - def test_live_time_clamped_to_0(self): - """Test that live-time fraction < 0 is clamped to 0.""" - n_channels = 4096 - energies = np.arange(n_channels, dtype=np.float64) - counts = np.exp(-energies / 500) * 10 # background only, no pulser peak - spectrum = ExpSpectrum(energies=energies, counts=counts) - - determiner = LivetimeDetermination( - spectrum, - pulser_frequency_hz=100.0, - real_time_sec=3600.0, - ) - result = determiner.run() - # With no pulser peak, the fitter may find a tiny spurious peak - # but the observed rate will be negligible -> live_time ~ 0 - assert result.live_time_fraction >= 0.0 - assert result.live_time_fraction <= 0.01 - - -# --- Pipeline validation tests --- - - -class TestPipelineValidation: - def test_calib_chi2_error(self): - """Test that poor calibration fit (chi2 > 3) produces an error message.""" - pipeline = SpectrumPipeline(calibration_path="/dev/null") - spectrum = ExpSpectrum( - energies=np.arange(100, dtype=np.float64), - counts=np.ones(100) * 10, - ) - bad_calib = CalibrationResult( - coefficients=np.array([0.1, 0.4]), - order=1, - chi2_per_dof=5.0, # poor fit - n_points=2, - covariance=np.eye(2), - ) - messages = pipeline._validate(spectrum, bad_calib, None) - assert any("ERROR" in m for m in messages) - - def test_calib_chi2_warning(self): - """Test that marginal calibration fit (1.5 < chi2 < 3) produces a warning.""" - pipeline = SpectrumPipeline(calibration_path="/dev/null") - spectrum = ExpSpectrum( - energies=np.arange(100, dtype=np.float64), - counts=np.ones(100) * 10, - ) - marginal_calib = CalibrationResult( - coefficients=np.array([0.1, 0.4]), - order=1, - chi2_per_dof=2.0, # marginal fit - n_points=2, - covariance=np.eye(2), - ) - messages = pipeline._validate(spectrum, marginal_calib, None) - assert any("WARNING" in m for m in messages) - assert not any("ERROR" in m for m in messages) - - def test_insufficient_counts(self): - """Test that low counts produce an error.""" - pipeline = SpectrumPipeline(calibration_path="/dev/null") - spectrum = ExpSpectrum( - energies=np.arange(10, dtype=np.float64), - counts=np.array([1.0] * 10), # only 10 counts total - ) - good_calib = CalibrationResult( - coefficients=np.array([0.1, 0.4]), - order=1, - chi2_per_dof=0.5, - n_points=2, - covariance=np.eye(2), - ) - messages = pipeline._validate(spectrum, good_calib, None) - assert any("ERROR" in m and "counts" in m for m in messages) - - def test_negative_energy(self): - """Test that negative energies are caught after calibration.""" - pipeline = SpectrumPipeline(calibration_path="/dev/null") - # Create a spectrum with valid positive energies first - spectrum = ExpSpectrum( - energies=np.array([1.0, 2.0, 3.0, 4.0]), - counts=np.array([10.0] * 4), - ) - # Manually set energies to negative (simulating bad calibration) - spectrum.energies = np.array([-1.0, 0.0, 1.0, 2.0]) - good_calib = CalibrationResult( - coefficients=np.array([0.1, 0.4]), - order=1, - chi2_per_dof=0.5, - n_points=2, - covariance=np.eye(2), - ) - messages = pipeline._validate(spectrum, good_calib, None) - assert any("ERROR" in m and "negative" in m.lower() for m in messages) - - def test_non_monotonic_energies(self): - """Test that non-monotonic energies are caught.""" - pipeline = SpectrumPipeline(calibration_path="/dev/null") - spectrum = ExpSpectrum( - energies=np.array([0.0, 2.0, 1.0, 3.0]), - counts=np.array([10.0] * 4), - ) - good_calib = CalibrationResult( - coefficients=np.array([0.1, 0.4]), - order=1, - chi2_per_dof=0.5, - n_points=2, - covariance=np.eye(2), - ) - messages = pipeline._validate(spectrum, good_calib, None) - assert any("ERROR" in m and "monoton" in m for m in messages) - - def test_low_live_time_fraction_warning(self): - """Test that low live-time fraction produces a warning.""" - pipeline = SpectrumPipeline(calibration_path="/dev/null") - spectrum = ExpSpectrum( - energies=np.arange(100, dtype=np.float64), - counts=np.ones(100) * 10, - ) - good_calib = CalibrationResult( - coefficients=np.array([0.1, 0.4]), - order=1, - chi2_per_dof=0.5, - n_points=2, - covariance=np.eye(2), - ) - bad_livetime = type( - "MockLivetime", - (), - { - "live_time_fraction": 0.3, - }, - )() - messages = pipeline._validate(spectrum, good_calib, bad_livetime) - assert any("WARNING" in m and "live-time" in m.lower() for m in messages) - - def test_all_checks_pass(self): - """Test that good data produces no errors.""" - pipeline = SpectrumPipeline(calibration_path="/dev/null") - spectrum = ExpSpectrum( - energies=np.arange(100, dtype=np.float64), - counts=np.ones(100) * 100, - ) - good_calib = CalibrationResult( - coefficients=np.array([0.1, 0.4]), - order=1, - chi2_per_dof=0.5, - n_points=2, - covariance=np.eye(2), - ) - good_livetime = type( - "MockLivetime", - (), - { - "live_time_fraction": 0.95, - }, - )() - messages = pipeline._validate(spectrum, good_calib, good_livetime) - assert not any("ERROR" in m for m in messages) - - -# --- ROOT I/O integration tests --- - - -class TestRootIOIntegration: - @pytest.mark.skipif( - not _has_root_io(), - reason="Neither PyROOT nor uproot available", - ) - def test_round_trip(self, full_spectrum, tmp_path): - """Test write and read round-trip preserves all fields.""" - from exp_data.root_io import write_spectrum, read_spectrum - - root_path = tmp_path / "test.root" - write_spectrum(str(root_path), full_spectrum) - - read_back = read_spectrum(str(root_path)) - - np.testing.assert_array_equal(read_back.energies, full_spectrum.energies) - np.testing.assert_array_equal(read_back.counts, full_spectrum.counts) - assert read_back.dead_time == full_spectrum.dead_time - assert read_back.source == full_spectrum.source - assert read_back.run_id == full_spectrum.run_id - assert read_back.date == full_spectrum.date - - @pytest.mark.skipif( - not _has_root_io(), - reason="Neither PyROOT nor uproot available", - ) - def test_metadata_round_trip(self, full_spectrum, tmp_path): - """Test that metadata survives round-trip.""" - from exp_data.root_io import write_spectrum, read_spectrum - - root_path = tmp_path / "test.root" - write_spectrum(str(root_path), full_spectrum) - read_back = read_spectrum(str(root_path)) - - assert "calibration" in read_back.metadata - assert read_back.metadata["calibration"]["type"] == "linear" - assert "livetime" in read_back.metadata - assert read_back.metadata["livetime"]["fraction"] == 0.95 - - -def _has_root_io(): - """Check if any ROOT I/O library is available.""" - try: - import ROOT # noqa: F401 - - return True - except ImportError: - pass - try: - import uproot # noqa: F401 - - return True - except ImportError: - return False