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

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

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

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

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

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

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

Bug fixes:
  - Fix multi-branch spectrum plotting: universal component keys
    (Fermi, Screening, etc.) don't have per-branch prefix; fixed
    fallback lookup.
  - Fix negative phase space beyond endpoint: BetaSpectrum returns
    negative values for W > W0; clamp branch spectra to zero.
This commit is contained in:
unzhakov
2026-05-17 18:58:04 +03:00
parent 745bbae4cd
commit 38b44b56e6
29 changed files with 3538 additions and 2010 deletions
Download Patch File Download Diff File Expand all files Collapse all files
beta_spectrum/__init__.py
-8
View File
@@ -13,8 +13,6 @@ from .components.detector_response import DetectorResponse
from .spectrum import BetaSpectrum, SpectrumConfig, BranchConfig
from .visualize import BetaSpectrumAnalyzer
from .fitter import CurveFitter, FitConfig, FitResult
from .cw_extractor import CWExtractor, CWExtractionResult, GVAExtractionResult
from .nuclear_data import (
DecayInfo,
BranchInfo,
@@ -46,12 +44,6 @@ __all__ = [
"SpectrumConfig",
"BetaSpectrumAnalyzer",
"BranchConfig",
"CurveFitter",
"FitConfig",
"FitResult",
"CWExtractor",
"CWExtractionResult",
"GVAExtractionResult",
"DecayInfo",
"BranchInfo",
"get_decay_info_from_paceENSDF",
beta_spectrum/cw_extractor.py
-704
View File
@@ -1,704 +0,0 @@
# mypy: disable-error-code = no-any-return, arg-type, assignment, call-arg, return-value, no-untyped-def, misc
# beta_spectrum/cw_extractor.py
# ruff: noqa: ANN201, ANN001, ANN204
"""
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 TYPE_CHECKING, Dict, List, Optional, Tuple
if TYPE_CHECKING:
from beta_spectrum.components.detector_response import DetectorResponse
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_float = float(T_to_W(self.config.endpoint_MeV))
rad = RadiativeCorrection(
W0=W0_float,
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_float = float(T_to_W(endpoint_keV / 1000.0))
else:
W0_float = float(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_float - 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_float - 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_float, 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_float,
endpoint_keV=float(W_to_T(W0_float)) * 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(Z=self.config.Z_parent, A=self.config.A_number)(W)
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 float(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(Z=self.config.Z_parent, A=self.config.A_number)(W)
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=[float(np.mean(kurie)), float(-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
# mypy: ignore-errors
beta_spectrum/fitter.py
-565
View File
@@ -1,565 +0,0 @@
# 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 Any, 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[..., Any]] = 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 float(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 np.asarray(self.covariance / np.outer(diag, diag), dtype=np.float64)
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[..., Any],
x_data: np.ndarray,
y_data: np.ndarray,
uncertainties: Optional[np.ndarray] = None,
config: Optional[FitConfig] = None,
) -> 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.ndarray = np.asarray(x_data, dtype=np.float64)
self.y_data: np.ndarray = np.asarray(y_data, dtype=np.float64)
self.model: Callable[..., Any] = 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
self._last_result: Optional[FitResult] = None
def _residuals(self, params: np.ndarray) -> np.ndarray:
"""Compute weighted residuals for least_squares."""
y_model = self.model(self.x_data, *params)
return np.asarray(
(self.y_data - y_model) / self.uncertainties, dtype=np.float64
)
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_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, object] = {
"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 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,
)
self._last_result = fit_result
return fit_result
def _compute_covariance(self, result: Any) -> 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 = 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 = cov * chi2 / ndof
return cov
def residuals_plot(
self,
save_path: Optional[str] = None,
) -> 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,
)
if self._last_result is None:
raise RuntimeError("No fit result available. Call fit() first.")
model_values = self._last_result.model_values
residuals = self._last_result.residuals
chi2_per_dof = self._last_result.chi2_per_dof
# 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,
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 = 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 = {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 self._last_result is None:
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 self._last_result is None:
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
beta_spectrum/nuclear_data.py
+1 -1
View File
@@ -616,7 +616,7 @@ def load_json_input(filepath: str) -> Dict[str, Any]:
data[key] = default
validate_json_input(data)
return data # type: ignore[no-any-return]
return data
# Fix the return type issue by explicitly typing the function
beta_spectrum/visualize/__init__.py
+23 -9
View File
@@ -163,7 +163,14 @@ class BetaSpectrumAnalyzer:
def _get_branch_spectrum(self, i: int) -> np.ndarray:
"""Get the spectrum array for branch i."""
if i < len(self.spectrum.branch_spectra):
return self.spectrum.branch_spectra[i]
bs_obj = self.spectrum.branch_spectra[i]
# branch_spectra stores BetaSpectrum objects — call them to get arrays
if hasattr(bs_obj, "__call__"):
val = bs_obj(self.W)
# Clamp to zero beyond endpoint (phase space goes negative)
W0 = T_to_W(self._unified_branches[i].endpoint_MeV)
return np.where(self.W <= W0, val, 0.0)
return np.asarray(bs_obj)
# Single-branch: the spectrum itself is the only branch
return self.spectrum(self.W)
@@ -276,15 +283,19 @@ class BetaSpectrumAnalyzer:
# Compute branch contributions (universal * intensity * branch_spectrum)
branch_contribs = []
# Universal components are shared (no per-branch key prefix)
universal_names = [
"Fermi", "Screening", "FiniteSizeL0",
"ChargeDistributionU", "Exchange",
]
for i in range(self._n_branches):
universal_product = np.ones_like(self.W)
for name in [
"Fermi", "Screening", "FiniteSizeL0",
"ChargeDistributionU", "Exchange",
]:
for name in universal_names:
key = self._get_component_key(i, name)
if key in components:
universal_product *= components[key]
elif name in components: # fallback: universal components
universal_product *= components[name]
bs = self._get_branch_spectrum(i)
contrib = universal_product * self._unified_branches[i].intensity * bs
branch_contribs.append(contrib)
@@ -348,15 +359,18 @@ class BetaSpectrumAnalyzer:
# Panel 1: Total spectrum + per-branch spectra
branch_contribs = []
universal_names = [
"Fermi", "Screening", "FiniteSizeL0",
"ChargeDistributionU", "Exchange",
]
for i in range(n_branches):
universal_product = np.ones_like(self.W)
for name in [
"Fermi", "Screening", "FiniteSizeL0",
"ChargeDistributionU", "Exchange",
]:
for name in universal_names:
key = self._get_component_key(i, name)
if key in components:
universal_product *= components[key]
elif name in components:
universal_product *= components[name]
bs = self._get_branch_spectrum(i)
contrib = universal_product * self._unified_branches[i].intensity * bs
branch_contribs.append(contrib)