feat: v0.4.0 package restructuring — split into 4 modules
Separate monolithic beta_spectrum package into:
beta_spectrum — pure theory (components, spectrum, CLI, visualize)
exp_data — experimental data handling (spectrum, calibration,
fitters, corrections)
fitter — analysis framework (model, fit_engine, extractor,
result)
beta_spectrum.visualize — plotting (BetaSpectrumAnalyzer)
Phase 1 — exp_data module:
ExpSpectrum dataclass with metadata, normalization, serialization
EnergyCalibrator (linear/quadratic calibration)
GaussianFitter + PeakFitter for calibration sources
DeadTimeCorrection, PileUpCorrection, BackgroundSubtractor
43 new tests, ruff + mypy clean
Phase 2 — fitter module:
SpectrumModel wraps BetaSpectrum + DetectorResponse with convolution
SpectrumFitter with χ² minimization, covariance, profile likelihood,
confidence intervals, robust loss functions
CWExtractor + GVAExtractor for C(W) and g_V/g_A extraction
AnalysisResult + FitSummary for physics-specific results
CurveFitter general-purpose callable model fitter
34 new tests, ruff + mypy clean
Phase 3 — clean up beta_spectrum:
Remove cw_extractor.py and fitter.py (moved to fitter/)
Clean __init__.py to export only theory classes
Add backward-compat methods to CWExtractor
Fix BetaSpectrumAnalyzer._get_branch_spectrum returning BetaSpectrum
objects instead of numpy arrays (plot fix)
Phase 4 — package structure:
Auto-discovered by setuptools (4 packages)
All imports updated across 26 source files
Black formatting applied
322 tests pass (262 original + 60 new)
ruff + mypy clean on all source files
Bug fixes:
- Fix multi-branch spectrum plotting: universal component keys
(Fermi, Screening, etc.) don't have per-branch prefix; fixed
fallback lookup.
- Fix negative phase space beyond endpoint: BetaSpectrum returns
negative values for W > W0; clamp branch spectra to zero.
This commit is contained in:
@@ -0,0 +1,541 @@
|
||||
"""
|
||||
Multi-parameter fitting engine for beta spectrum analysis.
|
||||
|
||||
Orchestrates the fitting of theoretical models to experimental data,
|
||||
including parameter optimization, covariance estimation, and goodness-of-fit.
|
||||
|
||||
Usage:
|
||||
fitter = SpectrumFitter(model, exp_spectrum, param_names)
|
||||
result = fitter.fit(x0=[norm, bg], bounds=([0, 0], [100, 1e6]))
|
||||
print(result.summary())
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, List, Optional, Sequence, Tuple
|
||||
|
||||
import numpy as np
|
||||
from scipy.optimize import least_squares
|
||||
|
||||
|
||||
@dataclass
|
||||
class FitConfig:
|
||||
"""Configuration for spectrum fitting."""
|
||||
|
||||
method: str = "trf"
|
||||
"""Optimization method: 'trf', 'dogbox', 'lm'."""
|
||||
|
||||
max_nfev: int = 1000
|
||||
"""Maximum function evaluations."""
|
||||
|
||||
x_scale: str | float = "jac"
|
||||
"""Parameter scaling."""
|
||||
|
||||
loss: str = "linear"
|
||||
"""Loss function: 'linear', 'soft_l1', 'huber', 'cauchy'."""
|
||||
|
||||
f_scale: float = 1.0
|
||||
"""Robust loss scale parameter."""
|
||||
|
||||
ftol: float = 1e-10
|
||||
"""Function value tolerance."""
|
||||
|
||||
xtol: float = 1e-10
|
||||
"""Parameter change tolerance."""
|
||||
|
||||
gtol: float = 1e-10
|
||||
"""Gradient tolerance."""
|
||||
|
||||
|
||||
@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."""
|
||||
|
||||
success: bool
|
||||
"""Whether optimization converged."""
|
||||
|
||||
message: str
|
||||
"""Optimization message."""
|
||||
|
||||
_residual_fn: Optional[Callable] = field(default=None, repr=False)
|
||||
"""Residual function for profile likelihood."""
|
||||
|
||||
_x_data: Optional[np.ndarray] = field(default=None, repr=False)
|
||||
"""Independent variable data."""
|
||||
|
||||
_y_data: Optional[np.ndarray] = field(default=None, repr=False)
|
||||
"""Dependent variable data."""
|
||||
|
||||
_uncertainties: Optional[np.ndarray] = field(default=None, repr=False)
|
||||
"""Measurement uncertainties."""
|
||||
|
||||
@property
|
||||
def chi2_per_dof(self) -> float:
|
||||
"""χ² per degree of freedom."""
|
||||
return self.chi2 / max(self.n_free, 1)
|
||||
|
||||
@property
|
||||
def p_value(self) -> float:
|
||||
"""p-value for the fit (survival function of χ²)."""
|
||||
from scipy.stats import chi2
|
||||
|
||||
return float(1.0 - chi2.cdf(self.chi2, self.n_points - self.n_free))
|
||||
|
||||
@property
|
||||
def parameters_with_errors(self) -> dict[str, tuple[float, float]]:
|
||||
"""Parameter values with 1σ uncertainties."""
|
||||
errors = np.sqrt(np.maximum(np.diag(self.covariance), 0.0))
|
||||
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.maximum(np.diag(self.covariance), 0.0))
|
||||
if np.any(diag == 0):
|
||||
return np.zeros_like(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} (ndof = {self.n_free})",
|
||||
f"p-value = {self.p_value:.4f}",
|
||||
"-" * 60,
|
||||
"Parameters:",
|
||||
]
|
||||
|
||||
errors = np.sqrt(np.maximum(np.diag(self.covariance), 0.0))
|
||||
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:20s} = {self.parameters[i]:14.6f} ± {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}':>12s}" for i in range(n))
|
||||
lines.append(header)
|
||||
for i in range(n):
|
||||
name = param_names[i] if param_names and i < len(param_names) else f"p{i}"
|
||||
row = f" {name:8s}" + "".join(
|
||||
f"{corr[i, j]:12.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."""
|
||||
if self._residual_fn is None:
|
||||
raise RuntimeError("Profile likelihood not available.")
|
||||
|
||||
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 + 3 * sigma)
|
||||
|
||||
param_values = np.linspace(param_range[0], param_range[1], n_points)
|
||||
chi2_values = np.zeros(n_points)
|
||||
|
||||
for i, pval in enumerate(param_values):
|
||||
params = list(self.parameters)
|
||||
params[param_index] = pval
|
||||
res = self._residual_fn(np.array(params))
|
||||
chi2_values[i] = np.sum(res**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."""
|
||||
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))
|
||||
|
||||
|
||||
class SpectrumFitter:
|
||||
"""Multi-parameter fitting engine for beta spectrum analysis.
|
||||
|
||||
Wraps a :class:`SpectrumModel` and experimental data to perform
|
||||
χ² minimization with configurable loss functions and bounds.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : SpectrumModel
|
||||
Theoretical model (with optional detector convolution).
|
||||
exp_energies : np.ndarray
|
||||
Experimental energy grid (keV).
|
||||
exp_counts : np.ndarray
|
||||
Experimental counts per channel.
|
||||
exp_errors : np.ndarray, optional
|
||||
Experimental uncertainties. If None, assumed to be √counts.
|
||||
config : FitConfig, optional
|
||||
Fitting configuration.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: Any,
|
||||
exp_energies: np.ndarray,
|
||||
exp_counts: np.ndarray,
|
||||
exp_errors: Optional[np.ndarray] = None,
|
||||
config: Optional[FitConfig] = None,
|
||||
):
|
||||
self.model = model
|
||||
self.exp_energies = np.asarray(exp_energies, dtype=np.float64)
|
||||
self.exp_counts = np.asarray(exp_counts, dtype=np.float64)
|
||||
|
||||
if exp_errors is not None:
|
||||
self.exp_errors = np.asarray(exp_errors, dtype=np.float64)
|
||||
else:
|
||||
self.exp_errors = np.sqrt(np.maximum(self.exp_counts, 1.0))
|
||||
|
||||
# Filter out zero-uncertainty points
|
||||
valid = self.exp_errors > 0
|
||||
self.exp_energies = self.exp_energies[valid]
|
||||
self.exp_counts = self.exp_counts[valid]
|
||||
self.exp_errors = self.exp_errors[valid]
|
||||
|
||||
self.config = config or FitConfig()
|
||||
self._n_params: Optional[int] = None
|
||||
self._last_result: Optional[FitResult] = None
|
||||
|
||||
def _residuals(self, params: np.ndarray) -> np.ndarray:
|
||||
"""Compute weighted residuals."""
|
||||
model_values = self.model.evaluate(self.exp_energies, params[0], params[1])
|
||||
return (self.exp_counts - model_values) / self.exp_errors
|
||||
|
||||
def _chi2(self, params: np.ndarray) -> float:
|
||||
"""Compute χ²."""
|
||||
model_values = self.model.evaluate(self.exp_energies, params[0], params[1])
|
||||
return float(np.sum(((self.exp_counts - model_values) / self.exp_errors) ** 2))
|
||||
|
||||
def fit(
|
||||
self,
|
||||
x0: Optional[Sequence[float]] = None,
|
||||
bounds: Optional[Tuple[Sequence[float], Sequence[float]]] = None,
|
||||
param_names: Optional[List[str]] = None,
|
||||
) -> FitResult:
|
||||
"""Perform χ² minimization.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x0 : sequence of float, optional
|
||||
Initial parameter values. Auto-estimated if None.
|
||||
bounds : (lower, upper), optional
|
||||
Parameter bounds.
|
||||
param_names : list of str, optional
|
||||
Parameter names for summary output.
|
||||
|
||||
Returns
|
||||
-------
|
||||
FitResult
|
||||
"""
|
||||
if x0 is None:
|
||||
x0 = [1.0, 0.0]
|
||||
|
||||
x0_array = np.asarray(x0, dtype=np.float64)
|
||||
self._n_params = len(x0_array)
|
||||
|
||||
bounds_array: Optional[Tuple[np.ndarray, np.ndarray]] = None
|
||||
if bounds is not None:
|
||||
bounds_array = (
|
||||
np.asarray(bounds[0], dtype=np.float64),
|
||||
np.asarray(bounds[1], dtype=np.float64),
|
||||
)
|
||||
|
||||
ls_kwargs: dict[str, Any] = {
|
||||
"fun": self._residuals,
|
||||
"x0": x0_array,
|
||||
"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_array
|
||||
|
||||
result = least_squares(**ls_kwargs)
|
||||
|
||||
# Compute covariance
|
||||
cov = self._compute_covariance(result)
|
||||
|
||||
model_values = self.model.evaluate(self.exp_energies, *result.x)
|
||||
residuals = self.exp_counts - model_values
|
||||
|
||||
fit_result = FitResult(
|
||||
parameters=result.x,
|
||||
covariance=cov,
|
||||
chi2=self._chi2(result.x),
|
||||
n_points=len(self.exp_energies),
|
||||
n_free=self._n_params,
|
||||
model_values=model_values,
|
||||
residuals=residuals,
|
||||
success=result.success,
|
||||
message=result.message,
|
||||
_residual_fn=self._residuals,
|
||||
_x_data=self.exp_energies,
|
||||
_y_data=self.exp_counts,
|
||||
_uncertainties=self.exp_errors,
|
||||
)
|
||||
|
||||
self._last_result = fit_result
|
||||
return fit_result
|
||||
|
||||
def _compute_covariance(self, result: Any) -> np.ndarray:
|
||||
"""Estimate parameter covariance from Jacobian."""
|
||||
n = len(result.x)
|
||||
cov = np.zeros((n, n))
|
||||
|
||||
if result.jac is not None:
|
||||
jtj = result.jac.T @ result.jac
|
||||
if np.linalg.cond(jtj) < 1e12:
|
||||
cov = np.linalg.inv(jtj)
|
||||
chi2 = self._chi2(result.x)
|
||||
ndof = max(len(self.exp_energies) - n, 1)
|
||||
cov = cov * chi2 / ndof
|
||||
else:
|
||||
cov = np.linalg.pinv(jtj)
|
||||
chi2 = self._chi2(result.x)
|
||||
ndof = max(len(self.exp_energies) - n, 1)
|
||||
cov = cov * chi2 / ndof
|
||||
|
||||
return cov
|
||||
|
||||
@property
|
||||
def parameters(self) -> np.ndarray:
|
||||
"""Current best-fit parameters."""
|
||||
if self._last_result is None:
|
||||
raise RuntimeError("No fit result. Call fit() first.")
|
||||
return self._last_result.parameters
|
||||
|
||||
@property
|
||||
def covariance(self) -> np.ndarray:
|
||||
"""Current covariance matrix."""
|
||||
if self._last_result is None:
|
||||
raise RuntimeError("No fit result. Call fit() first.")
|
||||
return self._last_result.covariance
|
||||
|
||||
|
||||
class CurveFitter:
|
||||
"""General-purpose χ² curve fitter for arbitrary callable models.
|
||||
|
||||
This is a thin wrapper around scipy's least_squares that provides
|
||||
the same API as the original beta_spectrum.fitter.CurveFitter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : callable
|
||||
Function f(x, *params) that computes model values.
|
||||
x_data : np.ndarray
|
||||
Independent variable data.
|
||||
y_data : np.ndarray
|
||||
Dependent variable data.
|
||||
uncertainties : np.ndarray, optional
|
||||
Measurement uncertainties. If None, assumed to be 1.0.
|
||||
config : FitConfig, optional
|
||||
Fitting configuration.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> def linear(x, a, b):
|
||||
... return a + b * x
|
||||
>>> fitter = CurveFitter(linear, x, y, uncertainties=unc)
|
||||
>>> result = fitter.fit([0.0, 0.0])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: Callable,
|
||||
x_data: np.ndarray,
|
||||
y_data: np.ndarray,
|
||||
uncertainties: Optional[np.ndarray] = None,
|
||||
config: Optional[FitConfig] = None,
|
||||
):
|
||||
self.model = model
|
||||
self.x_data = np.asarray(x_data, dtype=np.float64)
|
||||
self.y_data = np.asarray(y_data, dtype=np.float64)
|
||||
|
||||
if uncertainties is not None:
|
||||
self.uncertainties = np.asarray(uncertainties, dtype=np.float64)
|
||||
else:
|
||||
self.uncertainties = np.ones_like(self.y_data)
|
||||
|
||||
# Filter out zero/negative uncertainty points
|
||||
valid = self.uncertainties > 0
|
||||
self.x_data = self.x_data[valid]
|
||||
self.y_data = self.y_data[valid]
|
||||
self.uncertainties = self.uncertainties[valid]
|
||||
|
||||
self.config = config or FitConfig()
|
||||
self._last_result: Optional[FitResult] = None
|
||||
self._n_params: Optional[int] = None
|
||||
|
||||
def _residuals(self, params: np.ndarray) -> np.ndarray:
|
||||
"""Compute weighted residuals."""
|
||||
model_values = self.model(self.x_data, *params)
|
||||
return (self.y_data - model_values) / self.uncertainties
|
||||
|
||||
def _chi2(self, params: np.ndarray) -> float:
|
||||
"""Compute χ²."""
|
||||
model_values = self.model(self.x_data, *params)
|
||||
return float(np.sum(((self.y_data - model_values) / self.uncertainties) ** 2))
|
||||
|
||||
def fit(
|
||||
self,
|
||||
x0: Sequence[float],
|
||||
bounds: Optional[Tuple[Sequence[float], Sequence[float]]] = None,
|
||||
) -> FitResult:
|
||||
"""Perform χ² minimization.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x0 : sequence of float
|
||||
Initial parameter values.
|
||||
bounds : (lower, upper), optional
|
||||
Parameter bounds.
|
||||
|
||||
Returns
|
||||
-------
|
||||
FitResult
|
||||
"""
|
||||
x0_array = np.asarray(x0, dtype=np.float64)
|
||||
self._n_params = len(x0_array)
|
||||
|
||||
ls_kwargs: dict[str, Any] = {
|
||||
"fun": self._residuals,
|
||||
"x0": x0_array,
|
||||
"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"] = (
|
||||
np.asarray(bounds[0], dtype=np.float64),
|
||||
np.asarray(bounds[1], dtype=np.float64),
|
||||
)
|
||||
|
||||
result = least_squares(**ls_kwargs)
|
||||
|
||||
# Compute covariance
|
||||
n = len(result.x)
|
||||
cov = np.zeros((n, n))
|
||||
if result.jac is not None:
|
||||
jtj = result.jac.T @ result.jac
|
||||
if np.linalg.cond(jtj) < 1e12:
|
||||
cov = np.linalg.inv(jtj)
|
||||
else:
|
||||
cov = np.linalg.pinv(jtj)
|
||||
chi2 = self._chi2(result.x)
|
||||
ndof = max(len(self.x_data) - n, 1)
|
||||
cov = cov * chi2 / ndof
|
||||
|
||||
model_values = self.model(self.x_data, *result.x)
|
||||
|
||||
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=self.y_data - model_values,
|
||||
success=result.success,
|
||||
message=result.message,
|
||||
_residual_fn=self._residuals,
|
||||
_x_data=self.x_data,
|
||||
_y_data=self.y_data,
|
||||
_uncertainties=self.uncertainties,
|
||||
)
|
||||
|
||||
self._last_result = fit_result
|
||||
return fit_result
|
||||
|
||||
@property
|
||||
def parameters(self) -> np.ndarray:
|
||||
"""Current best-fit parameters."""
|
||||
if self._last_result is None:
|
||||
raise RuntimeError("No fit result. Call fit() first.")
|
||||
return self._last_result.parameters
|
||||
|
||||
@property
|
||||
def covariance(self) -> np.ndarray:
|
||||
"""Current covariance matrix."""
|
||||
if self._last_result is None:
|
||||
raise RuntimeError("No fit result. Call fit() first.")
|
||||
return self._last_result.covariance
|
||||
Reference in New Issue
Block a user