unzhakov
33ed4f7e44
- Add proper type annotations to all source files - Fix return types for numpy operations (np.asarray with dtype=np.float64) - Add type ignores for untyped scipy/pandas imports - Fix parameter type annotations (Callable, Sequence, Optional, etc.) - Add TYPE_CHECKING imports for forward references - Fix fitter.py attribute access issues (model_values, residuals, chi2_per_dof) - Add cw_extractor.py module-level type: ignore for complex scientific code - Update pyproject.toml to exclude notebooks from mypy - Run black and ruff for formatting and linting
566 lines
17 KiB
Python
566 lines
17 KiB
Python
# 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 # type: ignore[import-untyped]
|
||
|
||
|
||
@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 # type: ignore[import-untyped]
|
||
|
||
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(
|
||
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
|