# spectrum.py from __future__ import annotations import logging from datetime import datetime, timezone from typing import TYPE_CHECKING import numpy as np import matplotlib.pyplot as plt if TYPE_CHECKING: from matplotlib.axes import Axes from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from dataclasses import dataclass 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.logging_utils import get_git_short_hash from beta_spectrum.utils import T_to_W from beta_spectrum.constants import ME_MEV if TYPE_CHECKING: from beta_spectrum.components.detector_response import DetectorResponse @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: set detector_model, detector_sigma_a_keV, and related parameters 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 # 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 # Detector response convolution (analytical model) use_detector_response: bool = False detector_model: str = "gaussian" detector_sigma_a_keV: float = 1.0 detector_sigma_b: float = 0.0 detector_tail_fraction: float = 0.0 detector_tau_keV: float = 5.0 detector_fano_factor: float = 0.12 detector_n_channels: int = 4096 detector_channel_energy_range: tuple[float, float] = ( 0.0, 0.35, ) # in m_e units (total energy) # 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) ) 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: components.append( PhaseSpace(W0=W0, transition_type=config.transition_type, 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] = [] 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, # 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 @staticmethod def create_detector_from_config( config: SpectrumConfig, ) -> "DetectorResponse": """ Create a DetectorResponse from SpectrumConfig detector parameters. Converts keV-based resolution parameters to m_e units internally. Parameters ---------- config : SpectrumConfig Configuration with detector response parameters set. Returns ------- DetectorResponse Detector response object ready for convolution. """ from beta_spectrum.components.detector_response import DetectorResponse from beta_spectrum.utils import T_to_W sigma_a_me = config.detector_sigma_a_keV / ME_MEV tau_me = config.detector_tau_keV / ME_MEV W0 = T_to_W(config.endpoint_MeV) channel_range = config.detector_channel_energy_range # Extend range to cover endpoint if channel_range[1] < W0: channel_range = (channel_range[0], float(W0 + 0.05)) detector = DetectorResponse.from_gaussian_params( channel_energy_range=channel_range, n_channels=config.detector_n_channels, sigma_a=sigma_a_me, sigma_b=config.detector_sigma_b, tail_fraction=config.detector_tail_fraction, tau=tau_me, model=config.detector_model, fano_factor=config.detector_fano_factor, ) return detector def convolve_with_detector( self, detector_response: "DetectorResponse", W: Optional[np.ndarray] = None, config: Optional[SpectrumConfig] = None, ) -> np.ndarray: """ Convolve theoretical spectrum with detector response. Returns the predicted measured spectrum after detector smearing. Parameters ---------- detector_response : DetectorResponse Detector response object (analytical or tabulated). W : np.ndarray, optional Energy grid in m_e units. If None, generated from config. config : SpectrumConfig, optional Used to generate energy grid if W is None. Returns ------- np.ndarray Convolved spectrum (predicted measured counts per channel). """ if W is None: if config is None: raise ValueError("Either W or config must be provided") W, _ = self.get_energy_grid(config) theoretical_spectrum = self(W) if self._logger: self._logger.debug( "Convolution: %d channels, %d spectrum points", len(detector_response.channel_energies), len(W), ) convolved = detector_response.convolve(W, theoretical_spectrum, normalize=True) return convolved def convolve_detector( self, config: SpectrumConfig, W: Optional[np.ndarray] = None, detector_response: Optional["DetectorResponse"] = None, ) -> np.ndarray: """ Convolve theoretical spectrum with detector response from config. Convenience method: creates DetectorResponse from config parameters and convolves the spectrum in a single call. Parameters ---------- config : SpectrumConfig Configuration with detector response parameters set. W : np.ndarray, optional Energy grid in m_e units. If None, generated from config. detector_response : DetectorResponse, optional Pre-built detector response. If None, created from config. Returns ------- np.ndarray Convolved spectrum (predicted measured counts per channel). Examples -------- >>> config = SpectrumConfig( ... Z_parent=43, Z_daughter=44, A_number=99, endpoint_MeV=0.294, ... use_detector_response=True, ... detector_sigma_a_keV=1.0, ... ) >>> spectrum = BetaSpectrum.from_config(config) >>> W, _ = spectrum.get_energy_grid(config) >>> convolved = spectrum.convolve_detector(config, W=W) """ if detector_response is None: detector_response = self.create_detector_from_config(config) return self.convolve_with_detector(detector_response, W=W, config=config) 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 @property def components(self) -> Dict[str, np.ndarray]: if self._components_cache is None: self._components_cache = self.spectrum.calculate_components(self.W) 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 def convolved_spectrum( self, detector_response: Optional["DetectorResponse"] = None, normalize: bool = True, ) -> np.ndarray: """ Calculate spectrum convolved with detector response. If detector_response is None and use_detector_response is True in config, creates a detector response from config parameters. Parameters ---------- detector_response : DetectorResponse, optional Detector response to convolve with. Created from config if None. normalize : bool If True, normalize to unit area. Returns ------- np.ndarray Convolved spectrum. """ if detector_response is None and self.config.use_detector_response: detector_response = BetaSpectrum.create_detector_from_config(self.config) if detector_response is None: raise ValueError( "detector_response must be provided, or set " "use_detector_response=True in config" ) convolved = self.spectrum.convolve_with_detector( detector_response, W=self.W, config=self.config ) if normalize: integral = np.trapezoid(convolved, detector_response.channel_energies) if integral > 0: convolved = convolved / integral return convolved def plot_analysis( self, save_path: Optional[str] = None, show_components: bool = True, ) -> None: """ Create visualization of the spectrum and all correction factors. In single-branch mode: standard 4-panel debug view or single spectrum plot. In multi-branch mode: spectrum plot with individual branches shown, and debug view with vertical layout for branch comparison. Parameters ---------- save_path : str, optional Path to save the figure. show_components : bool If True, show full debug view. If False, show only the spectrum plot with nuclear data header. If True, use logarithmic y-scale. Set to False for linear scale, which is better when branch intensities are similar. """ total = self.total_spectrum(normalize=True) components = self.components commit = get_git_short_hash(6) timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") if self._is_multi: if show_components: self._plot_multi_branch_debug( total, components, save_path, commit, timestamp ) else: self._plot_multi_branch_spectra( total, components, save_path, commit, timestamp ) else: if show_components: self._plot_debug_view(total, components, save_path, commit, timestamp) else: self._plot_spectrum_only(total, save_path, commit, timestamp) def _add_id_textbox(self, ax: "Axes", commit: str, timestamp: str) -> None: """Add ID text box with commit hash and timestamp to plot.""" id_text = f"commit: {commit} | {timestamp}" ax.text( 0.01, 0.01, id_text, transform=ax.transAxes, fontsize=8, color="gray", verticalalignment="bottom", horizontalalignment="left", fontfamily="monospace", ) def _add_nuclear_data_header(self, ax: "Axes") -> None: """Add nuclear data information header to plot.""" parent = self._element_symbol(self.config.Z_parent) daughter = self._element_symbol(self.config.Z_daughter) enabled_corrections = [] if self.config.use_phase_space: enabled_corrections.append("phase_space") if self.config.use_fermi: enabled_corrections.append("fermi") if self.config.use_screening: enabled_corrections.append("screening") if self.config.use_finite_size: enabled_corrections.append("finite_size") if self.config.use_charge_dist: enabled_corrections.append("charge_dist") if self.config.use_radiative: enabled_corrections.append("radiative") if self.config.use_exchange: enabled_corrections.append("exchange") header_lines = [ f"Nuclide: {parent}{self.config.A_number} -> {daughter}{self.config.A_number}", f"Endpoint: {self.config.endpoint_MeV * 1000:.1f} keV", f"Transition: {self.config.transition_type}", f"Corrections: {', '.join(enabled_corrections)}", ] if self.config.use_detector_response: header_lines.append( f"Detector: {self.config.detector_model} " f"(σ={self.config.detector_sigma_a_keV} keV)" ) header_text = "\n".join(header_lines) ax.text( 0.99, 0.99, header_text, transform=ax.transAxes, fontsize=8, verticalalignment="top", horizontalalignment="right", bbox=dict( boxstyle="round,pad=0.3", facecolor="white", alpha=0.8, edgecolor="none" ), ) def _plot_spectrum_only( self, total: np.ndarray, save_path: Optional[str], commit: str, timestamp: str, ) -> None: """Plot only the total spectrum with nuclear data header.""" fig, ax = plt.subplots(figsize=(10, 6)) ax.plot(self.energies_MeV, total, "b-", lw=2, label="Normalized spectrum") ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=11) ax.set_ylabel("Normalized Counts", fontsize=11) ax.set_title( f"Beta-decay: {self.config.Z_parent} -> {self.config.Z_daughter}, A={self.config.A_number}", fontsize=13, ) ax.set_yscale("log") ax.grid(True, alpha=0.3) ax.legend(loc="best", fontsize=10) self._add_nuclear_data_header(ax) self._add_id_textbox(ax, commit, timestamp) plt.tight_layout() if save_path: fig.savefig(save_path, dpi=150, bbox_inches="tight") print(f"Figure saved to {save_path}") plt.show() def _plot_debug_view( self, total: np.ndarray, components: Dict[str, np.ndarray], save_path: Optional[str], commit: str, timestamp: str, ) -> None: """Full debug view with 4-panel spectrum analysis.""" fig = plt.figure(figsize=(14, 10)) # 1. Main spectrum plot (top-left) ax1 = plt.subplot(2, 2, 1) ax1.plot( self.energies_MeV, total, "b-", lw=2, label="Total spectrum (log scale)" ) ax1.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=10) ax1.set_ylabel( f"Normalized Counts per {self.config.e_step_MeV:.3e} MeV", fontsize=10 ) ax1.set_title( f"Beta-decay: Z={self.config.Z_parent} -> {self.config.Z_daughter}, A={self.config.A_number}", fontsize=12, ) ax1.set_yscale("log") ax1.grid(True, alpha=0.3) ax1.legend() # 2. Correction plots (top-right) ax2 = plt.subplot(2, 2, 2) # Set color order for correction plots component_order = [ "PhaseSpace", "Fermi", "Screening", "Exchange", "FiniteSizeL0", "ChargeDistributionU", "Radiative", ] colors = ["gray", "red", "green", "blue", "orange", "purple", "brown"] for name, color in zip(component_order, colors): if name in components: values = components[name] # For Fermi and PhaseSpace, scale for visibility if name in ["Fermi", "PhaseSpace"]: values = values / np.max(values) label = f"{name} (norm)" else: label = name ax2.plot( self.energies_MeV, values, color=color, lw=1.5, label=label, alpha=0.8, ) ax2.set_xlabel("Electron kinetic energy E [MeV]", fontsize=10) ax2.set_ylabel("Correction factor", fontsize=10) ax2.set_title("Spectrum components", fontsize=12) ax2.grid(True, alpha=0.3) ax2.legend(loc="best", fontsize=8) # 3. Cumulative effect (botom-left) ax3 = plt.subplot(2, 2, 3) # Cumulative product cumulative = np.ones_like(self.W) baseline_norm = np.ones_like(cumulative) # Plot baseline ax3.plot( self.energies_MeV, baseline_norm, "k--", lw=1.5, label="Baseline", alpha=0.7 ) # Add components one by one for name, color in zip(component_order, colors): if name in components: cumulative *= components[name] norm_cumulative = cumulative / np.max(cumulative) ax3.plot( self.energies_MeV, norm_cumulative, color=color, lw=1.5, label=f"+ {name}", alpha=0.7, ) ax3.set_xlabel("Electron kinetic energy E [MeV]", fontsize=10) ax3.set_ylabel("Normalized spectrum", fontsize=10) ax3.set_title("Cumulative effect", fontsize=12) ax3.grid(True, alpha=0.3) ax3.legend(loc="best", fontsize=8) # 4. Deviation from unity (bottom-right) ax4 = plt.subplot(2, 2, 4) for name, color in zip(component_order, colors): if name in components and name not in ["Fermi", "PhaseSpace"]: deviation = components[name] - 1.0 ax4.plot( self.energies_MeV, deviation, color=color, lw=1.5, label=f"{name}", alpha=0.7, ) ax4.set_xlabel("Electron kinetic energy E [MeV]", fontsize=10) ax4.set_ylabel("Deviation from unity", fontsize=10) ax4.set_title("Correction deviations (C - 1)", fontsize=12) ax4.grid(True, alpha=0.3) ax4.legend(loc="best", fontsize=8) ax4.axhline(y=0, color="k", linestyle="-", lw=0.5) plt.suptitle( f"Beta-decay spectrum {self.config.A_number}{self._element_symbol(self.config.Z_parent)} -> {self.config.A_number}{self._element_symbol(self.config.Z_daughter)}", fontsize=14, fontweight="bold", ) plt.tight_layout() if save_path: fig.savefig(save_path, dpi=150, bbox_inches="tight") print(f"Figure saved to {save_path}") plt.show() def _element_symbol(self, Z: int) -> str: 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", "Fm", ] return symbols[Z] if Z < len(symbols) else f"Z{Z}" # --------------------------------------------------------------------------- # Multi-branch plotting # --------------------------------------------------------------------------- def _get_branch_label(self, i: int) -> str: """Generate branch label for legend.""" if i >= len(self.spectrum.branches): return f"Branch {i+1}" branch = self.spectrum.branches[i] return f"Branch {i+1}: {branch.transition_type}, E₀={branch.endpoint_MeV*1000:.1f} keV" def _plot_multi_branch_spectra( self, total: np.ndarray, components: Dict[str, np.ndarray], save_path: Optional[str], commit: str, timestamp: str, ) -> None: """Plot total spectrum with individual branch decomposition. All curves use raw (un-normalized) values so that the total is the arithmetic sum of branch contributions, and branch amplitudes are proportional to their intensities. """ fig, ax = plt.subplots(figsize=(12, 7)) branch_colors = [ "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf", ] # Compute branch contributions (universal * intensity * branch_spectrum) branch_contribs = [] for i, branch_spec in enumerate(self.spectrum.branch_spectra): universal_product = np.ones_like(self.W) for name in [ "Fermi", "Screening", "FiniteSizeL0", "ChargeDistributionU", "Exchange", ]: if name in components: universal_product *= components[name] bs = self.spectrum.get_branch_spectra(self.W) contrib = universal_product * self.spectrum.branches[i].intensity * bs[i] branch_contribs.append(contrib) # Total = arithmetic sum of all branches (raw, un-normalized) total_raw = sum(branch_contribs) ax.plot( self.energies_MeV, total_raw, "k-", lw=2.5, label="Total spectrum", zorder=10, ) # Plot each branch contribution (raw, un-normalized) for i, (contrib, branch) in enumerate( zip(branch_contribs, self.spectrum.branches) ): color = branch_colors[i % len(branch_colors)] ax.plot( self.energies_MeV, contrib, color=color, lw=1.5, alpha=0.7, label=self._get_branch_label(i), ) ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=12) ax.set_ylabel("Normalized Counts", fontsize=12) parent = self._element_symbol(self.config.Z_parent) daughter = self._element_symbol(self.config.Z_daughter) ax.set_title( f"Beta-decay: {parent}{self.config.A_number} -> {daughter}{self.config.A_number} — {len(self.spectrum.branches)} branches", fontsize=14, ) ax.set_yscale("log") ax.grid(True, alpha=0.3) ax.legend(loc="best", fontsize=10, ncol=min(2, len(self.spectrum.branches) + 1)) self._add_nuclear_data_header(ax) self._add_id_textbox(ax, commit, timestamp) plt.tight_layout() if save_path: fig.savefig(save_path, dpi=150, bbox_inches="tight") print(f"Figure saved to {save_path}") def _plot_multi_branch_debug( self, total: np.ndarray, components: Dict[str, np.ndarray], save_path: Optional[str], commit: str, timestamp: str, ) -> None: """Debug view for multi-branch: vertical layout with 4 panels. Panel 1: Total spectrum + intensity-scaled branch spectra (log scale) Panel 2: Fermi function + all phase spaces (shape comparison, no intensity) Panel 3: Universal W0-independent corrections (Fermi, Screening, etc.) Panel 4: W0-dependent corrections (Radiative for each branch) """ n_branches = len(self.spectrum.branches) fig, axes = plt.subplots(4, 1, figsize=(12, 16), sharex=True) branch_colors = [ "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf", ] parent = self._element_symbol(self.config.Z_parent) daughter = self._element_symbol(self.config.Z_daughter) # Panel 1: Total spectrum + per-branch spectra # All curves use raw (un-normalized) values so total = sum of branches branch_contribs = [] for i, branch_spec in enumerate(self.spectrum.branch_spectra): universal_product = np.ones_like(self.W) for name in [ "Fermi", "Screening", "FiniteSizeL0", "ChargeDistributionU", "Exchange", ]: if name in components: universal_product *= components[name] bs = self.spectrum.get_branch_spectra(self.W) contrib = universal_product * self.spectrum.branches[i].intensity * bs[i] branch_contribs.append(contrib) total_raw = sum(branch_contribs) ax = axes[0] ax.plot( self.energies_MeV, total_raw, "k-", lw=2, label="Total spectrum", zorder=10, ) for i, (contrib, branch) in enumerate( zip(branch_contribs, self.spectrum.branches) ): color = branch_colors[i % len(branch_colors)] ax.plot( self.energies_MeV, contrib, color=color, lw=1.5, alpha=0.7, label=f"Branch {i+1}: {branch.transition_type}, " f"E₀={branch.endpoint_MeV*1000:.1f} keV", ) ax.set_yscale("log") ax.set_ylabel("Normalized Counts", fontsize=10) ax.set_title( f"Multi-branch: {parent}{self.config.A_number} -> {daughter}{self.config.A_number} ({n_branches} branches)", fontsize=12, fontweight="bold", ) ax.grid(True, alpha=0.3) ax.legend(fontsize=8, loc="best", ncol=min(2, n_branches + 1)) # Panel 2: Fermi function + all phase spaces (shape comparison) # Phase space is masked beyond endpoint before normalizing to avoid # negative values from the sqrt term ax = axes[1] if "Fermi" in components: fermi_vals = components["Fermi"] fermi_scaled = fermi_vals / np.max(fermi_vals) ax.plot( self.energies_MeV, fermi_scaled, "r-", lw=1.5, label="Fermi (normalized)", alpha=0.8, ) for i, (branch_spec, branch) in enumerate( zip(self.spectrum.branch_spectra, self.spectrum.branches) ): color = branch_colors[i % len(branch_colors)] if f"branch_{i}.PhaseSpace" in components: ps_vals = components[f"branch_{i}.PhaseSpace"] # Mask beyond endpoint to avoid negative values W0_branch = T_to_W(branch.endpoint_MeV) mask = self.W <= W0_branch ps_masked = np.where(mask, ps_vals, 0.0) ps_max = np.max(ps_masked) if ps_max > 0: ps_scaled = ps_masked / ps_max else: ps_scaled = ps_masked ax.plot( self.energies_MeV, ps_scaled, color=color, lw=1.5, alpha=0.7, label=f"PS Branch {i+1} (normalized)", ) ax.set_ylabel("Normalized Factor", fontsize=10) ax.set_title( "Fermi Function & Phase Space Shapes", fontsize=11, fontweight="bold" ) ax.grid(True, alpha=0.3) ax.legend(fontsize=8, loc="best") # Panel 3: Universal W0-independent corrections ax = axes[2] universal_components = [ "Screening", "FiniteSizeL0", "ChargeDistributionU", "Exchange", ] colors = ["green", "orange", "purple", "brown"] for name, color in zip(universal_components, colors): if name in components: values = components[name] values_scaled = values / np.max(values) ax.plot( self.energies_MeV, values_scaled, color=color, lw=1.5, label=f"{name} (normalized)", alpha=0.8, ) ax.set_ylabel("Normalized Factor", fontsize=10) ax.set_title( "Universal W0-Independent Corrections", fontsize=11, fontweight="bold" ) ax.grid(True, alpha=0.3) ax.legend(fontsize=8, loc="best") # Panel 4: W0-dependent corrections (Radiative for each branch) ax = axes[3] for i, (branch_spec, branch) in enumerate( zip(self.spectrum.branch_spectra, self.spectrum.branches) ): color = branch_colors[i % len(branch_colors)] if f"branch_{i}.Radiative" in components: rad_vals = components[f"branch_{i}.Radiative"] ax.plot( self.energies_MeV, rad_vals, color=color, lw=1.5, alpha=0.7, label=f"Branch {i+1} Radiative (E₀={branch.endpoint_MeV*1000:.1f} keV)", ) ax.set_ylabel("Radiative Correction", fontsize=10) ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=10) ax.set_title( "W0-Dependent Corrections (Radiative)", fontsize=11, fontweight="bold" ) ax.grid(True, alpha=0.3) ax.legend(fontsize=8, loc="best", ncol=min(2, n_branches + 1)) self._add_id_textbox(axes[3], commit, timestamp) plt.tight_layout() if save_path: fig.savefig(save_path, dpi=150, bbox_inches="tight") print(f"Figure saved to {save_path}") def export_to_csv(self, filename: str, source_type: str = "unknown") -> None: """ Export spectrum data to CSV with metadata header. Parameters ---------- filename : str Output CSV file path. source_type : str Data source type (paceENSDF, json, cli) for the header. """ total = self.total_spectrum(normalize=True) components = self.components # Build enabled corrections list enabled = [] if self.config.use_phase_space: enabled.append("phase_space") if self.config.use_fermi: enabled.append("fermi") if self.config.use_screening: enabled.append("screening") if self.config.use_finite_size: enabled.append("finite_size") if self.config.use_charge_dist: enabled.append("charge_dist") if self.config.use_radiative: enabled.append("radiative") if self.config.use_exchange: enabled.append("exchange") # Build metadata header from datetime import datetime, timezone def _element_symbol(Z: int) -> str: 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", "Fm", ] return symbols[Z] if Z < len(symbols) else f"Z{Z}" parent_symbol = _element_symbol(self.config.Z_parent) daughter_symbol = _element_symbol(self.config.Z_daughter) header_lines = [ f"# beta-spectrum v{__import__('beta_spectrum').__version__}", f"# timestamp: {datetime.now(timezone.utc).isoformat()}", f"# source: {source_type}", f"# nuclide: {parent_symbol}{self.config.A_number} -> {daughter_symbol}{self.config.A_number} (Z={self.config.Z_parent}->{self.config.Z_daughter})", f"# endpoint: {self.config.endpoint_MeV * 1000:.1f} keV", f"# transition: {self.config.transition_type}", f"# corrections: {', '.join(enabled)}", f"# e_step: {self.config.e_step_MeV:.4f} MeV", ] if self._is_multi and self.spectrum.branches: header_lines.append(f"# branches: {len(self.spectrum.branches)}") for i, branch in enumerate(self.spectrum.branches): header_lines.append( f"# branch_{i+1}: E₀={branch.endpoint_MeV*1000:.1f} keV, " f"transition={branch.transition_type}, " f"intensity={branch.intensity:.4f}" ) else: header_lines.append("# branches: 1 (single)") if self.config.use_detector_response: header_lines.append( f"# detector: {self.config.detector_model} " f"(sigma={self.config.detector_sigma_a_keV} keV, " f"tail={self.config.detector_tail_fraction})" ) else: header_lines.append("# detector: disabled") header_lines.append(f"# git_commit: {get_git_short_hash()}") # Write CSV with header import pandas as pd data: Dict[str, np.ndarray] = { "energy_MeV": self.energies_MeV, "spectrum": total, } if self._is_multi and self.spectrum.branches: # Multi-branch mode: total spectrum + per-branch spectra + components # (branch info — intensity, endpoint, transition — is in the header, not repeated columns) for i, branch_spec in enumerate(self.spectrum.branch_spectra): data[f"branch_{i+1}_spectrum"] = branch_spec(self.W) # Per-branch components (PhaseSpace, Radiative, etc.) for i, branch_spec in enumerate(self.spectrum.branch_spectra): branch_comps = branch_spec.calculate_components(self.W) for comp_name, values in branch_comps.items(): data[f"branch_{i+1}_{comp_name}"] = values # Universal components only (skip per-branch keys) for name, values in components.items(): if not name.startswith("branch_"): data[f"universal_{name}"] = values else: # Single-branch mode: standard export for name, values in components.items(): data[name] = values df = pd.DataFrame(data) df.to_csv(filename, index=False, float_format="%.4e") # Prepend header with open(filename, "r") as f: content = f.read() with open(filename, "w") as f: f.write("\n".join(header_lines) + "\n" + content) if self._logger: mode = "multi-branch" if self._is_multi else "single-branch" self._logger.info("CSV exported to %s (%s mode)", filename, mode) 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, }