""" 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.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._branch_spectra_cache: Optional[List[np.ndarray]] = None 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 @property def branch_spectra(self) -> List[np.ndarray]: """Intensity-weighted branch spectra.""" if self._branch_spectra_cache is None: self._branch_spectra_cache = self.spectrum.get_branch_normalized_spectra( self.W ) assert self._branch_spectra_cache is not None return self._branch_spectra_cache 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) 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_branch_spectrum(self, i: int) -> np.ndarray: """Get the spectrum array for branch i.""" if i < len(self.spectrum.branch_spectra): 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) 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 _has_component(self, i: int, name: str) -> bool: """Check if branch i has a given component.""" key = self._get_component_key(i, name) return key in self.components def _get_component(self, i: int, name: str) -> np.ndarray: """Get a per-branch component array.""" key = self._get_component_key(i, name) return self.components[key] 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