feat: experimental data calibration pipeline with ROOT I/O
- Fix chi2 calculation bug in calibration.py (remove incorrect coeffs[::-1] reversal in np.polyval) - Fix broken beta_spectrum/cli.py (remove orphaned detector response add_argument calls and unused imports) - Add ExpSpectrum.apply_calibration() and EnergyCalibrator.from_am241_spectrum() - Add LivetimeDetermination class for pulser-based livetime - Add SpectrumPipeline orchestration (read→subtract→calibrate→livetime→validate) - Add ROOT I/O with PyROOT primary + uproot fallback - Add bs_exp CLI with process, livetime, export subcommands - Add exp_calibration notebook with pre-calibration visualization - Add comprehensive tests (73 passing) - Add bs_exp console script and [root] extra to pyproject.toml - Remove obsolete detector_response module and unused notebooks
This commit is contained in:
+18
-1
@@ -6,20 +6,37 @@ experimental spectra from detector readouts.
|
||||
"""
|
||||
|
||||
from .spectrum import ExpSpectrum
|
||||
from .calibration import EnergyCalibrator
|
||||
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",
|
||||
]
|
||||
|
||||
+109
-1
@@ -13,9 +13,11 @@ Supports:
|
||||
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
|
||||
@@ -183,7 +185,9 @@ class EnergyCalibrator:
|
||||
cov = np.eye(order + 1) * 1e-10
|
||||
|
||||
# Compute χ²
|
||||
y_fit = np.polyval(coeffs[::-1], x)
|
||||
# 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))
|
||||
@@ -289,3 +293,107 @@ class EnergyCalibrator:
|
||||
]
|
||||
)
|
||||
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.8, 17.8, 26.3].
|
||||
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 peaks in low-energy region
|
||||
counts_low = counts.copy()
|
||||
counts_low[300:] = 0
|
||||
peaks, _ = find_peaks(counts_low, height=10, distance=10)
|
||||
peak_energies = np.array(peaks, dtype=float)
|
||||
|
||||
# Known Am-241 gamma lines
|
||||
if known_energies is None:
|
||||
known_energies = [13.8, 17.8, 26.3]
|
||||
|
||||
# Match peaks to known energies.
|
||||
# Am-241 has well-known lines: 13.8, 17.8, 26.3 keV (and 59.54 keV).
|
||||
# Strategy: find the two closest peaks with the largest energy gap,
|
||||
# assign the known energies to them proportionally, then verify
|
||||
# with remaining peaks.
|
||||
channel_to_energy: dict[float, float] = {}
|
||||
|
||||
if len(peak_energies) >= 2 and len(known_energies) >= 2:
|
||||
# Find the two peaks with the largest separation — most reliable
|
||||
# for gain estimation
|
||||
max_separation = 0
|
||||
best_pair = (0, 1)
|
||||
for i in range(len(peak_energies)):
|
||||
for j in range(i + 1, len(peak_energies)):
|
||||
sep = peak_energies[j] - peak_energies[i]
|
||||
if sep > max_separation:
|
||||
max_separation = sep
|
||||
best_pair = (i, j)
|
||||
|
||||
# Use the full known energy range to estimate gain
|
||||
e_min, e_max = min(known_energies), max(known_energies)
|
||||
if e_max != e_min:
|
||||
gain = (peak_energies[best_pair[1]] - peak_energies[best_pair[0]]) / (
|
||||
e_max - e_min
|
||||
)
|
||||
intercept = peak_energies[best_pair[0]] - gain * e_min
|
||||
else:
|
||||
gain = 1.0
|
||||
intercept = 0.0
|
||||
|
||||
# Map each known energy to a channel and verify against peaks
|
||||
for target_energy in known_energies:
|
||||
expected_channel = gain * target_energy + intercept
|
||||
if 0 <= expected_channel < n_channels:
|
||||
best_idx = np.argmin(np.abs(peak_energies - expected_channel))
|
||||
detected_channel = float(peak_energies[best_idx])
|
||||
if np.abs(detected_channel - expected_channel) < 10.0:
|
||||
channel_to_energy[detected_channel] = target_energy
|
||||
|
||||
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
|
||||
calibrator = cls()
|
||||
calibrator.set_calibration_points(list(channel_to_energy.items()))
|
||||
return calibrator.fit_calibration(order=1)
|
||||
|
||||
@@ -0,0 +1,421 @@
|
||||
"""
|
||||
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()
|
||||
@@ -0,0 +1,219 @@
|
||||
"""
|
||||
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()
|
||||
@@ -0,0 +1,307 @@
|
||||
"""
|
||||
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
|
||||
)
|
||||
@@ -163,14 +163,16 @@ def calibrate_from_am241(
|
||||
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)
|
||||
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"]
|
||||
_peak_heights = props["peak_heights"] # noqa: F841
|
||||
|
||||
# Known Am-241 gamma lines (keV)
|
||||
am241_lines = {13.8: None, 17.8: None, 26.3: None}
|
||||
@@ -221,7 +223,7 @@ def read_am241_calibration(
|
||||
# Am-241 known peaks (keV): 13.8, 17.8, 26.3, 59.54
|
||||
# The spectrum typically has a few prominent peaks
|
||||
with open(filepath, "rb") as f:
|
||||
header = f.read(128)
|
||||
_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)
|
||||
|
||||
@@ -0,0 +1,287 @@
|
||||
"""
|
||||
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,
|
||||
)
|
||||
+42
-1
@@ -8,7 +8,10 @@ and calibration information.
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from exp_data.calibration import CalibrationResult
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -140,6 +143,44 @@ class ExpSpectrum:
|
||||
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 {
|
||||
|
||||
Reference in New Issue
Block a user