From f05a2de1b8dc4bda63aec918a5c09e3c763511e0 Mon Sep 17 00:00:00 2001 From: Evgeniy Unzhakov Date: Sat, 16 May 2026 20:41:58 +0300 Subject: [PATCH 1/4] refactor: move BetaSpectrumAnalyzer to separate analyzer module --- beta_spectrum/__init__.py | 3 +- beta_spectrum/analyzer.py | 1027 +++++++++++++++++++++++++++ beta_spectrum/cli.py | 2 +- beta_spectrum/spectrum.py | 1017 +------------------------- tests/quality/test_logging_utils.py | 35 +- 5 files changed, 1042 insertions(+), 1042 deletions(-) create mode 100644 beta_spectrum/analyzer.py diff --git a/beta_spectrum/__init__.py b/beta_spectrum/__init__.py index b836a81..0b28154 100644 --- a/beta_spectrum/__init__.py +++ b/beta_spectrum/__init__.py @@ -11,7 +11,8 @@ from .components.exchange import ExchangeCorrection from .components.radiative import RadiativeCorrection from .components.detector_response import DetectorResponse -from .spectrum import BetaSpectrum, SpectrumConfig, BetaSpectrumAnalyzer, BranchConfig +from .spectrum import BetaSpectrum, SpectrumConfig, BranchConfig +from .analyzer import BetaSpectrumAnalyzer from .fitter import CurveFitter, FitConfig, FitResult from .cw_extractor import CWExtractor, CWExtractionResult, GVAExtractionResult from .nuclear_data import ( diff --git a/beta_spectrum/analyzer.py b/beta_spectrum/analyzer.py new file mode 100644 index 0000000..ac852f4 --- /dev/null +++ b/beta_spectrum/analyzer.py @@ -0,0 +1,1027 @@ +""" +Analyzer module — BetaSpectrumAnalyzer class. + +Introspection, analysis, and debugging tools for BetaSpectrum. +""" + +from __future__ import annotations + +import logging +import pandas as pd +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import numpy as np +import matplotlib.pyplot as plt + +from beta_spectrum.logging_utils import get_git_short_hash +from beta_spectrum.utils import T_to_W + +if TYPE_CHECKING: + from matplotlib.axes import Axes + + from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig + from beta_spectrum.components.detector_response import DetectorResponse + + +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 + 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 + 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 (same for all branches — no prefix needed) + for name, values in components.items(): + if not name.startswith("branch_"): + data[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, + } diff --git a/beta_spectrum/cli.py b/beta_spectrum/cli.py index 19b1e40..d50350f 100644 --- a/beta_spectrum/cli.py +++ b/beta_spectrum/cli.py @@ -196,10 +196,10 @@ def _run(args: argparse.Namespace) -> None: ) from beta_spectrum.spectrum import ( BetaSpectrum, - BetaSpectrumAnalyzer, BranchConfig, SpectrumConfig, ) + from beta_spectrum.analyzer import BetaSpectrumAnalyzer # Determine logging level if args.quiet: diff --git a/beta_spectrum/spectrum.py b/beta_spectrum/spectrum.py index 5622847..f5dfa8b 100644 --- a/beta_spectrum/spectrum.py +++ b/beta_spectrum/spectrum.py @@ -2,16 +2,10 @@ from __future__ import annotations import logging -from datetime import datetime, timezone -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from dataclasses import dataclass 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 @@ -20,8 +14,6 @@ from beta_spectrum.components.finite_size import FiniteSizeL0, ChargeDistributio 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 @@ -608,1008 +600,3 @@ class BetaSpectrum: 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 (same for all branches — no prefix needed) - for name, values in components.items(): - if not name.startswith("branch_"): - data[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, - } diff --git a/tests/quality/test_logging_utils.py b/tests/quality/test_logging_utils.py index 78399a4..f610c80 100644 --- a/tests/quality/test_logging_utils.py +++ b/tests/quality/test_logging_utils.py @@ -130,11 +130,8 @@ class TestCSVMetadataHeader: """Test CSV export metadata header.""" def test_header_contains_version(self): - from beta_spectrum.spectrum import ( - BetaSpectrum, - BetaSpectrumAnalyzer, - SpectrumConfig, - ) + from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig + from beta_spectrum.analyzer import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, @@ -161,11 +158,8 @@ class TestCSVMetadataHeader: os.unlink(filename) def test_header_contains_timestamp(self): - from beta_spectrum.spectrum import ( - BetaSpectrum, - BetaSpectrumAnalyzer, - SpectrumConfig, - ) + from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig + from beta_spectrum.analyzer import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, @@ -191,11 +185,8 @@ class TestCSVMetadataHeader: os.unlink(filename) def test_header_contains_nuclide_info(self): - from beta_spectrum.spectrum import ( - BetaSpectrum, - BetaSpectrumAnalyzer, - SpectrumConfig, - ) + from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig + from beta_spectrum.analyzer import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, @@ -222,11 +213,8 @@ class TestCSVMetadataHeader: os.unlink(filename) def test_header_contains_git_commit(self): - from beta_spectrum.spectrum import ( - BetaSpectrum, - BetaSpectrumAnalyzer, - SpectrumConfig, - ) + from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig + from beta_spectrum.analyzer import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, @@ -252,11 +240,8 @@ class TestCSVMetadataHeader: os.unlink(filename) def test_header_contains_corrections(self): - from beta_spectrum.spectrum import ( - BetaSpectrum, - BetaSpectrumAnalyzer, - SpectrumConfig, - ) + from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig + from beta_spectrum.analyzer import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, From b3d6abb4d5f2c7359a0662a5c8aef56b27e443f3 Mon Sep 17 00:00:00 2001 From: Evgeniy Unzhakov Date: Sat, 16 May 2026 21:21:28 +0300 Subject: [PATCH 2/4] =?UTF-8?q?polish:=20unify=20plot=20output=20=E2=80=94?= =?UTF-8?q?=20external=20headers,=20compact=20legends,=20consistent=20titl?= =?UTF-8?q?es?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- beta_spectrum/analyzer.py | 419 +++++++++++++------------------------- 1 file changed, 139 insertions(+), 280 deletions(-) diff --git a/beta_spectrum/analyzer.py b/beta_spectrum/analyzer.py index ac852f4..6ed5530 100644 --- a/beta_spectrum/analyzer.py +++ b/beta_spectrum/analyzer.py @@ -18,7 +18,7 @@ from beta_spectrum.logging_utils import get_git_short_hash from beta_spectrum.utils import T_to_W if TYPE_CHECKING: - from matplotlib.axes import Axes + from matplotlib.figure import Figure from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig from beta_spectrum.components.detector_response import DetectorResponse @@ -127,6 +127,50 @@ class BetaSpectrumAnalyzer: return convolved + def _build_figure_header(self, fig: "Figure") -> None: + """Add a compact header with nuclide info, e_step, commit and timestamp outside axes.""" + parent = self._element_symbol(self.config.Z_parent) + daughter = self._element_symbol(self.config.Z_daughter) + e_step_keV = self.config.e_step_MeV * 1000 + header = ( + f"{parent}{self.config.A_number} -> {daughter}{self.config.A_number} " + f"| E₀={self.config.endpoint_MeV * 1000:.1f} keV " + f"| ΔE={e_step_keV:.1f} keV " + f"| {self.config.transition_type}" + ) + fig.text( + 0.01, 0.98, header, + transform=fig.transFigure, fontsize=9, + verticalalignment="top", horizontalalignment="left", + fontfamily="monospace", + ) + id_text = f"{get_git_short_hash(6)} | {self._timestamp}" + fig.text( + 0.01, 0.01, id_text, + transform=fig.transFigure, fontsize=8, + verticalalignment="bottom", horizontalalignment="left", + fontfamily="monospace", color="gray", + ) + + def _build_figure_title(self, fig: "Figure") -> str: + """Build a unified figure title and return the figure for suptitle.""" + parent = self._element_symbol(self.config.Z_parent) + daughter = self._element_symbol(self.config.Z_daughter) + n = self.config.A_number + title = f"Beta-decay: {n}{parent} -> {n}{daughter}" + if self._is_multi and self.spectrum.branches: + title += f" ({len(self.spectrum.branches)} branches)" + return title + + def _format_branch_label(self, i: int) -> str: + """Abbreviated branch label for legend.""" + if i >= len(self.spectrum.branches): + return f"Br. {i+1}" + branch = self.spectrum.branches[i] + e0 = branch.endpoint_MeV * 1000 + intensity_pct = branch.intensity * 100 + return f"Br. {i+1}: {e0:.1f} keV ({intensity_pct:.1f}%)" + def plot_analysis( self, save_path: Optional[str] = None, @@ -151,110 +195,44 @@ class BetaSpectrumAnalyzer: """ 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") + self._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 + total, components, save_path ) else: self._plot_multi_branch_spectra( - total, components, save_path, commit, timestamp + total, components, save_path ) else: if show_components: - self._plot_debug_view(total, components, save_path, commit, timestamp) + self._plot_debug_view(total, components, save_path) 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" - ), - ) + self._plot_spectrum_only(total, save_path) 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.""" + """Plot only the total spectrum.""" 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_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=12) + ax.set_ylabel( + f"Normalized Counts / {self.config.e_step_MeV * 1000:.1f} keV", fontsize=12 + ) + fig.suptitle( + self._build_figure_title(fig), fontsize=14, fontweight="bold" ) 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) + self._build_figure_header(fig) plt.tight_layout() @@ -269,33 +247,30 @@ class BetaSpectrumAnalyzer: 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)) + fig.suptitle( + self._build_figure_title(fig), fontsize=14, fontweight="bold" + ) + + e_step_keV = self.config.e_step_MeV * 1000 + # 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.plot(self.energies_MeV, total, "b-", lw=2, label="Total spectrum") 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, + f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=10 ) ax1.set_yscale("log") ax1.grid(True, alpha=0.3) - ax1.legend() + ax1.legend(fontsize=8) # 2. Correction plots (top-right) ax2 = plt.subplot(2, 2, 2) - # Set color order for correction plots component_order = [ "PhaseSpace", "Fermi", @@ -310,58 +285,44 @@ class BetaSpectrumAnalyzer: 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)" + label = name else: label = name - ax2.plot( - self.energies_MeV, - values, - color=color, - lw=1.5, - label=label, - alpha=0.8, + 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.set_title("Spectrum components", fontsize=11) ax2.grid(True, alpha=0.3) ax2.legend(loc="best", fontsize=8) - # 3. Cumulative effect (botom-left) + # 3. Cumulative effect (bottom-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 + self.energies_MeV, np.ones_like(cumulative), "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, + 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.set_title("Cumulative effect", fontsize=11) ax3.grid(True, alpha=0.3) ax3.legend(loc="best", fontsize=8) @@ -372,26 +333,19 @@ class BetaSpectrumAnalyzer: 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, + self.energies_MeV, deviation, color=color, lw=1.5, + label=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.set_title("Correction deviations (C - 1)", fontsize=11) 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", - ) + self._build_figure_header(fig) + plt.tight_layout() if save_path: @@ -510,40 +464,21 @@ class BetaSpectrumAnalyzer: # 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. - """ + """Plot total spectrum with individual branch decomposition.""" fig, ax = plt.subplots(figsize=(12, 7)) branch_colors = [ - "#1f77b4", - "#ff7f0e", - "#2ca02c", - "#d62728", - "#9467bd", - "#8c564b", - "#e377c2", - "#7f7f7f", - "#bcbd22", - "#17becf", + "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", + "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", + "#bcbd22", "#17becf", ] # Compute branch contributions (universal * intensity * branch_spectrum) @@ -551,11 +486,8 @@ class BetaSpectrumAnalyzer: 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", + "Fermi", "Screening", "FiniteSizeL0", + "ChargeDistributionU", "Exchange", ]: if name in components: universal_product *= components[name] @@ -563,46 +495,35 @@ class BetaSpectrumAnalyzer: 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, + 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) - ): + for i, contrib in enumerate(branch_contribs): 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), + self.energies_MeV, contrib, color=color, lw=1.5, + alpha=0.7, label=self._format_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_ylabel( + f"Normalized Counts / {self.config.e_step_MeV * 1000:.1f} keV", fontsize=12 + ) + fig.suptitle( + self._build_figure_title(fig), fontsize=14, fontweight="bold" ) ax.set_yscale("log") ax.grid(True, alpha=0.3) - ax.legend(loc="best", fontsize=10, ncol=min(2, len(self.spectrum.branches) + 1)) + ax.legend( + loc="best", fontsize=9, + ncol=min(2, len(self.spectrum.branches) + 1), + ) - self._add_nuclear_data_header(ax) - self._add_id_textbox(ax, commit, timestamp) + self._build_figure_header(fig) plt.tight_layout() @@ -615,46 +536,29 @@ class BetaSpectrumAnalyzer: 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) - """ + """Debug view for multi-branch: vertical layout with 4 panels.""" n_branches = len(self.spectrum.branches) + e_step_keV = self.config.e_step_MeV * 1000 fig, axes = plt.subplots(4, 1, figsize=(12, 16), sharex=True) + fig.suptitle( + self._build_figure_title(fig), fontsize=14, fontweight="bold" + ) branch_colors = [ - "#1f77b4", - "#ff7f0e", - "#2ca02c", - "#d62728", - "#9467bd", - "#8c564b", - "#e377c2", - "#7f7f7f", - "#bcbd22", - "#17becf", + "#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", + "Fermi", "Screening", "FiniteSizeL0", + "ChargeDistributionU", "Exchange", ]: if name in components: universal_product *= components[name] @@ -666,59 +570,36 @@ class BetaSpectrumAnalyzer: ax = axes[0] ax.plot( - self.energies_MeV, - total_raw, - "k-", - lw=2, - label="Total spectrum", - zorder=10, + 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) - ): + for i, contrib in enumerate(branch_contribs): 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", + self.energies_MeV, contrib, color=color, lw=1.5, + alpha=0.7, label=self._format_branch_label(i), ) 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.set_ylabel( + f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=10 ) + ax.set_title("Spectrum", fontsize=11) 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 + # Panel 2: Fermi function + all phase spaces 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, + self.energies_MeV, fermi_scaled, "r-", lw=1.5, + label="Fermi", alpha=0.8, ) - for i, (branch_spec, branch) in enumerate( - zip(self.spectrum.branch_spectra, self.spectrum.branches) - ): + for i, branch in enumerate(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) @@ -728,27 +609,19 @@ class BetaSpectrumAnalyzer: 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)", + self.energies_MeV, ps_scaled, color=color, lw=1.5, + alpha=0.7, label=f"PS Br. {i+1}", ) ax.set_ylabel("Normalized Factor", fontsize=10) - ax.set_title( - "Fermi Function & Phase Space Shapes", fontsize=11, fontweight="bold" - ) + ax.set_title("Fermi & Phase Space", fontsize=11) 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", + "Screening", "FiniteSizeL0", + "ChargeDistributionU", "Exchange", ] colors = ["green", "orange", "purple", "brown"] for name, color in zip(universal_components, colors): @@ -756,45 +629,31 @@ class BetaSpectrumAnalyzer: 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, + self.energies_MeV, values_scaled, color=color, lw=1.5, + label=name, alpha=0.8, ) ax.set_ylabel("Normalized Factor", fontsize=10) - ax.set_title( - "Universal W0-Independent Corrections", fontsize=11, fontweight="bold" - ) + ax.set_title("Universal Corrections", fontsize=11) 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) - ): + for i, branch in enumerate(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)", + self.energies_MeV, rad_vals, color=color, lw=1.5, + alpha=0.7, label=f"Br. {i+1}", ) 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.set_title("Radiative Corrections", fontsize=11) 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) + self._build_figure_header(fig) plt.tight_layout() From 55cfb30b2b05f4b78b9d52e2c070482f6e940fda Mon Sep 17 00:00:00 2001 From: Evgeniy Unzhakov Date: Sat, 16 May 2026 22:59:48 +0300 Subject: [PATCH 3/4] =?UTF-8?q?refactor:=20unify=20plotting=20=E2=80=94=20?= =?UTF-8?q?single-branch=20treated=20as=20multi-branch=20with=20one=20bran?= =?UTF-8?q?ch;=20remove=20old=20single-branch=20methods?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- beta_spectrum/analyzer.py | 612 ++++++++++---------------------------- 1 file changed, 154 insertions(+), 458 deletions(-) diff --git a/beta_spectrum/analyzer.py b/beta_spectrum/analyzer.py index 6ed5530..330449e 100644 --- a/beta_spectrum/analyzer.py +++ b/beta_spectrum/analyzer.py @@ -2,6 +2,8 @@ Analyzer module — BetaSpectrumAnalyzer class. Introspection, analysis, and debugging 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 @@ -9,6 +11,7 @@ from __future__ import annotations import logging import pandas as pd from datetime import datetime, timezone +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, List, Optional import numpy as np @@ -55,10 +58,33 @@ class BetaSpectrumAnalyzer: self._logger = logger self._is_multi = spectrum._is_multi + # 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: - self._components_cache = self.spectrum.calculate_components(self.W) + 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 @@ -73,9 +99,7 @@ class BetaSpectrumAnalyzer: return self._branch_spectra_cache def total_spectrum(self, normalize: bool = True) -> np.ndarray: - """ - Calculate the total_spectrum - """ + """Calculate the total spectrum.""" total = self.spectrum(self.W) if normalize: @@ -127,6 +151,42 @@ class BetaSpectrumAnalyzer: return convolved + # --------------------------------------------------------------------------- + # 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): + return self.spectrum.branch_spectra[i] + # 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] + + # --------------------------------------------------------------------------- + # Plotting helpers + # --------------------------------------------------------------------------- + def _build_figure_header(self, fig: "Figure") -> None: """Add a compact header with nuclide info, e_step, commit and timestamp outside axes.""" parent = self._element_symbol(self.config.Z_parent) @@ -153,24 +213,26 @@ class BetaSpectrumAnalyzer: ) def _build_figure_title(self, fig: "Figure") -> str: - """Build a unified figure title and return the figure for suptitle.""" + """Build a unified figure title.""" parent = self._element_symbol(self.config.Z_parent) daughter = self._element_symbol(self.config.Z_daughter) n = self.config.A_number title = f"Beta-decay: {n}{parent} -> {n}{daughter}" - if self._is_multi and self.spectrum.branches: - title += f" ({len(self.spectrum.branches)} branches)" + if self._n_branches > 1: + title += f" ({self._n_branches} branches)" return title def _format_branch_label(self, i: int) -> str: """Abbreviated branch label for legend.""" - if i >= len(self.spectrum.branches): - return f"Br. {i+1}" - branch = self.spectrum.branches[i] + branch = self._unified_branches[i] e0 = branch.endpoint_MeV * 1000 intensity_pct = branch.intensity * 100 return f"Br. {i+1}: {e0:.1f} keV ({intensity_pct:.1f}%)" + # --------------------------------------------------------------------------- + # Plotting — unified multi-branch approach + # --------------------------------------------------------------------------- + def plot_analysis( self, save_path: Optional[str] = None, @@ -179,294 +241,24 @@ class BetaSpectrumAnalyzer: """ 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. + If False, show only the spectrum plot. """ total = self.total_spectrum(normalize=True) components = self.components self._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 - ) - else: - self._plot_multi_branch_spectra( - total, components, save_path - ) + if show_components: + self._plot_debug_view(total, components, save_path) else: - if show_components: - self._plot_debug_view(total, components, save_path) - else: - self._plot_spectrum_only(total, save_path) + self._plot_spectra_view(total, components, save_path) - def _plot_spectrum_only( - self, - total: np.ndarray, - save_path: Optional[str], - ) -> None: - """Plot only the total spectrum.""" - 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=12) - ax.set_ylabel( - f"Normalized Counts / {self.config.e_step_MeV * 1000:.1f} keV", fontsize=12 - ) - fig.suptitle( - self._build_figure_title(fig), fontsize=14, fontweight="bold" - ) - ax.set_yscale("log") - ax.grid(True, alpha=0.3) - ax.legend(loc="best", fontsize=10) - - self._build_figure_header(fig) - - 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], - ) -> None: - """Full debug view with 4-panel spectrum analysis.""" - fig = plt.figure(figsize=(14, 10)) - - fig.suptitle( - self._build_figure_title(fig), fontsize=14, fontweight="bold" - ) - - e_step_keV = self.config.e_step_MeV * 1000 - - # 1. Main spectrum plot (top-left) - ax1 = plt.subplot(2, 2, 1) - ax1.plot(self.energies_MeV, total, "b-", lw=2, label="Total spectrum") - ax1.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=10) - ax1.set_ylabel( - f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=10 - ) - ax1.set_yscale("log") - ax1.grid(True, alpha=0.3) - ax1.legend(fontsize=8) - - # 2. Correction plots (top-right) - ax2 = plt.subplot(2, 2, 2) - - 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] - if name in ["Fermi", "PhaseSpace"]: - values = values / np.max(values) - label = name - 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=11) - ax2.grid(True, alpha=0.3) - ax2.legend(loc="best", fontsize=8) - - # 3. Cumulative effect (bottom-left) - ax3 = plt.subplot(2, 2, 3) - - cumulative = np.ones_like(self.W) - - ax3.plot( - self.energies_MeV, np.ones_like(cumulative), "k--", lw=1.5, - label="Baseline", alpha=0.7, - ) - - 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=11) - 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=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=11) - ax4.grid(True, alpha=0.3) - ax4.legend(loc="best", fontsize=8) - ax4.axhline(y=0, color="k", linestyle="-", lw=0.5) - - self._build_figure_header(fig) - - 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 _plot_multi_branch_spectra( + def _plot_spectra_view( self, total: np.ndarray, components: Dict[str, np.ndarray], @@ -481,18 +273,21 @@ class BetaSpectrumAnalyzer: "#bcbd22", "#17becf", ] + e_step_keV = self.config.e_step_MeV * 1000 + # Compute branch contributions (universal * intensity * branch_spectrum) branch_contribs = [] - for i, branch_spec in enumerate(self.spectrum.branch_spectra): + for i in range(self._n_branches): 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] + key = self._get_component_key(i, name) + if key in components: + universal_product *= components[key] + bs = self._get_branch_spectrum(i) + contrib = universal_product * self._unified_branches[i].intensity * bs branch_contribs.append(contrib) total_raw = sum(branch_contribs) @@ -511,7 +306,7 @@ class BetaSpectrumAnalyzer: ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=12) ax.set_ylabel( - f"Normalized Counts / {self.config.e_step_MeV * 1000:.1f} keV", fontsize=12 + f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=12 ) fig.suptitle( self._build_figure_title(fig), fontsize=14, fontweight="bold" @@ -520,7 +315,7 @@ class BetaSpectrumAnalyzer: ax.grid(True, alpha=0.3) ax.legend( loc="best", fontsize=9, - ncol=min(2, len(self.spectrum.branches) + 1), + ncol=min(2, self._n_branches + 1), ) self._build_figure_header(fig) @@ -531,14 +326,14 @@ class BetaSpectrumAnalyzer: fig.savefig(save_path, dpi=150, bbox_inches="tight") print(f"Figure saved to {save_path}") - def _plot_multi_branch_debug( + def _plot_debug_view( self, total: np.ndarray, components: Dict[str, np.ndarray], save_path: Optional[str], ) -> None: - """Debug view for multi-branch: vertical layout with 4 panels.""" - n_branches = len(self.spectrum.branches) + """Debug view: vertical layout with 4 panels.""" + n_branches = self._n_branches e_step_keV = self.config.e_step_MeV * 1000 fig, axes = plt.subplots(4, 1, figsize=(12, 16), sharex=True) @@ -554,16 +349,17 @@ class BetaSpectrumAnalyzer: # Panel 1: Total spectrum + per-branch spectra branch_contribs = [] - for i, branch_spec in enumerate(self.spectrum.branch_spectra): + for i in range(n_branches): 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] + key = self._get_component_key(i, name) + if key in components: + universal_product *= components[key] + bs = self._get_branch_spectrum(i) + contrib = universal_product * self._unified_branches[i].intensity * bs branch_contribs.append(contrib) total_raw = sum(branch_contribs) @@ -596,10 +392,12 @@ class BetaSpectrumAnalyzer: self.energies_MeV, fermi_scaled, "r-", lw=1.5, label="Fermi", alpha=0.8, ) - for i, branch in enumerate(self.spectrum.branches): + for i in range(n_branches): color = branch_colors[i % len(branch_colors)] - if f"branch_{i}.PhaseSpace" in components: - ps_vals = components[f"branch_{i}.PhaseSpace"] + branch = self._unified_branches[i] + key = self._get_component_key(i, "PhaseSpace") + if key in components: + ps_vals = components[key] W0_branch = T_to_W(branch.endpoint_MeV) mask = self.W <= W0_branch ps_masked = np.where(mask, ps_vals, 0.0) @@ -639,10 +437,11 @@ class BetaSpectrumAnalyzer: # Panel 4: W0-dependent corrections (Radiative for each branch) ax = axes[3] - for i, branch in enumerate(self.spectrum.branches): + for i in range(n_branches): color = branch_colors[i % len(branch_colors)] - if f"branch_{i}.Radiative" in components: - rad_vals = components[f"branch_{i}.Radiative"] + key = self._get_component_key(i, "Radiative") + if key in components: + rad_vals = components[key] ax.plot( self.energies_MeV, rad_vals, color=color, lw=1.5, alpha=0.7, label=f"Br. {i+1}", @@ -661,21 +460,31 @@ class BetaSpectrumAnalyzer: 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. + 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}" - Parameters - ---------- - filename : str - Output CSV file path. - source_type : str - Data source type (paceENSDF, json, cli) for the header. - """ + # --------------------------------------------------------------------------- + # Export + # --------------------------------------------------------------------------- + + def export_to_csv(self, filename: str, source_type: str = "unknown") -> None: + """Export spectrum data to CSV with metadata 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") @@ -692,115 +501,8 @@ class BetaSpectrumAnalyzer: if self.config.use_exchange: enabled.append("exchange") - # Build metadata header - 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) + parent_symbol = self._element_symbol(self.config.Z_parent) + daughter_symbol = self._element_symbol(self.config.Z_daughter) header_lines = [ f"# beta-spectrum v{__import__('beta_spectrum').__version__}", @@ -813,16 +515,13 @@ class BetaSpectrumAnalyzer: 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)") + # Branch info — unified view + for i, branch in enumerate(self._unified_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}" + ) if self.config.use_detector_response: header_lines.append( @@ -835,45 +534,34 @@ class BetaSpectrumAnalyzer: header_lines.append(f"# git_commit: {get_git_short_hash()}") - # Write CSV with header 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 spectra and components + for i in range(self._n_branches): + data[f"branch_{i+1}_spectrum"] = self._get_branch_spectrum(i) + for comp_name in ["PhaseSpace", "Radiative"]: + key = self._get_component_key(i, comp_name) + if key in components: + data[f"branch_{i+1}_{comp_name}"] = components[key] - # 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 (same for all branches — no prefix needed) - for name, values in components.items(): - if not name.startswith("branch_"): - data[name] = values - else: - # Single-branch mode: standard export - for name, values in components.items(): + # Universal components (no prefix) + for name, values in components.items(): + if not name.startswith("branch_"): 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) + self._logger.info("CSV exported to %s", filename) def get_data(self) -> Dict[str, Any]: """Get all numerical data for custom analysis.""" @@ -884,3 +572,11 @@ class BetaSpectrumAnalyzer: "components": self.components, "config": self.config, } + + +@dataclass +class _UnifiedBranch: + """Unified branch view for both single- and multi-branch modes.""" + endpoint_MeV: float + transition_type: str + intensity: float From 64785f4aa981ea85516eca7194231b88de35623e Mon Sep 17 00:00:00 2001 From: Evgeniy Unzhakov Date: Sat, 16 May 2026 23:24:33 +0300 Subject: [PATCH 4/4] =?UTF-8?q?style:=20remove=20transition=20type=20from?= =?UTF-8?q?=20plot=20header=20=E2=80=94=20shown=20in=20branch=20legends=20?= =?UTF-8?q?instead?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- beta_spectrum/analyzer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/beta_spectrum/analyzer.py b/beta_spectrum/analyzer.py index 330449e..c9c9b1a 100644 --- a/beta_spectrum/analyzer.py +++ b/beta_spectrum/analyzer.py @@ -195,8 +195,7 @@ class BetaSpectrumAnalyzer: header = ( f"{parent}{self.config.A_number} -> {daughter}{self.config.A_number} " f"| E₀={self.config.endpoint_MeV * 1000:.1f} keV " - f"| ΔE={e_step_keV:.1f} keV " - f"| {self.config.transition_type}" + f"| ΔE={e_step_keV:.1f} keV" ) fig.text( 0.01, 0.98, header,