""" Theoretical model for fitting — wraps BetaSpectrum + DetectorResponse. Provides a callable interface that evaluates the theoretical spectrum, optionally convolved with detector response, on an arbitrary energy grid. Usage: model = SpectrumModel(beta_spectrum, detector_response, channel_energies) y_model = model(channel_energies, normalization, background) """ from __future__ import annotations from typing import TYPE_CHECKING, Any, Optional import numpy as np if TYPE_CHECKING: from beta_spectrum.spectrum import BetaSpectrum from exp_data.spectrum import ExpSpectrum class SpectrumModel: """Theoretical beta spectrum model for fitting to experimental data. Wraps a :class:`beta_spectrum.BetaSpectrum` and optionally a :class:`beta_spectrum.DetectorResponse` to produce a callable model that can be fitted against experimental spectra. The model computes: y(E) = A · [S(E) ⊗ R(E)] + B where S(E) is the theoretical spectrum, R(E) is the detector response, ⊗ denotes convolution, A is a normalization factor, and B is a constant background level. Parameters ---------- spectrum : BetaSpectrum Theoretical spectrum calculator. detector_response : optional Detector response object exposing ``channel_energies`` and ``convolve`` (duck-typed: the DetectorResponse module was removed from beta_spectrum and is pending a restore-or-delete decision). If None, no convolution is applied. channel_energies : np.ndarray, optional Energy values for detector channels (keV). Used to set up the response if detector_response is None. """ def __init__( self, spectrum: "BetaSpectrum", detector_response: Optional[Any] = None, channel_energies: Optional[np.ndarray] = None, ): self.spectrum = spectrum self.detector_response = detector_response self.channel_energies = channel_energies # Cache for pre-computed unconvolved spectrum self._cached_W: Optional[np.ndarray] = None self._cached_spectral: Optional[np.ndarray] = None def evaluate_unconvolved( self, energies_keV: Optional[np.ndarray] = None, ) -> np.ndarray: """Evaluate the theoretical spectrum (unconvolved). Parameters ---------- energies_keV : np.ndarray, optional Energy grid in keV. Uses the spectrum's internal grid if None. Returns ------- np.ndarray Theoretical spectrum values. """ if energies_keV is not None: W = self._keV_to_W(energies_keV) return self.spectrum(W) w_grid = getattr(self.spectrum, "_W_grid", None) if w_grid is None: return np.array([], dtype=np.float64) return self.spectrum(w_grid) def _keV_to_W(self, energies_keV: np.ndarray) -> np.ndarray: """Convert kinetic energy in keV to total energy in m_e units.""" return energies_keV / 510.998950 + 1.0 def convolve(self, energies_keV: np.ndarray) -> np.ndarray: """Convolve theoretical spectrum with detector response. Parameters ---------- energies_keV : np.ndarray Energy grid in keV for the detector channels. Returns ------- np.ndarray Convolved spectrum on the given energy grid. Raises ------ RuntimeError If no detector response is available. """ if self.detector_response is None: raise RuntimeError( "No detector response configured. " "Provide a DetectorResponse to SpectrumModel." ) # The detector's channel_energies are in m_e units W = self.detector_response.channel_energies # Clip W to valid range (W >= 1.0) before spectrum evaluation W_safe = np.maximum(W, 1.0 + 1e-12) spectral = self.spectrum(W_safe) # Clip negative/NaN values from spectrum (numerical artifacts) spectral = np.maximum( np.nan_to_num(spectral, nan=0.0, posinf=0.0, neginf=0.0), 0.0 ) # Convolve: detector_response.convolve(W, spectrum, normalize) return self.detector_response.convolve(W, spectral, normalize=True) def evaluate( self, energies_keV: np.ndarray, normalization: float = 1.0, background: float = 0.0, ) -> np.ndarray: """Evaluate the full theoretical model (with optional convolution). Parameters ---------- energies_keV : np.ndarray Energy grid in keV. normalization : float Multiplicative normalization factor. background : float Constant background level. Returns ------- np.ndarray Model prediction: A · (S ⊗ R) + B """ if self.detector_response is not None: spectral = self.convolve(energies_keV) else: spectral = self.spectrum(self._keV_to_W(energies_keV)) return normalization * spectral + background def residual( self, params: np.ndarray, exp_energies: np.ndarray, exp_counts: np.ndarray, exp_errors: np.ndarray, ) -> np.ndarray: """Compute weighted residuals for fitting. Parameters ---------- params : np.ndarray Fit parameters [normalization, background, ...]. exp_energies : np.ndarray Experimental energy grid (keV). exp_counts : np.ndarray Experimental counts. exp_errors : np.ndarray Experimental uncertainties. Returns ------- np.ndarray Weighted residuals: (exp - model) / σ """ model_values = self.evaluate(exp_energies, params[0], params[1]) return (exp_counts - model_values) / exp_errors @staticmethod def default_params( exp_spectrum: "ExpSpectrum", ) -> tuple[float, float]: """Estimate default fit parameters from experimental data. Parameters ---------- exp_spectrum : ExpSpectrum Experimental spectrum. Returns ------- (normalization, background) Initial parameter estimates. """ # Background from low-energy sideband mask_low = exp_spectrum.energies < np.percentile(exp_spectrum.energies, 5) background = ( float(np.mean(exp_spectrum.counts[mask_low])) if np.any(mask_low) else 0.0 ) # Normalization from total counts / integral of unconvolved model total_counts = float(np.sum(exp_spectrum.counts)) unconvolved = exp_spectrum.counts - background norm = total_counts / max( float(np.trapezoid(unconvolved, exp_spectrum.energies)), 1e-10 ) return (norm, background) @staticmethod def default_bounds( n_channels: int, ) -> tuple[list[float], list[float]]: """Estimate default parameter bounds. Parameters ---------- n_channels : int Number of energy channels. Returns ------- (lower_bounds, upper_bounds) Parameter bounds for [normalization, background]. """ return ( [0.01, 0.0], [100.0, 1e6], )