""" Core analyzer module — BetaSpectrumAnalyzer base class. Core introspection and analysis tools for BetaSpectrum. All nuclides (single- and multi-branch) are treated uniformly as one or more branches — single-branch is just a special case with one branch. """ from __future__ import annotations import logging from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, List, Optional import numpy as np from beta_spectrum.components.phase_space import PhaseSpace from beta_spectrum.components.radiative import RadiativeCorrection from beta_spectrum.utils import T_to_W if TYPE_CHECKING: from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig class BetaSpectrumAnalyzer: """ Introspection, analysis and debugging tools for BetaSpectrum. Parameters ---------- spectrum : BetaSpectrum The spectrum to analyze. config : SpectrumConfig Configuration used to create the spectrum. logger : logging.Logger, optional Logger for progress output. """ def __init__( self, spectrum: BetaSpectrum, config: SpectrumConfig, logger: Optional[logging.Logger] = None, ): """ Initialize analyzer with a spectrum and configuration. """ self.spectrum = spectrum self.config = config self.W, self.energies_MeV = spectrum.get_energy_grid(config) self._components_cache: Optional[Dict[str, np.ndarray]] = None self._universal_cache: Optional[np.ndarray] = None self._physical_branch_cache: Dict[int, np.ndarray] = {} self._logger = logger self._is_multi = spectrum._is_multi self._timestamp: Optional[str] = None # Unified branch view: single-branch treated as one-branch multi-branch if not self._is_multi and self.spectrum.branch_spectra == []: self._unified_branches: List[_UnifiedBranch] = [ _UnifiedBranch( endpoint_MeV=config.endpoint_MeV, transition_type=config.transition_type, intensity=1.0, ) ] else: self._unified_branches = [ _UnifiedBranch( endpoint_MeV=b.endpoint_MeV, transition_type=b.transition_type, intensity=b.intensity, ) for b in self.spectrum.branches ] @property def components(self) -> Dict[str, np.ndarray]: if self._components_cache is None: raw = self.spectrum.calculate_components(self.W) # Normalize keys: single-branch has bare names, # multi-branch has branch_0., branch_1., etc. # We keep the raw keys for export_to_csv compatibility. self._components_cache = raw assert self._components_cache is not None return self._components_cache def universal_factor(self) -> np.ndarray: """Product of the components shared by all branches. The universal factor always contains the Fermi function (the core 0-order Coulomb correction) and every enabled higher-order correction that does not depend on the branch endpoint (screening, finite size, charge distribution, exchange). In single-branch mode this equals the product of all components except the per-branch phase space and radiative corrections. """ if self._universal_cache is None: prod: Optional[np.ndarray] = None for comp in self.spectrum.components: if not self._is_multi and isinstance( comp, (PhaseSpace, RadiativeCorrection) ): continue values = comp(self.W) prod = ( np.array(values, dtype=np.float64) if prod is None else prod * values ) if prod is None: prod = np.ones_like(self.W, dtype=np.float64) self._universal_cache = prod assert self._universal_cache is not None return self._universal_cache def physical_branch_spectrum(self, i: int) -> np.ndarray: """Physical spectrum of branch i: core × corrections, unnormalized. The physical branch spectrum is core(phase_space × fermi) × all enabled higher-order corrections, masked to zero beyond the branch endpoint. This is the canonical branch accessor: every consumer (CSV export, plots) must use it so that the Fermi function is always present in branch-level spectra. """ if i in self._physical_branch_cache: return self._physical_branch_cache[i] if i < len(self.spectrum.branch_spectra): raw = self.spectrum.branch_spectra[i](self.W) W0 = T_to_W(self._unified_branches[i].endpoint_MeV) raw = np.where(self.W <= W0, raw, 0.0) physical = self.universal_factor() * raw else: # Single-branch mode: the spectrum itself is the full physical branch physical = np.asarray(self.spectrum(self.W), dtype=np.float64) self._physical_branch_cache[i] = physical return physical def total_integral(self) -> float: """Integral of the raw total spectrum over kinetic energy. Raises ------ ValueError If the integral is not positive (e.g. all branches are degenerate because the neutrino mass exceeds every endpoint). """ total_raw = self.spectrum(self.W) integral = float(np.trapezoid(total_raw, self.energies_MeV)) if integral <= 0.0: raise ValueError( "Total spectrum integral is non-positive: every branch is " "degenerate (check neutrino_mass vs branch endpoints and " "intensity_cutoff)" ) return integral def total_spectrum(self, normalize: bool = True) -> np.ndarray: """Calculate the total spectrum.""" total = self.spectrum(self.W) if normalize: integral = np.trapezoid(total, self.energies_MeV) if integral <= 0.0: raise ValueError( "Total spectrum integral is non-positive: every branch " "is degenerate (check neutrino_mass vs branch endpoints " "and intensity_cutoff)" ) total = total / integral return total # --------------------------------------------------------------------------- # Unified branch view # --------------------------------------------------------------------------- @property def _n_branches(self) -> int: """Number of branches (unified view).""" return len(self._unified_branches) def _get_component_key(self, i: int, name: str) -> str: """Get the canonical key for a per-branch component.""" if self._is_multi: return f"branch_{i}.{name}" return name def _element_symbol(self, Z: int) -> str: """Return chemical symbol for atomic number Z.""" symbols = [ "", "H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg", "Al", "Si", "P", "S", "Cl", "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni", "Cu", "Zn", "Ga", "Ge", "As", "Se", "Br", "Kr", "Rb", "Sr", "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd", "In", "Sn", "Sb", "Te", "I", "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt", "Au", "Hg", "Tl", "Pb", "Bi", "Po", "At", "Rn", "Fr", "Ra", "Ac", "Th", "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk", "Cf", "Es", ] return symbols[Z] if Z < len(symbols) else f"Z{Z}" def get_data(self) -> Dict[str, Any]: """Get all numerical data for custom analysis.""" return { "energies_MeV": self.energies_MeV, "energies_W": self.W, "spectrum": self.total_spectrum(normalize=True), "components": self.components, "config": self.config, } def export_to_csv( self, filename: str, source_type: str = "unknown", csv_verbosity: str = "full" ) -> None: """Export spectrum data to CSV with metadata header.""" from .csv_export import CSVEmitter CSVEmitter.export_to_csv(self, filename, source_type, csv_verbosity) def plot_analysis( self, save_path: Optional[str] = None, show_components: bool = True, ) -> None: """Create visualization of the spectrum and all correction factors.""" from .visualization import Visualizer Visualizer.plot_analysis(self, save_path, show_components) @dataclass class _UnifiedBranch: """Unified branch view for both single- and multi-branch modes.""" endpoint_MeV: float transition_type: str intensity: float