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
beta_spectrum/__init__.py
+9 -1
View File
@@ -12,6 +12,8 @@ from .components.radiative import RadiativeCorrection
from .components.detector_response import DetectorResponse from .components.detector_response import DetectorResponse
from .spectrum import BetaSpectrum, SpectrumConfig, BetaSpectrumAnalyzer from .spectrum import BetaSpectrum, SpectrumConfig, BetaSpectrumAnalyzer
from .fitter import CurveFitter, FitConfig, FitResult
from .cw_extractor import CWExtractor, CWExtractionResult, GVAExtractionResult
__all__ = [ __all__ = [
"ALPHA", "ALPHA",
@@ -32,6 +34,12 @@ __all__ = [
"BetaSpectrum", "BetaSpectrum",
"SpectrumConfig", "SpectrumConfig",
"BetaSpectrumAnalyzer", "BetaSpectrumAnalyzer",
"CurveFitter",
"FitConfig",
"FitResult",
"CWExtractor",
"CWExtractionResult",
"GVAExtractionResult",
] ]
__version__ = "0.2.0" __version__ = "0.3.0"
beta_spectrum/cw_extractor.py
+697
View File
@@ -0,0 +1,697 @@
# beta_spectrum/cw_extractor.py
"""
C(W) shape factor extraction from beta spectra.
Provides methods to:
- Extract C(W) from measured beta spectra
- Perform Kurie plot analysis for endpoint determination
- Fit C(W) parametrization to extract g_V, g_A
- Compare extracted C(W) with theoretical predictions
The beta decay spectrum is:
dΓ/dE = C₀ · p·E·(W₀-E)² · F(Z,E) · C(W) · [1 + δ(E)]
where:
p, E = momentum, total energy of electron
W₀ = endpoint total energy in m_e units
F(Z,E) = Fermi function
C(W) = shape factor (what we extract)
δ(E) = radiative + screening + other corrections
For allowed transitions with C(W) = 1, the Kurie plot
K(E) = sqrt(dΓ/dE / (p·E·F(Z,E)·[1+δ(E)]))
is linear in E with slope -(W₀-E₀)/E₀.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import numpy as np
from .fitter import CurveFitter, FitResult
from .spectrum import SpectrumConfig
from .utils import T_to_W, W_to_T, momentum
@dataclass
class CWExtractionResult:
"""Container for C(W) extraction results."""
energies_W: np.ndarray
"""Total energy grid in m_e units."""
energies_keV: np.ndarray
"""Kinetic energy grid in keV."""
cw_values: np.ndarray
"""Extracted C(W) values."""
cw_errors: np.ndarray
"""Uncertainties on C(W) values."""
flux: np.ndarray
"""Measured flux (counts per energy bin)."""
flux_errors: np.ndarray
"""Uncertainties on flux."""
endpoint_W: float
"""Endpoint energy W₀ in m_e units."""
endpoint_keV: float
"""Endpoint kinetic energy in keV."""
fit_result: Optional[FitResult] = None
"""Fit result from Kurie plot endpoint determination."""
def kurie_plot(self) -> Tuple[np.ndarray, np.ndarray]:
"""
Compute Kurie plot values.
For allowed transitions (C(W)=1), K(E) should be linear.
Deviations from linearity indicate shape factor effects.
"""
# K(E) = sqrt(flux / (p * E * F * phase_space * corrections))
# We compute this from the extracted C(W) values
# K(E) ∝ (W0 - E) / sqrt(C(W))
kurie = np.sqrt(self.flux) / (
momentum(self.energies_W)
* np.sqrt(self.energies_W)
* (self.endpoint_W - self.energies_W)
)
kurie = kurie / np.max(kurie) # Normalize
return self.energies_keV, kurie
@dataclass
class GVAExtractionResult:
"""Container for g_V, g_A extraction results."""
g_V: float
"""Extracted vector coupling constant."""
g_V_error: float
"""Uncertainty on g_V."""
g_A: float
"""Extracted axial-vector coupling constant."""
g_A_error: float
"""Uncertainty on g_A."""
fit_result: FitResult
"""Fit result with full covariance matrix."""
cw_data: np.ndarray
"""C(W) data used in fit."""
cw_errors: np.ndarray
"""C(W) uncertainties."""
energy_W: np.ndarray
"""Energy grid for C(W) data."""
def summary(self) -> str:
"""Human-readable summary of g_V, g_A extraction."""
lines = [
"=" * 60,
"g_V and g_A Extraction Results",
"=" * 60,
f"g_V = {self.g_V:.4f} ± {self.g_V_error:.4f}",
f"g_A = {self.g_A:.4f} ± {self.g_A_error:.4f}",
f"χ²/ndf = {self.fit_result.chi2_per_dof:.3f}",
f"p-value = {self.fit_result.p_value:.4f}",
"-" * 60,
"Comparison with standard values:",
f" g_V / g_A ratio = {self.g_V / self.g_A:.4f} "
f"(± {self._ratio_error():.4f})",
" Standard model: g_V/g_A ≈ 1.0 for pure Fermi",
"=" * 60,
]
return "\n".join(lines)
def _ratio_error(self) -> float:
"""Error on g_V/g_A ratio using error propagation."""
r = self.g_V / self.g_A
cov = self.fit_result.covariance
dg_V = np.sqrt(cov[0, 0])
dg_A = np.sqrt(cov[1, 1])
dcorr = cov[0, 1]
return abs(r) * np.sqrt(
(dg_V / self.g_V) ** 2
+ (dg_A / self.g_A) ** 2
- 2 * dcorr / (self.g_V * self.g_A)
)
class CWExtractor:
"""
Extract C(W) shape factor from measured beta spectra.
The extraction works by dividing the measured spectrum by all known
theoretical factors:
C(W) = (dΓ/dE)_measured / [C₀ · p·E·(W₀-E)² · F(Z,E) · (1+δ)]
where C₀ is an overall normalization constant.
For allowed transitions, C(W) should be approximately constant (=1).
For forbidden transitions, C(W) has energy dependence that encodes
nuclear structure information.
"""
def __init__(
self,
config: SpectrumConfig,
flux: np.ndarray,
energies_keV: np.ndarray,
flux_errors: Optional[np.ndarray] = None,
detector_response: Optional["DetectorResponse"] = None, # noqa: F821
):
"""
Initialize extractor.
Parameters
----------
config : SpectrumConfig
Decay configuration (Z, A, endpoint, corrections).
flux : np.ndarray
Measured spectrum (counts per energy bin).
energies_keV : np.ndarray
Kinetic energy grid in keV.
flux_errors : np.ndarray, optional
Uncertainties on flux. If None, assumed to be sqrt(flux).
detector_response : DetectorResponse, optional
Detector response for deconvolution (optional).
"""
self.config = config
self.flux = np.asarray(flux, dtype=np.float64)
self.energies_keV = np.asarray(energies_keV, dtype=np.float64)
self.energies_W = np.array([T_to_W(e / 1000.0) for e in self.energies_keV])
self.detector_response = detector_response
if flux_errors is not None:
self.flux_errors = np.asarray(flux_errors, dtype=np.float64)
else:
# Poisson statistics
self.flux_errors = np.sqrt(np.maximum(self.flux, 1.0))
def _theoretical_factors(self, W: np.ndarray) -> Dict[str, np.ndarray]:
"""Compute all known theoretical factors."""
factors = {}
# Phase space: p * E
p = momentum(W)
factors["phase_space"] = p * W
# Fermi function
from .components.fermi import FermiFunction
fermi = FermiFunction(self.config.Z_parent, self.config.A_number)
factors["fermi"] = fermi(W)
# Screening
if self.config.use_screening:
from .components.screening import ScreeningCorrection
screening = ScreeningCorrection(
FermiFunction(self.config.Z_parent, self.config.A_number)
)
factors["screening"] = screening(W)
else:
factors["screening"] = np.ones_like(W)
# Exchange
if self.config.use_exchange:
from .components.exchange import ExchangeCorrection
exchange = ExchangeCorrection(Z=self.config.Z_parent)
factors["exchange"] = exchange(W)
else:
factors["exchange"] = np.ones_like(W)
# Finite size
if self.config.use_finite_size:
from .components.finite_size import (
FiniteSizeL0,
ChargeDistributionU,
)
fs_l0 = FiniteSizeL0(Z=self.config.Z_daughter, A=self.config.A_number)
fs_u = ChargeDistributionU(Z=self.config.Z_daughter, A=self.config.A_number)
factors["finite_size_L0"] = fs_l0(W)
factors["charge_distribution"] = fs_u(W)
else:
factors["finite_size_L0"] = np.ones_like(W)
factors["charge_distribution"] = np.ones_like(W)
# Radiative corrections
if self.config.use_radiative:
from .components.radiative import RadiativeCorrection
W0 = T_to_W(self.config.endpoint_MeV)
rad = RadiativeCorrection(
W0=W0,
Z=self.config.Z_parent,
A=self.config.A_number,
use_endpoint_resummation=True,
delta_cut=1e-3,
)
factors["radiative"] = rad(W)
else:
factors["radiative"] = np.ones_like(W)
return factors
def extract_CW(
self,
endpoint_keV: Optional[float] = None,
normalize_method: str = "endpoint",
) -> CWExtractionResult:
"""
Extract C(W) from measured spectrum.
Parameters
----------
endpoint_keV : float, optional
Endpoint energy in keV. If None, determined from Kurie plot.
normalize_method : str
Method for normalization: 'endpoint' or 'fit'.
Returns
-------
CWExtractionResult
Extracted C(W) values with uncertainties.
"""
if endpoint_keV is not None:
W0 = T_to_W(endpoint_keV / 1000.0)
else:
W0 = self._determine_endpoint()
# Compute theoretical factors (excluding C(W))
factors = self._theoretical_factors(self.energies_W)
# Product of all known factors
product = factors["phase_space"] * factors["fermi"]
for key in [
"screening",
"exchange",
"finite_size_L0",
"charge_distribution",
"radiative",
]:
product *= factors[key]
# Endpoint factor: (W0 - W)^2
endpoint_factor = (W0 - self.energies_W) ** 2
# Overall normalization
if normalize_method == "endpoint":
# Use flux at lowest energy where C(W) ≈ 1 for allowed decay
low_E_mask = self.energies_W < W0 - 0.02
if np.any(low_E_mask):
C0 = np.median(
self.flux[low_E_mask]
/ (product[low_E_mask] * endpoint_factor[low_E_mask])
)
else:
C0 = 1.0
else:
# Fit normalization
C0 = self._fit_normalization(W0, product, endpoint_factor)
# Extract C(W)
cw = self.flux / (C0 * product * endpoint_factor)
# Propagate uncertainties
cw_errors = np.maximum(np.abs(cw), 1e-10) * np.sqrt(
(self.flux_errors / np.maximum(self.flux, 1e-30)) ** 2
+ 0.01**2 # ~1% systematic on theoretical factors
)
# Clip unphysical values
cw = np.clip(cw, 0.01, 10.0)
cw_errors = np.clip(cw_errors, 0.0, np.inf)
return CWExtractionResult(
energies_W=self.energies_W.copy(),
energies_keV=self.energies_keV.copy(),
cw_values=cw,
cw_errors=cw_errors,
flux=self.flux.copy(),
flux_errors=self.flux_errors.copy(),
endpoint_W=W0,
endpoint_keV=W_to_T(W0) * 1000.0,
fit_result=None,
)
def kurie_analysis(
self,
endpoint_keV: Optional[float] = None,
) -> Tuple[CWExtractionResult, FitResult]:
"""
Perform Kurie plot analysis for endpoint determination.
For allowed transitions, the Kurie plot K(E) = sqrt(Φ/(p·E·F))
should be linear in E. The endpoint is where K(E) → 0.
Parameters
----------
endpoint_keV : float, optional
If provided, use this endpoint. Otherwise fit for it.
Returns
-------
cw_result : CWExtractionResult
Extracted C(W) values.
fit_result : FitResult
Kurie plot fit result.
"""
# First extract C(W) with a rough endpoint
cw_result = self.extract_CW(
endpoint_keV=endpoint_keV,
normalize_method="fit",
)
if endpoint_keV is None:
# Fit Kurie plot to determine endpoint
fit_result = self._fit_kurie_endpoint(cw_result)
else:
# Dummy fit result
fit_result = FitResult(
parameters=np.array([cw_result.endpoint_keV]),
covariance=np.eye(1) * 1.0,
chi2=0.0,
n_points=len(cw_result.energies_keV),
n_free=1,
model_values=np.zeros_like(cw_result.energies_keV),
residuals=np.zeros_like(cw_result.energies_keV),
success=True,
message="Endpoint fixed, no fit needed",
)
cw_result.fit_result = fit_result
return cw_result, fit_result
def fit_CW_parametrization(
self,
cw_result: CWExtractionResult,
parametrization: str = "linear",
bounds: Optional[Tuple] = None,
) -> FitResult:
"""
Fit C(W) to a parametrized form.
Supported parametrizations:
- "constant": C(W) = a₀ (allowed transition)
- "linear": C(W) = a₀ + a₁·(W - 1) (first-order correction)
- "quadratic": C(W) = a₀ + a₁·(W - 1) + a₂·(W - 1)²
- "gV_gA": C(W) = g_V²·|M_F|² + g_A²·|M_GT|² + interference terms
Parameters
----------
cw_result : CWExtractionResult
Extracted C(W) data.
parametrization : str
Parametrization form.
bounds : tuple, optional
Parameter bounds (lower, upper).
Returns
-------
FitResult
Fit results.
"""
# Define parametrization models
models = {
"constant": self._model_constant,
"linear": self._model_linear,
"quadratic": self._model_quadratic,
}
if parametrization not in models:
raise ValueError(
f"Unknown parametrization: {parametrization}. "
f"Use one of: {list(models.keys())}"
)
n_params = {
"constant": 1,
"linear": 2,
"quadratic": 3,
}[parametrization]
# Initial guesses
x0 = self._initial_guesses(parametrization, cw_result)
# Filter out points near endpoint (large uncertainties)
mask = cw_result.energies_keV < cw_result.endpoint_keV * 0.95
energies = cw_result.energies_W[mask]
cw = cw_result.cw_values[mask]
cw_err = cw_result.cw_errors[mask]
if len(energies) < n_params + 2:
raise ValueError(
f"Not enough data points for {parametrization} fit. "
f"Need at least {n_params + 2}, got {len(energies)}"
)
fitter = CurveFitter(
model=models[parametrization],
x_data=energies,
y_data=cw,
uncertainties=cw_err,
)
return fitter.fit(x0, bounds, param_names=[f"a{i}" for i in range(n_params)])
def fit_gV_gA(
self,
cw_result: CWExtractionResult,
M_F: float = 1.0,
M_GT: float = 1.0,
interference: float = 0.0,
bounds: Optional[Tuple] = None,
) -> GVAExtractionResult:
"""
Extract g_V and g_A from C(W) data.
For mixed transitions:
C(W) = g_V²·|M_F|² + g_A²·|M_GT|² + b·g_V·g_A·interference
where M_F and M_GT are Fermi and Gamow-Teller matrix elements.
Parameters
----------
cw_result : CWExtractionResult
Extracted C(W) data.
M_F : float
Fermi matrix element (default: 1.0 for pure Fermi).
M_GT : float
Gamow-Teller matrix element (default: 1.0 for pure GT).
interference : float
Interference term coefficient (b coefficient).
bounds : tuple, optional
Parameter bounds (g_V_low, g_A_low, g_V_high, g_A_high).
Returns
-------
GVAExtractionResult
Extracted g_V, g_A with uncertainties.
"""
# Filter points near endpoint
mask = cw_result.energies_keV < cw_result.endpoint_keV * 0.95
energies = cw_result.energies_W[mask]
cw = cw_result.cw_values[mask]
cw_err = cw_result.cw_errors[mask]
if len(energies) < 4:
raise ValueError(
f"Not enough data points for g_V/g_A fit. "
f"Need at least 4, got {len(energies)}"
)
# g_V, g_A model
def gVgA_model(W, g_V, g_A):
return g_V**2 * M_F**2 + g_A**2 * M_GT**2 + interference * g_V * g_A
# Initial guesses: use Paulsen et al. values as starting point
x0 = [0.376, 0.574]
if bounds is None:
bounds = (
[0.0, 0.0],
[2.0, 2.0],
)
fitter = CurveFitter(
model=gVgA_model,
x_data=energies,
y_data=cw,
uncertainties=cw_err,
)
fit_result = fitter.fit(
x0,
bounds,
param_names=["g_V", "g_A"],
)
return GVAExtractionResult(
g_V=float(fit_result.parameters[0]),
g_V_error=float(np.sqrt(fit_result.covariance[0, 0])),
g_A=float(fit_result.parameters[1]),
g_A_error=float(np.sqrt(fit_result.covariance[1, 1])),
fit_result=fit_result,
cw_data=cw,
cw_errors=cw_err,
energy_W=energies,
)
def _determine_endpoint(self) -> float:
"""
Determine endpoint energy from Kurie plot linearity.
Scans endpoint values and finds the one that gives the most
linear Kurie plot (highest R² for linear fit).
"""
# Scan endpoint in keV
E_min = self.energies_keV[-1] - 50 # At least 50 keV below max
E_max = self.energies_keV[-1] + 10 # Allow slight overshoot
best_R2 = -np.inf
best_endpoint = E_max
for E_trial in np.linspace(E_min, E_max, 30):
W0 = T_to_W(E_trial / 1000.0)
# Compute Kurie values
mask = self.energies_W < W0 - 0.01
if not np.any(mask):
continue
W = self.energies_W[mask]
flux = self.flux[mask]
# K(E) = sqrt(flux / (p*E*F))
p = momentum(W)
from .components.fermi import FermiFunction
F = FermiFunction()(W, self.config.Z_parent)
kurie = np.sqrt(flux / (p * W * F))
# Linear fit: K(E) = a + b*E
E_keV = self.energies_keV[mask]
try:
coeffs = np.polyfit(E_keV, kurie, 1)
# R²
kurie_fit = np.polyval(coeffs, E_keV)
SS_res = np.sum((kurie - kurie_fit) ** 2)
SS_tot = np.sum((kurie - np.mean(kurie)) ** 2)
R2 = 1 - SS_res / SS_tot if SS_tot > 0 else 0
if R2 > best_R2:
best_R2 = R2
best_endpoint = E_trial
except np.linalg.LinAlgError:
continue
return T_to_W(best_endpoint / 1000.0)
def _fit_kurie_endpoint(self, cw_result: CWExtractionResult) -> FitResult:
"""Fit Kurie plot to determine endpoint energy."""
W0 = cw_result.endpoint_W
# Kurie plot: K(E) = sqrt(flux / (p*E*F))
mask = self.energies_W < W0 - 0.01
W = self.energies_W[mask]
flux = self.flux[mask]
flux_err = self.flux_errors[mask]
p = momentum(W)
from .components.fermi import FermiFunction
F = FermiFunction()(W, self.config.Z_parent)
kurie = np.sqrt(flux / (p * W * F))
E_keV = self.energies_keV[mask]
# K(E) = a + b*(E - E0), where K(E0) = 0
# So a = -b*(E0 - E0) = 0 at the true endpoint
# Linear fit: K(E) = a + b*E
def kurie_model(E, a, b):
return a + b * E
fitter = CurveFitter(
model=kurie_model,
x_data=E_keV,
y_data=kurie,
uncertainties=flux_err / (2 * kurie * p * W * F),
)
return fitter.fit(
x0=[np.mean(kurie), -np.mean(kurie) / E_keV[-1]],
param_names=["a", "b"],
)
def _fit_normalization(
self, W0: float, product: np.ndarray, endpoint_factor: np.ndarray
) -> float:
"""Fit overall normalization C₀."""
mask = self.energies_W < W0 - 0.02
if not np.any(mask):
return 1.0
def norm_model(E, C0):
return C0 * product[mask] * endpoint_factor[mask]
fitter = CurveFitter(
model=norm_model,
x_data=self.energies_W[mask],
y_data=self.flux[mask],
uncertainties=self.flux_errors[mask],
)
result = fitter.fit(x0=[1.0])
return float(result.parameters[0])
def _model_constant(self, W: np.ndarray, a0: float) -> np.ndarray:
"""C(W) = a₀ (constant)."""
return np.full_like(W, a0)
def _model_linear(self, W: np.ndarray, a0: float, a1: float) -> np.ndarray:
"""C(W) = a₀ + a₁·(W - 1)."""
return a0 + a1 * (W - 1.0)
def _model_quadratic(
self, W: np.ndarray, a0: float, a1: float, a2: float
) -> np.ndarray:
"""C(W) = a₀ + a₁·(W - 1) + a₂·(W - 1)²."""
return a0 + a1 * (W - 1.0) + a2 * (W - 1.0) ** 2
def _initial_guesses(
self, parametrization: str, cw_result: CWExtractionResult
) -> List[float]:
"""Generate initial guesses for fit parameters."""
# Average C(W) in mid-energy region
mid_mask = (cw_result.energies_keV > cw_result.endpoint_keV * 0.2) & (
cw_result.energies_keV < cw_result.endpoint_keV * 0.7
)
if np.any(mid_mask):
avg_cw = float(np.mean(cw_result.cw_values[mid_mask]))
else:
avg_cw = 1.0
if parametrization == "constant":
return [avg_cw]
elif parametrization == "linear":
return [avg_cw, 0.0]
elif parametrization == "quadratic":
return [avg_cw, 0.0, 0.0]
else:
return [avg_cw] * 3
beta_spectrum/fitter.py
+552
View File
@@ -0,0 +1,552 @@
# beta_spectrum/fitter.py
"""
χ² minimization framework for beta spectrum fitting.
Provides infrastructure for fitting theoretical spectra to experimental data,
including parameter optimization, covariance estimation, and goodness-of-fit
metrics.
Usage:
fitter = CurveFitter(model, data, uncertainties)
result = fitter.fit(x0=[...], bounds=[...])
print(f"Best fit: {result.parameters}")
print(f"Covariance: {result.covariance}")
print(f"χ²/ndf: {result.chi2_per_dof:.3f}")
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
import numpy as np
from scipy.optimize import least_squares
@dataclass
class FitResult:
"""Container for fit results."""
parameters: np.ndarray
"""Best-fit parameter values."""
covariance: np.ndarray
"""Parameter covariance matrix."""
chi2: float
"""Minimum χ² value."""
n_points: int
"""Number of data points."""
n_free: int
"""Number of free parameters."""
model_values: np.ndarray
"""Model values at best-fit parameters."""
residuals: np.ndarray
"""Data minus model at best-fit parameters."""
success: bool
"""Whether optimization converged."""
message: str
"""Optimization success/failure message."""
_model: Optional[Callable] = field(default=None, repr=False)
"""Model function (for profile likelihood)."""
_x_data: Optional[np.ndarray] = field(default=None, repr=False)
"""Independent variable data (for profile likelihood)."""
_y_data: Optional[np.ndarray] = field(default=None, repr=False)
"""Dependent variable data (for profile likelihood)."""
_uncertainties: Optional[np.ndarray] = field(default=None, repr=False)
"""Measurement uncertainties (for profile likelihood)."""
@property
def chi2_per_dof(self) -> float:
"""χ² per degree of freedom."""
return self.chi2 / self.n_free
@property
def p_value(self) -> float:
"""
p-value for the fit (survival function of χ² distribution).
High p-value (>0.05) suggests good fit; low p-value (<0.01) suggests
poor fit or underestimated uncertainties.
"""
from scipy.stats import chi2
return 1.0 - chi2.cdf(self.chi2, self.n_points - self.n_free)
@property
def parameters_with_errors(self) -> Dict[str, Tuple[float, float]]:
"""
Return parameter names with values and 1σ uncertainties.
Returns dict mapping parameter names to (value, error) tuples.
"""
errors = np.sqrt(np.diag(self.covariance))
return {
f"p{i}": (float(self.parameters[i]), float(errors[i]))
for i in range(len(self.parameters))
}
@property
def correlation_matrix(self) -> np.ndarray:
"""Parameter correlation matrix."""
diag = np.sqrt(np.diag(self.covariance))
return self.covariance / np.outer(diag, diag)
def summary(self, param_names: Optional[List[str]] = None) -> str:
"""Human-readable fit summary."""
lines = [
"=" * 60,
"Fit Results",
"=" * 60,
f"Converged: {self.success}",
f"Message: {self.message}",
f"χ² = {self.chi2:.4f}",
f"χ²/ndof = {self.chi2_per_dof:.3f} " f"(ndof = {self.n_free})",
f"p-value = {self.p_value:.4f}",
"-" * 60,
"Parameters:",
]
errors = np.sqrt(np.diag(self.covariance))
for i in range(len(self.parameters)):
name = param_names[i] if param_names and i < len(param_names) else f"p{i}"
lines.append(
f" {name:15s} = {self.parameters[i]:12.6f} " f"± {errors[i]:10.6f}"
)
lines.append("-" * 60)
lines.append("Correlation matrix (off-diagonal):")
corr = self.correlation_matrix
n = len(self.parameters)
header = " " + "".join(f"{f'p{i}':>10s}" for i in range(n))
lines.append(header)
for i in range(n):
row = f" p{i!s:<8s}" + "".join(
f"{corr[i, j]:10.3f}" for j in range(n) if j != i
)
lines.append(row)
lines.append("=" * 60)
return "\n".join(lines)
def profile_likelihood(
self,
param_index: int,
n_points: int = 50,
param_range: Optional[Tuple[float, float]] = None,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Compute profile likelihood for a single parameter.
Parameters
----------
param_index : int
Index of parameter to profile over.
n_points : int
Number of points to evaluate.
param_range : (low, high), optional
Range to scan. Uses ±3σ from fit result if None.
Returns
-------
param_values : np.ndarray
Parameter values scanned.
chi2_values : np.ndarray
χ² values at each point.
normalized : np.ndarray
χ² normalized to minimum (Δχ² = 0 at best fit).
"""
if self._model is None or self._x_data is None:
raise RuntimeError(
"Cannot compute profile likelihood: model or data not stored. "
"Use CurveFitter.profile_likelihood() instead."
)
if param_range is None:
sigma = np.sqrt(abs(self.covariance[param_index, param_index]))
center = self.parameters[param_index]
param_range = (center - 3 * sigma, center + sigma * 3)
param_values = np.linspace(param_range[0], param_range[1], n_points)
chi2_values = np.zeros(n_points)
unc = (
self._uncertainties
if self._uncertainties is not None
else np.ones_like(self._x_data)
)
for i, pval in enumerate(param_values):
params = list(self.parameters)
params[param_index] = pval
y_model = self._model(self._x_data, *params)
residuals = (self._y_data - y_model) / unc
chi2_values[i] = np.sum(residuals**2)
min_chi2 = np.min(chi2_values)
normalized = chi2_values - min_chi2
return param_values, chi2_values, normalized
def confidence_interval(
self,
param_index: int,
confidence: float = 0.95,
n_points: int = 100,
) -> Tuple[float, float]:
"""
Compute confidence interval for a parameter using Δχ² method.
Parameters
----------
param_index : int
Index of parameter.
confidence : float
Confidence level (e.g., 0.95 for 95% CI).
n_points : int
Number of points for profile likelihood scan.
Returns
-------
(lower, upper) : float tuple
Confidence interval bounds.
"""
from scipy.stats import chi2
delta_chi2 = chi2.ppf(confidence, df=1)
param_values, chi2_values, normalized = self.profile_likelihood(
param_index, n_points=n_points
)
above = normalized >= delta_chi2
if not np.any(above):
return (
float(self.parameters[param_index]),
float(self.parameters[param_index]),
)
crossings = np.where(np.diff(above.astype(int)))[0]
if len(crossings) >= 2:
lower = param_values[crossings[0]]
upper = param_values[crossings[-1] + 1]
elif len(crossings) == 1:
lower = param_values[crossings[0]]
upper = param_values[-1]
else:
lower = param_values[0]
upper = param_values[-1]
return (float(lower), float(upper))
@dataclass
class FitConfig:
"""Configuration for curve fitting."""
method: str = "trf"
"""Optimization method passed to scipy.optimize.least_squares.
Options: 'trf', 'dogbox', 'lm'."""
max_nfev: int = 1000
"""Maximum number of function evaluations."""
x_scale: Union[str, float] = "jac"
"""Parameter scaling: 'jac', 'diag', or float."""
loss: str = "linear"
"""Loss function: 'linear' (χ²), 'soft_l1', 'huber', 'cauchy'."""
f_scale: float = 1.0
"""Parameter for robust loss functions."""
ftol: float = 1e-8
"""Termination tolerance for function value."""
xtol: float = 1e-8
"""Termination tolerance for parameter change."""
gtol: float = 1e-8
"""Termination tolerance for gradient."""
class CurveFitter:
"""
χ² minimization framework for fitting models to data.
The model function should have signature:
model(x, *params) -> y
where x is the independent variable and params are the fit parameters.
χ² is computed as:
χ² = Σ [(y_data[i] - model(x[i], params)) / σ[i]]²
Parameters are optimized using scipy.optimize.least_squares with
configurable loss function for robust fitting.
"""
def __init__(
self,
model: Callable,
x_data: np.ndarray,
y_data: np.ndarray,
uncertainties: Optional[np.ndarray] = None,
config: Optional[FitConfig] = None,
):
"""
Initialize fitter.
Parameters
----------
model : callable
Model function: model(x, *params) -> y
x_data : np.ndarray
Independent variable data points.
y_data : np.ndarray
Dependent variable data points.
uncertainties : np.ndarray, optional
Measurement uncertainties (σ). If None, assumed to be 1.0.
config : FitConfig, optional
Fitting configuration. Uses defaults if None.
"""
self.x_data = np.asarray(x_data, dtype=np.float64)
self.y_data = np.asarray(y_data, dtype=np.float64)
self.model = model
if uncertainties is not None:
self.uncertainties = np.asarray(uncertainties, dtype=np.float64)
else:
self.uncertainties = np.ones_like(self.y_data)
# Filter out points with zero or negative uncertainties
valid = self.uncertainties > 0
self.x_data = self.x_data[valid]
self.y_data = self.y_data[valid]
self.uncertainties = self.uncertainties[valid]
if config is not None:
self.config = config
else:
self.config = FitConfig()
self._n_params: Optional[int] = None
def _residuals(self, params: np.ndarray) -> np.ndarray:
"""Compute weighted residuals for least_squares."""
y_model = self.model(self.x_data, *params)
return (self.y_data - y_model) / self.uncertainties
def _chi2(self, params: np.ndarray) -> float:
"""Compute χ² for given parameters."""
residuals = self._residuals(params)
return float(np.sum(residuals**2))
def fit(
self,
x0: Sequence[float],
bounds: Optional[Tuple[Sequence[float], Sequence[float]]] = None,
param_names: Optional[List[str]] = None,
) -> FitResult:
"""
Perform χ² minimization.
Parameters
----------
x0 : sequence of float
Initial parameter values.
bounds : (lower, upper), optional
Parameter bounds as (lower_bounds, upper_bounds).
Use -np.inf / np.inf for unbounded parameters.
param_names : list of str, optional
Names for parameters (for summary output).
Returns
-------
FitResult
Fit results including parameters, covariance, χ², etc.
"""
x0 = np.asarray(x0, dtype=np.float64)
self._n_params = len(x0)
if bounds is not None:
lb = np.asarray(bounds[0], dtype=np.float64)
ub = np.asarray(bounds[1], dtype=np.float64)
bounds = (lb, ub)
ls_kwargs = dict(
fun=self._residuals,
x0=x0,
method=self.config.method,
max_nfev=self.config.max_nfev,
x_scale=self.config.x_scale,
loss=self.config.loss,
f_scale=self.config.f_scale,
ftol=self.config.ftol,
xtol=self.config.xtol,
gtol=self.config.gtol,
)
if bounds is not None:
ls_kwargs["bounds"] = bounds
result = least_squares(**ls_kwargs)
# Compute covariance matrix from Jacobian
# Cov = (J^T J)^(-1) * χ² / (n - p)
cov = self._compute_covariance(result)
model_values = self.model(self.x_data, *result.x)
residuals = self.y_data - model_values
fit_result = FitResult(
parameters=result.x,
covariance=cov,
chi2=self._chi2(result.x),
n_points=len(self.x_data),
n_free=self._n_params,
model_values=model_values,
residuals=residuals,
success=result.success,
message=result.message,
_model=self.model,
_x_data=self.x_data,
_y_data=self.y_data,
_uncertainties=self.uncertainties,
)
return fit_result
def _compute_covariance(self, result) -> np.ndarray:
"""Estimate parameter covariance matrix from Jacobian."""
n = len(result.x)
cov = np.zeros((n, n))
if result.jac is not None:
# J^T J approximation
jtj = result.jac.T @ result.jac
if np.linalg.cond(jtj) < 1e12:
cov = np.linalg.inv(jtj)
# Scale by χ²/ndof
chi2 = self._chi2(result.x)
ndof = max(len(self.x_data) - n, 1)
cov *= chi2 / ndof
else:
# Singular matrix — use pseudo-inverse
cov = np.linalg.pinv(jtj)
chi2 = self._chi2(result.x)
ndof = max(len(self.x_data) - n, 1)
cov *= chi2 / ndof
return cov
def residuals_plot(
self,
save_path: Optional[str] = None,
):
"""
Plot data, model, and residuals.
Parameters
----------
save_path : str, optional
Path to save figure.
"""
import matplotlib.pyplot as plt
fig, (ax1, ax2) = plt.subplots(
2,
1,
figsize=(8, 8),
gridspec_kw={"height_ratios": [3, 1]},
sharex=True,
)
# Data + model
ax1.errorbar(
self.x_data,
self.y_data,
yerr=self.uncertainties,
fmt="ko",
markersize=4,
capsize=2,
label="Data",
alpha=0.7,
)
ax1.plot(
self.x_data,
self.model_values,
"r-",
lw=1.5,
label="Model",
)
ax1.set_ylabel("Counts")
ax1.set_title("Data vs. Model")
ax1.legend()
ax1.grid(True, alpha=0.3)
# Residuals
residuals = self.residuals / self.uncertainties
ax2.plot(self.x_data, residuals, "bo", markersize=4, alpha=0.7)
ax2.axhline(y=0, color="k", linestyle="--", lw=0.5)
ax2.set_xlabel("Energy [MeV]")
ax2.set_ylabel("Residuals (σ)")
ax2.set_title(f"Residuals χ²/ndof = {self.chi2_per_dof:.3f}")
ax2.grid(True, alpha=0.3)
plt.tight_layout()
if save_path:
fig.savefig(save_path, dpi=150, bbox_inches="tight")
plt.show()
@property
def parameters(self) -> np.ndarray:
"""Current best-fit parameters (set by fit())."""
if not hasattr(self, "_last_result"):
raise RuntimeError("No fit result available. Call fit() first.")
return self._last_result.parameters
@property
def covariance(self) -> np.ndarray:
"""Current covariance matrix (set by fit())."""
if not hasattr(self, "_last_result"):
raise RuntimeError("No fit result available. Call fit() first.")
return self._last_result.covariance
def fit_and_store(
self,
x0: Sequence[float],
bounds: Optional[Tuple[Sequence[float], Sequence[float]]] = None,
param_names: Optional[List[str]] = None,
) -> FitResult:
"""
Fit and store result for property access.
Parameters
----------
x0 : sequence of float
Initial parameter values.
bounds : (lower, upper), optional
Parameter bounds.
param_names : list of str, optional
Parameter names for summary.
Returns
-------
FitResult
"""
result = self.fit(x0, bounds, param_names)
self._last_result = result
return result
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