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

feat: add χ² fitter and C(W) shape factor extraction modules

Browse Source 
Add CurveFitter with scipy least_squares, configurable loss functions
(soft_l1, huber, cauchy), profile likelihood, confidence intervals, and
goodness-of-fit metrics. Add CWExtractor for extracting shape factors
from beta spectra via Kurie analysis, parametrization fitting, and
g_V/g_A extraction. Includes 32 new tests (18 fitter + 14 CW extractor).
Bump version to 0.3.0.
This commit is contained in:
unzhakov
2026-05-02 09:40:38 +03:00
parent 97aaddd86b
commit 272841f50d
5 changed files with 2381 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_cw_extractor.py
+675
View File
@@ -0,0 +1,675 @@
# tests/test_cw_extractor.py — Tests for C(W) shape factor extraction
"""
Why test this?
--------------
The C(W) extractor is the core scientific tool for extracting the shape factor
from measured beta spectra. It must correctly handle:
1. Energy unit conversions (keV ↔ m_e units)
2. Theoretical factor computation (Fermi, screening, radiative, etc.)
3. Normalization of extracted C(W)
4. Kurie plot endpoint determination
5. g_V, g_A extraction from C(W) data
We test with synthetic data where the true answer is known, so we can verify
the extraction recovers the correct values.
"""
import numpy as np
import pytest
from beta_spectrum import (
BetaSpectrum,
CWExtractor,
CWExtractionResult,
GVAExtractionResult,
SpectrumConfig,
T_to_W,
W_to_T,
)
class TestCWExtractionBasic:
"""Basic C(W) extraction tests."""
def _make_synthetic_data(
self,
endpoint_keV: float = 294.0,
n_points: int = 100,
seed: int = 42,
):
"""Create synthetic beta spectrum with C(W) = 1 (allowed decay)."""
np.random.seed(seed)
# Energy grid
energies_keV = np.linspace(5, endpoint_keV * 0.98, n_points)
# Build theoretical spectrum with C(W) = 1
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=endpoint_keV / 1000.0,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
spectrum = BetaSpectrum.from_config(config)
W, E_MeV = spectrum.get_energy_grid(config)
# Interpolate to our energy grid
from scipy.interpolate import interp1d
flux_interp = interp1d(
E_MeV * 1000, spectrum(W), kind="linear", fill_value=0.0, bounds_error=False
)
flux_theory = flux_interp(energies_keV)
# Add noise (Poisson-like)
noise_std = np.sqrt(np.maximum(flux_theory, 1e-10))
flux = flux_theory + np.random.normal(0, noise_std, len(flux_theory))
flux = np.maximum(flux, 0.0) # Flux cannot be negative
flux_errors = np.sqrt(np.maximum(flux, 1e-10))
return energies_keV, flux, flux_errors, config
def test_CW_extraction_allowed_decay(self):
"""For allowed decay (C(W)=1), extracted C(W) should be ~1."""
energies_keV, flux, flux_errors, config = self._make_synthetic_data()
extractor = CWExtractor(config, flux, energies_keV, flux_errors)
result = extractor.extract_CW(endpoint_keV=294.0)
# C(W) should be close to 1 in mid-energy region
mid_mask = (energies_keV > 50) & (energies_keV < 200)
if np.any(mid_mask):
avg_cw = np.mean(result.cw_values[mid_mask])
assert (
abs(avg_cw - 1.0) < 0.3
), f"Allowed decay C(W) should be ~1, got {avg_cw}"
def test_CW_extraction_preserves_endpoint(self):
"""Extracted C(W) should report correct endpoint."""
energies_keV, flux, flux_errors, config = self._make_synthetic_data()
extractor = CWExtractor(config, flux, energies_keV, flux_errors)
result = extractor.extract_CW(endpoint_keV=294.0)
assert (
abs(result.endpoint_keV - 294.0) < 1.0
), f"Endpoint should be ~294 keV, got {result.endpoint_keV}"
def test_CW_extraction_errors_positive(self):
"""C(W) uncertainties should be positive."""
energies_keV, flux, flux_errors, config = self._make_synthetic_data()
extractor = CWExtractor(config, flux, energies_keV, flux_errors)
result = extractor.extract_CW(endpoint_keV=294.0)
assert np.all(result.cw_errors > 0), "All errors should be positive"
def test_CW_extraction_clipped(self):
"""C(W) values should be clipped to [0.01, 10.0]."""
energies_keV, flux, flux_errors, config = self._make_synthetic_data()
extractor = CWExtractor(config, flux, energies_keV, flux_errors)
result = extractor.extract_CW(endpoint_keV=294.0)
assert np.all(result.cw_values >= 0.01), "C(W) should be >= 0.01"
assert np.all(result.cw_values <= 10.0), "C(W) should be <= 10.0"
class TestCWParametrization:
"""Test C(W) parametrization fitting."""
def _make_CW_data(
self,
endpoint_keV: float = 294.0,
n_points: int = 60,
seed: int = 42,
):
"""Create synthetic C(W) data with known parametrization."""
np.random.seed(seed)
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=endpoint_keV / 1000.0,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
# Generate C(W) values with known form
W = np.linspace(T_to_W(0.05), T_to_W(endpoint_keV * 0.95), n_points)
E_keV = np.array([W_to_T(w) for w in W])
# Linear C(W): C(W) = 1.0 + 0.5*(W-1)
true_a0, true_a1 = 1.0, 0.5
cw_true = true_a0 + true_a1 * (W - 1.0)
cw_noise = cw_true + np.random.normal(0, 0.05, len(W))
# C(W) = flux / (C0 * phase_space * Fermi * corrections * (W0-W)^2)
# So: flux = C(W) * C0 * phase_space * Fermi * corrections * (W0-W)^2
p = np.sqrt(np.maximum(W**2 - 1.0, 0.0))
from beta_spectrum.components.fermi import FermiFunction
F = FermiFunction()(W, config.Z_parent)
C0 = 1.0 # Normalization
endpoint_factor = (T_to_W(endpoint_keV / 1000.0) - W) ** 2
flux_synthetic = cw_noise * C0 * p * W * F * endpoint_factor
flux_synthetic = np.maximum(flux_synthetic, 0)
flux_errors = np.sqrt(np.maximum(flux_synthetic, 1.0))
return E_keV, flux_synthetic, flux_errors, config, true_a0, true_a1
def test_constant_parametrization(self):
"""Fit C(W) = a₀ with known constant value."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 50)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
cw_true = np.full_like(W, 1.0)
cw = cw_true + np.random.normal(0, 0.05, len(W))
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
extractor = CWExtractor(
config, cw, energies_keV, flux_errors=np.ones_like(cw) * 0.05
)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.05,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.05,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
result = extractor.fit_CW_parametrization(cw_result, parametrization="constant")
assert (
abs(result.parameters[0] - 1.0) < 0.1
), f"Constant C(W) should be ~1, got {result.parameters[0]}"
def test_linear_parametrization(self):
"""Fit C(W) = a₀ + a₁·(W-1) with known coefficients."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 50)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
true_a0, true_a1 = 1.0, 0.5
cw = true_a0 + true_a1 * (W - 1.0) + np.random.normal(0, 0.05, len(W))
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
extractor = CWExtractor(
config, cw, energies_keV, flux_errors=np.ones_like(cw) * 0.05
)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.05,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.05,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
result = extractor.fit_CW_parametrization(cw_result, parametrization="linear")
assert (
abs(result.parameters[0] - true_a0) < 0.15
), f"a₀: expected {true_a0}, got {result.parameters[0]}"
assert (
abs(result.parameters[1] - true_a1) < 0.5
), f"a₁: expected {true_a1}, got {result.parameters[1]}"
def test_quadratic_parametrization(self):
"""Fit C(W) = a₀ + a₁·(W-1) + a₂·(W-1)² with known coefficients."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 60)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
true_a0, true_a1, true_a2 = 1.0, 0.3, -0.1
cw = true_a0 + true_a1 * (W - 1.0) + true_a2 * (W - 1.0) ** 2
cw += np.random.normal(0, 0.03, len(W))
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
extractor = CWExtractor(
config, cw, energies_keV, flux_errors=np.ones_like(cw) * 0.03
)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.03,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.03,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
result = extractor.fit_CW_parametrization(
cw_result, parametrization="quadratic"
)
assert abs(result.parameters[0] - true_a0) < 0.15
assert abs(result.parameters[1] - true_a1) < 0.5
assert abs(result.parameters[2] - true_a2) < 0.5
def test_invalid_parametrization(self):
"""Invalid parametrization should raise ValueError."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 30)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
cw = np.ones_like(W)
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
)
extractor = CWExtractor(config, cw, energies_keV)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.05,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.05,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
with pytest.raises(ValueError, match="Unknown parametrization"):
extractor.fit_CW_parametrization(cw_result, parametrization="invalid")
class TestGVAExtraction:
"""Test g_V, g_A extraction from C(W) data."""
def test_gV_gA_pure_fermi(self):
"""For pure Fermi transition, g_A ≈ 0."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 50)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
# Pure Fermi: C(W) = g_V² (constant)
true_gV = 0.5
cw = true_gV**2 + np.random.normal(0, 0.02, len(W))
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
extractor = CWExtractor(
config, cw, energies_keV, flux_errors=np.ones_like(cw) * 0.02
)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.02,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.02,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
result = extractor.fit_gV_gA(cw_result, M_F=1.0, M_GT=0.0)
assert isinstance(result, GVAExtractionResult)
assert (
abs(result.g_V - true_gV) < 0.15
), f"g_V: expected {true_gV}, got {result.g_V}"
def test_gV_gA_pure_gt(self):
"""For pure Gamow-Teller, g_V ≈ 0."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 50)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
# Pure GT: C(W) = g_A² (constant)
true_gA = 0.8
cw = true_gA**2 + np.random.normal(0, 0.02, len(W))
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
extractor = CWExtractor(
config, cw, energies_keV, flux_errors=np.ones_like(cw) * 0.02
)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.02,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.02,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
result = extractor.fit_gV_gA(cw_result, M_F=0.0, M_GT=1.0)
assert isinstance(result, GVAExtractionResult)
assert (
abs(result.g_A - true_gA) < 0.15
), f"g_A: expected {true_gA}, got {result.g_A}"
def test_gV_gA_mixed_transition(self):
"""For mixed transition, extract both g_V and g_A."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 50)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
# Mixed: C(W) = g_V² + g_A² (constant for simplicity)
true_gV, true_gA = 0.4, 0.6
cw = true_gV**2 + true_gA**2 + np.random.normal(0, 0.02, len(W))
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
extractor = CWExtractor(
config, cw, energies_keV, flux_errors=np.ones_like(cw) * 0.02
)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.02,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.02,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
result = extractor.fit_gV_gA(cw_result, M_F=1.0, M_GT=1.0)
assert isinstance(result, GVAExtractionResult)
# Check that sum of squares is correct
reconstructed = result.g_V**2 + result.g_A**2
expected = true_gV**2 + true_gA**2
assert (
abs(reconstructed - expected) < 0.15
), f"g_V²+g_A²: expected {expected}, got {reconstructed}"
def test_gV_gA_summary(self):
"""summary() should return a non-empty string with key info."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 50)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
cw = np.full_like(W, 0.5)
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
)
extractor = CWExtractor(config, cw, energies_keV)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.02,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.02,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
result = extractor.fit_gV_gA(cw_result, M_F=1.0, M_GT=0.0)
summary = result.summary()
assert "g_V" in summary
assert "g_A" in summary
assert "χ²" in summary
assert "p-value" in summary
class TestKurieAnalysis:
"""Test Kurie plot analysis."""
def test_kurie_analysis_allowed_decay(self):
"""Kurie analysis on allowed decay should produce reasonable C(W)."""
np.random.seed(42)
energies_keV = np.linspace(5, 280, 80)
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
transition_type="F1",
use_screening=True,
use_exchange=True,
use_finite_size=True,
use_radiative=True,
)
spectrum = BetaSpectrum.from_config(config)
W, E_MeV = spectrum.get_energy_grid(config)
flux_theory = spectrum(W)
# Interpolate to our energy grid
from scipy.interpolate import interp1d
flux_interp = interp1d(
E_MeV * 1000, flux_theory, kind="linear", fill_value=0.0, bounds_error=False
)
flux = flux_interp(energies_keV)
# Add noise
flux += np.random.normal(0, 0.02 * np.max(flux), len(flux))
extractor = CWExtractor(config, flux, energies_keV)
cw_result, fit_result = extractor.kurie_analysis(endpoint_keV=294.0)
assert isinstance(cw_result, CWExtractionResult)
assert cw_result.fit_result is not None
assert np.all(cw_result.cw_values > 0)
def test_kurie_plot_values(self):
"""Kurie plot should produce valid values."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 50)
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
cw = np.ones_like(W) * 1.0
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.05,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.05,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
E_plot, kurie = cw_result.kurie_plot()
assert len(E_plot) == len(energies_keV)
assert len(kurie) == len(energies_keV)
assert np.all(kurie >= 0), "Kurie values should be non-negative"
class TestEdgeCases:
"""Test edge cases and error handling."""
def test_not_enough_points_for_fit(self):
"""Fitting with too few points should raise ValueError."""
np.random.seed(42)
energies_keV = np.array([10.0, 20.0, 30.0])
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
cw = np.ones_like(W)
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
)
extractor = CWExtractor(config, cw, energies_keV)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.05,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.05,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
with pytest.raises(ValueError, match="Not enough data points"):
extractor.fit_CW_parametrization(cw_result, parametrization="quadratic")
def test_not_enough_points_for_gV_gA(self):
"""g_V/g_A fit with too few points should raise ValueError."""
np.random.seed(42)
energies_keV = np.array([10.0, 20.0, 30.0])
W = np.array([T_to_W(e / 1000.0) for e in energies_keV])
cw = np.ones_like(W)
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
)
extractor = CWExtractor(config, cw, energies_keV)
cw_result = CWExtractionResult(
energies_W=W.copy(),
energies_keV=energies_keV.copy(),
cw_values=cw.copy(),
cw_errors=np.ones_like(cw) * 0.05,
flux=cw.copy(),
flux_errors=np.ones_like(cw) * 0.05,
endpoint_W=T_to_W(0.294),
endpoint_keV=294.0,
)
with pytest.raises(ValueError, match="Not enough data points"):
extractor.fit_gV_gA(cw_result, M_F=1.0, M_GT=1.0)
def test_default_uncertainties(self):
"""Default uncertainties should be sqrt(flux)."""
np.random.seed(42)
energies_keV = np.linspace(10, 250, 30)
flux = np.array(
[
100.0,
200.0,
150.0,
180.0,
120.0,
90.0,
80.0,
70.0,
60.0,
50.0,
45.0,
40.0,
35.0,
30.0,
25.0,
20.0,
18.0,
15.0,
12.0,
10.0,
8.0,
7.0,
6.0,
5.0,
4.0,
3.0,
2.0,
1.5,
1.0,
0.5,
]
)
config = SpectrumConfig(
Z_parent=43,
Z_daughter=44,
A_number=99,
endpoint_MeV=0.294,
)
extractor = CWExtractor(config, flux, energies_keV)
expected_unc = np.sqrt(np.maximum(flux, 1.0))
np.testing.assert_allclose(extractor.flux_errors, expected_unc)
tests/test_fitter.py
+448
View File
@@ -0,0 +1,448 @@
# 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 beta_spectrum 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