# spectrum.py from __future__ import annotations import logging from typing import Dict, List, Optional, Tuple from dataclasses import dataclass import numpy as np from beta_spectrum.base import SpectrumComponent from beta_spectrum.components.phase_space import PhaseSpace from beta_spectrum.components.fermi import FermiFunction from beta_spectrum.components.finite_size import FiniteSizeL0, ChargeDistributionU from beta_spectrum.components.screening import ScreeningCorrection from beta_spectrum.components.exchange import ExchangeCorrection from beta_spectrum.components.radiative import RadiativeCorrection from beta_spectrum.utils import T_to_W from beta_spectrum.constants import ME_MEV @dataclass class BranchConfig: """ Configuration for a single decay branch. Each branch corresponds to a decay to a specific daughter nuclear state, with its own endpoint energy (Q_value - excitation_energy) and transition type. Attributes ---------- endpoint_MeV : float Effective endpoint energy for this branch (Q_value - excitation_energy). transition_type : str Forbiddenness classification for this branch (A, F1, F1U, F2, ...). intensity : float Branch intensity as a fraction of total decays (0.0–1.0). Intensities across all branches should sum to 1.0 (100%). """ endpoint_MeV: float transition_type: str = "A" intensity: float = 1.0 def __post_init__(self) -> None: if self.intensity < 0 or self.intensity > 1.0: raise ValueError( f"Branch intensity must be in [0, 1], got {self.intensity}" ) if self.endpoint_MeV <= 0: raise ValueError( f"Branch endpoint_MeV must be positive, got {self.endpoint_MeV}" ) @dataclass class SpectrumConfig: """ Configuration for beta spectrum calculation. Supports declarative detector response specification: to enable automatic detector smearing via convolve_with_detector(). Supports multi-branch decay via the `branches` field. When branches are provided, the total spectrum is the intensity-weighted sum of individual branch spectra. Each branch has its own endpoint energy and transition type. """ Z_parent: int Z_daughter: int A_number: int endpoint_MeV: float transition_type: str = "A" # should be [A, F1, F1U, F2, F2U, F3, F3U, F4] e_step_MeV: float = 0.001 neutrino_mass_MeV: float = 0.0 # Toggle components and corrections use_phase_space: bool = True use_fermi: bool = True use_screening: bool = True use_finite_size: bool = True use_charge_dist: bool = True use_radiative: bool = True use_exchange: bool = True # Multi-branch support branches: Optional[List[BranchConfig]] = None # None → single-branch mode intensity_cutoff: float = ( 0.0 # Ignore branches below this fraction (default: 0.0 = all) ) # CSV output options csv_separator: str = ',' # Default separator for CSV output normalize_to_one: bool = False # If True, normalize each branch to 1, ignoring intensity class BetaSpectrum: """ Precise beta spectrum calculator. Supports both single-branch and multi-branch decay modes. In single-branch mode (branches=None), the spectrum is the product of all enabled correction factors. In multi-branch mode (branches provided), each branch is calculated independently with its own endpoint energy and transition type, then summed with intensity weighting. Parameters ---------- components : List[SpectrumComponent], optional List of enabled universal spectral components (Fermi, screening, etc.). Only used in multi-branch mode or for universal components. branch_spectra : List[BetaSpectrum], optional List of per-branch spectrum calculators (each with its own PhaseSpace, etc.). Only used in multi-branch mode. branches : List[BranchConfig], optional Branch configuration metadata (intensities, endpoint energies). logger : logging.Logger, optional Logger for debug/info output. """ def __init__( self, components: Optional[List[SpectrumComponent]] = None, branch_spectra: Optional[List[BetaSpectrum]] = None, branches: Optional[List[BranchConfig]] = None, logger: Optional[logging.Logger] = None, ): self.components = components or [] self._component_names = [self._get_component_name(c) for c in self.components] self.branch_spectra = branch_spectra or [] self.branches = branches or [] self._logger = logger self._is_multi = len(self.branch_spectra) > 0 def __call__(self, W: np.ndarray) -> np.ndarray: """ Calculate spectrum. In single-branch mode: product of all components. In multi-branch mode: intensity-weighted sum of branch spectra. """ if self._is_multi: return self._multi_branch_call(W) return self._single_branch_call(W) def _single_branch_call(self, W: np.ndarray) -> np.ndarray: """Calculate spectrum as product of components (single branch).""" result: np.ndarray = np.ones_like(W, dtype=float) for comp in self.components: if self._logger: self._logger.debug( "Evaluating %s at %d energy points", comp.__class__.__name__, len(W) ) result = result * comp(W) return result.astype(np.float64) def _multi_branch_call(self, W: np.ndarray) -> np.ndarray: """Calculate total spectrum as intensity-weighted sum of branches. The total is: (universal_components_product) * sum(intensity_i * branch_spectrum_i) """ # Sum of intensity-weighted branch spectra branch_sum = np.zeros_like(W, dtype=float) for i, branch_spec in enumerate(self.branch_spectra): if self._logger: self._logger.debug( "Evaluating branch %d/%d (endpoint=%.3f MeV, intensity=%.4f)", i + 1, len(self.branch_spectra), self.branches[i].endpoint_MeV, self.branches[i].intensity, ) branch_spectrum = branch_spec(W) # Mask out energies beyond this branch's endpoint W0_branch = T_to_W(self.branches[i].endpoint_MeV) mask = W <= W0_branch branch_sum += self.branches[i].intensity * np.where( mask, branch_spectrum, 0.0 ) # Multiply by universal components result: np.ndarray = np.ones_like(W, dtype=float) for comp in self.components: result = result * comp(W) return np.asarray((result * branch_sum).astype(np.float64), dtype=np.float64) def _get_component_name(self, comp: SpectrumComponent) -> str: """Get human-readable component name.""" class_name = comp.__class__.__name__ return class_name.replace("Correction", "").replace("Function", "") def calculate_components(self, W: np.ndarray) -> Dict[str, np.ndarray]: """ Calculate individual components. In single-branch mode: returns per-component values. In multi-branch mode: returns universal components + per-branch branch-dependent components (PhaseSpace, Radiative). """ if self._is_multi: return self._multi_branch_components(W) result: Dict[str, np.ndarray] = {} for name, comp in zip(self._component_names, self.components): result[name] = comp(W) return result def _multi_branch_components(self, W: np.ndarray) -> Dict[str, np.ndarray]: """Return universal + per-branch components for multi-branch mode.""" result: Dict[str, np.ndarray] = {} # Universal components (same for all branches) for name, comp in zip(self._component_names, self.components): result[name] = comp(W) # Per-branch components (branch-dependent) for i, branch_spec in enumerate(self.branch_spectra): branch_name = f"branch_{i}" branch_components = branch_spec.calculate_components(W) for comp_name, values in branch_components.items(): result[f"{branch_name}.{comp_name}"] = values return result @classmethod def from_config( cls, config: SpectrumConfig, logger: Optional[logging.Logger] = None ) -> BetaSpectrum: """ Create a BetaSpectrum from configuration. In single-branch mode: creates a single spectrum calculator. In multi-branch mode: creates per-branch calculators + universal components. Parameters ---------- config : SpectrumConfig Configuration with decay parameters and correction toggles. logger : logging.Logger, optional Logger for progress output. """ if config.branches is not None and len(config.branches) > 1: return cls._from_config_multi(config, logger) return cls._from_config_single(config, logger) @classmethod def _from_config_single( cls, config: SpectrumConfig, logger: Optional[logging.Logger] = None ) -> BetaSpectrum: """Create single-branch spectrum calculator.""" components: List[SpectrumComponent] = [] W0 = float(T_to_W(config.endpoint_MeV)) if config.use_phase_space: # Convert neutrino mass from MeV to natural units m_nu_natural = config.neutrino_mass_MeV / ME_MEV components.append( PhaseSpace( W0=W0, transition_type=config.transition_type, m_nu=m_nu_natural, logger=logger, ) ) if config.use_fermi: components.append( FermiFunction(Z=config.Z_daughter, A=config.A_number, logger=logger) ) if config.use_finite_size: components.append( FiniteSizeL0(Z=config.Z_daughter, A=config.A_number, logger=logger) ) if config.use_charge_dist: components.append( ChargeDistributionU( Z=config.Z_daughter, A=config.A_number, logger=logger ) ) if config.use_screening: components.append( ScreeningCorrection( FermiFunction(Z=config.Z_parent, A=config.A_number, logger=logger), logger=logger, ) ) if config.use_exchange: components.append(ExchangeCorrection(Z=config.Z_parent, logger=logger)) if config.use_radiative: components.append( RadiativeCorrection( W0=W0, Z=config.Z_daughter, A=config.A_number, use_endpoint_resummation=True, logger=logger, ) ) return cls(components, logger=logger) @classmethod def _from_config_multi( cls, config: SpectrumConfig, logger: Optional[logging.Logger] = None ) -> BetaSpectrum: """Create multi-branch spectrum calculator. Creates per-branch calculators (each with its own PhaseSpace, Radiative, etc.) and universal components (Fermi, screening, finite_size, etc.) shared across all branches. Note: Radiative correction is branch-dependent (depends on W0), so it goes into each branch's calculator, not the universal components. """ if logger: logger.info("Multi-branch mode: %d branches", len(config.branches)) # Create per-branch calculators branch_spectra: List[BetaSpectrum] = [] if config.branches is None: raise ValueError("branches must be provided for multi-branch mode") for i, branch in enumerate(config.branches): if logger: logger.info( " Branch %d: endpoint=%.3f MeV, transition=%s, intensity=%.4f", i + 1, branch.endpoint_MeV, branch.transition_type, branch.intensity, ) # Create a per-branch config with the branch's endpoint and transition type branch_config = SpectrumConfig( Z_parent=config.Z_parent, Z_daughter=config.Z_daughter, A_number=config.A_number, endpoint_MeV=branch.endpoint_MeV, transition_type=branch.transition_type, e_step_MeV=config.e_step_MeV, neutrino_mass_MeV=config.neutrino_mass_MeV, # Branch-dependent components use_phase_space=config.use_phase_space, use_radiative=config.use_radiative, # Universal components disabled (handled separately) use_fermi=False, use_screening=False, use_finite_size=False, use_charge_dist=False, use_exchange=False, ) branch_spectrum = cls._from_config_single(branch_config, logger) branch_spectra.append(branch_spectrum) # Create universal components (shared across all branches) # Note: Radiative is NOT included here because it depends on W0 universal_components: List[SpectrumComponent] = [] if config.use_fermi: universal_components.append( FermiFunction(Z=config.Z_daughter, A=config.A_number, logger=logger) ) if config.use_finite_size: universal_components.append( FiniteSizeL0(Z=config.Z_daughter, A=config.A_number, logger=logger) ) if config.use_charge_dist: universal_components.append( ChargeDistributionU( Z=config.Z_daughter, A=config.A_number, logger=logger ) ) if config.use_screening: universal_components.append( ScreeningCorrection( FermiFunction(Z=config.Z_parent, A=config.A_number, logger=logger), logger=logger, ) ) if config.use_exchange: universal_components.append( ExchangeCorrection(Z=config.Z_parent, logger=logger) ) return cls( components=universal_components, branch_spectra=branch_spectra, branches=list(config.branches), logger=logger, ) def get_energy_grid(self, config: SpectrumConfig) -> Tuple[np.ndarray, np.ndarray]: """ Generate array of energy values for spectrum calculation. In multi-branch mode, uses the maximum endpoint across all branches. """ if self._is_multi and self.branches: max_endpoint = max(b.endpoint_MeV for b in self.branches) else: max_endpoint = config.endpoint_MeV kinetic_MeV = np.arange(config.e_step_MeV, max_endpoint, config.e_step_MeV) W = T_to_W(kinetic_MeV) return W, kinetic_MeV.astype(np.float64) def get_branch_spectra(self, W: np.ndarray) -> List[np.ndarray]: """ Get individual branch spectra for multi-branch mode. Each branch spectrum is masked to zero beyond its endpoint. Parameters ---------- W : np.ndarray Energy grid in m_e units. Returns ------- List[np.ndarray] Per-branch spectrum arrays (unweighted). Empty in single-branch mode. """ if not self._is_multi: return [] result = [] for i, branch_spec in enumerate(self.branch_spectra): branch_spectrum = branch_spec(W) # Use np.maximum to smoothly approach zero at endpoint, # and clamp negative values beyond endpoint to zero result.append(np.maximum(0.0, branch_spectrum)) return result def get_branch_normalized_spectra(self, W: np.ndarray) -> List[np.ndarray]: """ Get intensity-weighted branch spectra for multi-branch mode. Each branch spectrum is masked to zero beyond its endpoint. Parameters ---------- W : np.ndarray Energy grid in m_e units. Returns ------- List[np.ndarray] Per-branch spectrum arrays weighted by intensity. Empty in single-branch mode. """ if not self._is_multi: return [] result = [] for i, branch_spec in enumerate(self.branch_spectra): branch_spectrum = branch_spec(W) result.append(self.branches[i].intensity * np.maximum(0.0, branch_spectrum)) return result