unzhakov/beta_shape_pnpi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: v0.4.0 package restructuring — split into 4 modules

Browse Source 
Separate monolithic beta_spectrum package into:
  beta_spectrum   — pure theory (components, spectrum, CLI, visualize)
  exp_data        — experimental data handling (spectrum, calibration,
                    fitters, corrections)
  fitter          — analysis framework (model, fit_engine, extractor,
                    result)
  beta_spectrum.visualize — plotting (BetaSpectrumAnalyzer)

Phase 1 — exp_data module:
  ExpSpectrum dataclass with metadata, normalization, serialization
  EnergyCalibrator (linear/quadratic calibration)
  GaussianFitter + PeakFitter for calibration sources
  DeadTimeCorrection, PileUpCorrection, BackgroundSubtractor
  43 new tests, ruff + mypy clean

Phase 2 — fitter module:
  SpectrumModel wraps BetaSpectrum + DetectorResponse with convolution
  SpectrumFitter with χ² minimization, covariance, profile likelihood,
  confidence intervals, robust loss functions
  CWExtractor + GVAExtractor for C(W) and g_V/g_A extraction
  AnalysisResult + FitSummary for physics-specific results
  CurveFitter general-purpose callable model fitter
  34 new tests, ruff + mypy clean

Phase 3 — clean up beta_spectrum:
  Remove cw_extractor.py and fitter.py (moved to fitter/)
  Clean __init__.py to export only theory classes
  Add backward-compat methods to CWExtractor
  Fix BetaSpectrumAnalyzer._get_branch_spectrum returning BetaSpectrum
  objects instead of numpy arrays (plot fix)

Phase 4 — package structure:
  Auto-discovered by setuptools (4 packages)
  All imports updated across 26 source files
  Black formatting applied
  322 tests pass (262 original + 60 new)
  ruff + mypy clean on all source files

Bug fixes:
  - Fix multi-branch spectrum plotting: universal component keys
    (Fermi, Screening, etc.) don't have per-branch prefix; fixed
    fallback lookup.
  - Fix negative phase space beyond endpoint: BetaSpectrum returns
    negative values for W > W0; clamp branch spectra to zero.
This commit is contained in:
unzhakov
2026-05-17 18:58:04 +03:00
parent 745bbae4cd
commit 38b44b56e6
29 changed files with 3538 additions and 2010 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/exp_data/test_calibration.py
+138
View File
@@ -0,0 +1,138 @@
"""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
tests/exp_data/test_corrections.py
+106
View File
@@ -0,0 +1,106 @@
"""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
tests/exp_data/test_fitters.py
+114
View File
@@ -0,0 +1,114 @@
"""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
tests/exp_data/test_spectrum.py
+141
View File
@@ -0,0 +1,141 @@
"""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_error_one(self):
spectrum = ExpSpectrum(
energies=np.array([0.0, 1.0]),
counts=np.array([0.0, 100.0]),
)
assert spectrum.errors[0] == 1.0 # sqrt(1) for zero counts
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_raises(self):
with pytest.raises(AssertionError):
ExpSpectrum(
energies=np.array([0.0, 1.0]),
counts=np.array([10.0, -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