unzhakov
38b44b56e6
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.
449 lines
14 KiB
Python
449 lines
14 KiB
Python
# tests/test_fitter.py — Tests for χ² minimization framework
|
||
|
||
"""
|
||
Why test this?
|
||
--------------
|
||
The fitter is the core numerical tool for extracting physical parameters
|
||
from beta spectra. It must be correct because all downstream results
|
||
(g_V, g_A, C(W) extraction) depend on it.
|
||
|
||
We test:
|
||
1. Basic χ² minimization on known functions (quadratic, linear).
|
||
2. Covariance matrix estimation (should match input uncertainties).
|
||
3. Confidence interval computation via profile likelihood.
|
||
4. Goodness-of-fit metrics (χ²/ndf, p-value).
|
||
5. Bounds handling and convergence diagnostics.
|
||
6. Robust loss functions (soft_l1, huber) for outlier resistance.
|
||
"""
|
||
|
||
import numpy as np
|
||
|
||
from fitter.fit_engine import CurveFitter, FitConfig
|
||
|
||
|
||
class TestBasicFitting:
|
||
"""Test χ² minimization on simple known functions."""
|
||
|
||
def test_quadratic_fit(self):
|
||
"""Fit y = a + b·x + c·x² with known coefficients."""
|
||
np.random.seed(42)
|
||
x = np.linspace(0, 10, 50)
|
||
true_params = [1.0, 2.0, -0.1]
|
||
y = true_params[0] + true_params[1] * x + true_params[2] * x**2
|
||
y += np.random.normal(0, 0.5, len(x))
|
||
|
||
def model(x, a, b, c):
|
||
return a + b * x + c * x**2
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(y) * 0.5)
|
||
result = fitter.fit([0.0, 0.0, 0.0])
|
||
|
||
for i, name in enumerate(["a", "b", "c"]):
|
||
assert (
|
||
abs(result.parameters[i] - true_params[i]) < 0.5
|
||
), f"Parameter {name}: expected {true_params[i]}, got {result.parameters[i]}"
|
||
|
||
assert result.success, "Fit should converge"
|
||
assert result.chi2 > 0, "χ² must be positive"
|
||
|
||
def test_linear_fit(self):
|
||
"""Fit y = a + b·x with known coefficients."""
|
||
np.random.seed(123)
|
||
x = np.linspace(0, 5, 30)
|
||
true_a, true_b = 3.0, -1.5
|
||
y = true_a + true_b * x
|
||
y += np.random.normal(0, 0.2, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(y) * 0.2)
|
||
result = fitter.fit([0.0, 0.0])
|
||
|
||
assert abs(result.parameters[0] - true_a) < 0.3
|
||
assert abs(result.parameters[1] - true_b) < 0.3
|
||
assert result.success
|
||
|
||
def test_constant_fit(self):
|
||
"""Fit y = a (constant) with known value."""
|
||
np.random.seed(456)
|
||
x = np.linspace(0, 10, 40)
|
||
true_a = 5.0
|
||
y = np.full_like(x, true_a)
|
||
y += np.random.normal(0, 0.3, len(x))
|
||
|
||
def model(x, a):
|
||
return np.full_like(x, a)
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.3)
|
||
result = fitter.fit([0.0])
|
||
|
||
assert abs(result.parameters[0] - true_a) < 0.2
|
||
assert result.success
|
||
|
||
|
||
class TestCovarianceEstimation:
|
||
"""Test that covariance matrix is correctly estimated."""
|
||
|
||
def test_covariance_matches_uncertainty(self):
|
||
"""For a constant model, parameter error should match input σ."""
|
||
np.random.seed(789)
|
||
x = np.linspace(0, 10, 50)
|
||
true_a = 5.0
|
||
sigma = 0.5
|
||
y = np.full_like(x, true_a, dtype=float)
|
||
y += np.random.normal(0, sigma, len(x))
|
||
|
||
def model(x, a):
|
||
return np.full_like(x, a)
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * sigma)
|
||
result = fitter.fit([0.0])
|
||
|
||
error = np.sqrt(result.covariance[0, 0])
|
||
# Expected error: σ / sqrt(N)
|
||
expected_error = sigma / np.sqrt(len(x))
|
||
assert (
|
||
abs(error - expected_error) < expected_error * 0.3
|
||
), f"Covariance error {error:.4f} != expected {expected_error:.4f}"
|
||
|
||
def test_covariance_matrix_positive_definite(self):
|
||
"""Covariance matrix must be positive definite."""
|
||
np.random.seed(321)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 2.0 + 1.5 * x - 0.3 * x**2
|
||
y += np.random.normal(0, 0.2, len(x))
|
||
|
||
def model(x, a, b, c):
|
||
return a + b * x + c * x**2
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.2)
|
||
result = fitter.fit([0.0, 0.0, 0.0])
|
||
|
||
# Check eigenvalues are non-negative
|
||
eigvals = np.linalg.eigvalsh(result.covariance)
|
||
assert np.all(eigvals >= -1e-10), f"Covariance not positive definite: {eigvals}"
|
||
|
||
def test_correlation_matrix(self):
|
||
"""Correlation matrix should have 1.0 on diagonal."""
|
||
np.random.seed(654)
|
||
x = np.linspace(0, 10, 50)
|
||
y = 1.0 + 2.0 * x + 0.1 * x**2
|
||
y += np.random.normal(0, 0.5, len(x))
|
||
|
||
def model(x, a, b, c):
|
||
return a + b * x + c * x**2
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.5)
|
||
result = fitter.fit([0.0, 0.0, 0.0])
|
||
|
||
diag = np.diag(result.correlation_matrix)
|
||
assert np.allclose(
|
||
diag, 1.0, atol=1e-6
|
||
), f"Diagonal of correlation matrix: {diag}"
|
||
|
||
|
||
class TestGoodnessOfFit:
|
||
"""Test χ²/ndf and p-value calculations."""
|
||
|
||
def test_chi2_per_dof(self):
|
||
"""χ²/ndof should be positive and reasonable."""
|
||
np.random.seed(111)
|
||
x = np.linspace(0, 10, 50)
|
||
y = 2.0 * x
|
||
y += np.random.normal(0, 0.3, len(x))
|
||
|
||
def model(x, a):
|
||
return a * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.3)
|
||
result = fitter.fit([0.0])
|
||
|
||
assert result.chi2_per_dof > 0
|
||
assert result.n_points == 50
|
||
assert result.n_free == 1
|
||
assert result.n_free == 1
|
||
|
||
def test_p_value_range(self):
|
||
"""p-value should be in [0, 1]."""
|
||
np.random.seed(222)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 1.0 + x
|
||
y += np.random.normal(0, 0.1, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.1)
|
||
result = fitter.fit([0.0, 0.0])
|
||
|
||
assert 0.0 <= result.p_value <= 1.0
|
||
|
||
def test_p_value_good_fit(self):
|
||
"""Good fit should have reasonable p-value."""
|
||
np.random.seed(333)
|
||
x = np.linspace(0, 10, 100)
|
||
y = 3.0 + 1.5 * x - 0.05 * x**2
|
||
y += np.random.normal(0, 0.2, len(x))
|
||
|
||
def model(x, a, b, c):
|
||
return a + b * x + c * x**2
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.2)
|
||
result = fitter.fit([0.0, 0.0, 0.0])
|
||
|
||
assert (
|
||
result.p_value > 0.001
|
||
), f"Good fit should have reasonable p-value, got {result.p_value}"
|
||
|
||
|
||
class TestBounds:
|
||
"""Test parameter bounds handling."""
|
||
|
||
def test_lower_bound(self):
|
||
"""Test that lower bounds are respected."""
|
||
np.random.seed(444)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 2.0 + 1.0 * x
|
||
y += np.random.normal(0, 0.1, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.1)
|
||
result = fitter.fit([1.5, 0.5], bounds=([1.0, 0.0], [10.0, 5.0]))
|
||
|
||
assert (
|
||
result.parameters[0] >= 1.0
|
||
), f"Parameter a should be >= 1.0, got {result.parameters[0]}"
|
||
assert (
|
||
result.parameters[1] >= 0.0
|
||
), f"Parameter b should be >= 0.0, got {result.parameters[1]}"
|
||
|
||
def test_upper_bound(self):
|
||
"""Test that upper bounds are respected."""
|
||
np.random.seed(555)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 0.5 + 0.5 * x
|
||
y += np.random.normal(0, 0.1, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.1)
|
||
result = fitter.fit([0.5, 0.5], bounds=([0.0, 0.0], [1.0, 1.0]))
|
||
|
||
assert result.parameters[0] <= 1.0
|
||
assert result.parameters[1] <= 1.0
|
||
|
||
|
||
class TestRobustLoss:
|
||
"""Test robust loss functions for outlier resistance."""
|
||
|
||
def test_soft_l1_loss(self):
|
||
"""soft_l1 loss should handle outliers better than linear."""
|
||
np.random.seed(666)
|
||
x = np.linspace(0, 10, 50)
|
||
y = 1.0 + 2.0 * x
|
||
y[25] = 50.0 # Add outlier
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
# Linear loss — outlier pulls fit away
|
||
config_linear = FitConfig(loss="linear")
|
||
fitter_linear = CurveFitter(
|
||
model, x, y, uncertainties=np.ones_like(x) * 1.0, config=config_linear
|
||
)
|
||
result_linear = fitter_linear.fit([0.0, 0.0])
|
||
|
||
# soft_l1 loss — outlier has less influence
|
||
config_robust = FitConfig(loss="soft_l1")
|
||
fitter_robust = CurveFitter(
|
||
model, x, y, uncertainties=np.ones_like(x) * 1.0, config=config_robust
|
||
)
|
||
result_robust = fitter_robust.fit([0.0, 0.0])
|
||
|
||
# Robust fit should be closer to true slope
|
||
assert (
|
||
abs(result_robust.parameters[1] - 2.0)
|
||
< abs(result_linear.parameters[1] - 2.0) * 1.5
|
||
), "Robust loss should handle outliers better"
|
||
|
||
|
||
class TestFitResult:
|
||
"""Test FitResult dataclass methods."""
|
||
|
||
def test_parameters_with_errors(self):
|
||
"""parameters_with_errors should return dict of (value, error)."""
|
||
np.random.seed(777)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 2.0 + x
|
||
y += np.random.normal(0, 0.1, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.1)
|
||
result = fitter.fit([0.0, 0.0])
|
||
|
||
errors = result.parameters_with_errors
|
||
assert len(errors) == 2
|
||
for key, (val, err) in errors.items():
|
||
assert isinstance(val, float)
|
||
assert isinstance(err, float)
|
||
assert err >= 0
|
||
|
||
def test_summary(self):
|
||
"""summary() should return a non-empty string."""
|
||
np.random.seed(888)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 1.0 + 2.0 * x
|
||
y += np.random.normal(0, 0.1, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.1)
|
||
result = fitter.fit([0.0, 0.0])
|
||
|
||
summary = result.summary(param_names=["a", "b"])
|
||
assert "Fit Results" in summary
|
||
assert "a" in summary
|
||
assert "b" in summary
|
||
assert "χ²" in summary
|
||
assert "Converged" in summary
|
||
|
||
def test_residuals_property(self):
|
||
"""residuals should be data - model."""
|
||
np.random.seed(999)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 2.0 + x
|
||
y += np.random.normal(0, 0.1, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.1)
|
||
result = fitter.fit([0.0, 0.0])
|
||
|
||
expected_residuals = y - result.model_values
|
||
np.testing.assert_allclose(result.residuals, expected_residuals)
|
||
|
||
|
||
class TestFitConfig:
|
||
"""Test FitConfig options."""
|
||
|
||
def test_custom_config(self):
|
||
"""Custom FitConfig should be respected."""
|
||
np.random.seed(101)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 1.0 + x
|
||
y += np.random.normal(0, 0.1, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
config = FitConfig(
|
||
method="trf",
|
||
max_nfev=500,
|
||
loss="linear",
|
||
ftol=1e-10,
|
||
)
|
||
fitter = CurveFitter(
|
||
model, x, y, uncertainties=np.ones_like(x) * 0.1, config=config
|
||
)
|
||
result = fitter.fit([0.0, 0.0])
|
||
|
||
assert result.success
|
||
|
||
|
||
class TestConfidenceIntervals:
|
||
"""Test confidence interval computation."""
|
||
|
||
def test_95_percent_ci(self):
|
||
"""95% CI should contain the true parameter value."""
|
||
np.random.seed(202)
|
||
x = np.linspace(0, 10, 60)
|
||
true_a, true_b = 2.0, 1.5
|
||
y = true_a + true_b * x
|
||
y += np.random.normal(0, 0.3, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.3)
|
||
result = fitter.fit([0.0, 0.0])
|
||
|
||
ci_a = result.confidence_interval(0, confidence=0.95)
|
||
ci_b = result.confidence_interval(1, confidence=0.95)
|
||
|
||
assert ci_a[0] <= true_a <= ci_a[1], f"True a={true_a} not in CI {ci_a}"
|
||
assert ci_b[0] <= true_b <= ci_b[1], f"True b={true_b} not in CI {ci_b}"
|
||
|
||
def test_ci_width_increases_with_confidence(self):
|
||
"""Higher confidence level → wider interval."""
|
||
np.random.seed(303)
|
||
x = np.linspace(0, 5, 40)
|
||
y = 1.0 + 2.0 * x
|
||
y += np.random.normal(0, 0.2, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.2)
|
||
result = fitter.fit([0.0, 0.0])
|
||
|
||
ci_90 = result.confidence_interval(0, confidence=0.90)
|
||
ci_95 = result.confidence_interval(0, confidence=0.95)
|
||
|
||
width_90 = ci_90[1] - ci_90[0]
|
||
width_95 = ci_95[1] - ci_95[0]
|
||
|
||
assert (
|
||
width_95 > width_90
|
||
), f"95% CI width {width_95} should be > 90% width {width_90}"
|
||
|
||
|
||
class TestEdgeCases:
|
||
"""Test edge cases and error handling."""
|
||
|
||
def test_zero_uncertainties_filtered(self):
|
||
"""Points with zero uncertainty should be filtered out."""
|
||
x = np.array([1.0, 2.0, 3.0, 4.0])
|
||
y = np.array([2.0, 4.0, 6.0, 8.0])
|
||
unc = np.array([1.0, 0.0, 1.0, 0.0])
|
||
|
||
def model(x, a):
|
||
return a * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=unc)
|
||
assert len(fitter.x_data) == 2, "Zero-uncertainty points should be removed"
|
||
assert len(fitter.y_data) == 2
|
||
|
||
def test_negative_uncertainties_filtered(self):
|
||
"""Points with negative uncertainty should be filtered out."""
|
||
x = np.array([1.0, 2.0, 3.0])
|
||
y = np.array([2.0, 4.0, 6.0])
|
||
unc = np.array([1.0, -1.0, 1.0])
|
||
|
||
def model(x, a):
|
||
return a * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=unc)
|
||
assert len(fitter.x_data) == 2
|
||
|
||
def test_unbounded_parameters(self):
|
||
"""Unbounded parameters (±inf) should work."""
|
||
np.random.seed(404)
|
||
x = np.linspace(0, 5, 30)
|
||
y = 1.0 + 2.0 * x
|
||
y += np.random.normal(0, 0.1, len(x))
|
||
|
||
def model(x, a, b):
|
||
return a + b * x
|
||
|
||
fitter = CurveFitter(model, x, y, uncertainties=np.ones_like(x) * 0.1)
|
||
result = fitter.fit([0.0, 0.0], bounds=([-np.inf, -np.inf], [np.inf, np.inf]))
|
||
assert result.success
|