diff --git a/beta_spectrum/__init__.py b/beta_spectrum/__init__.py index 439eb56..574ef80 100644 --- a/beta_spectrum/__init__.py +++ b/beta_spectrum/__init__.py @@ -11,7 +11,7 @@ from .components.exchange import ExchangeCorrection from .components.radiative import RadiativeCorrection from .spectrum import BetaSpectrum, SpectrumConfig, BranchConfig -from .visualize import BetaSpectrumAnalyzer +from .output import BetaSpectrumAnalyzer from .nuclear_data import ( DecayInfo, BranchInfo, diff --git a/beta_spectrum/cli.py b/beta_spectrum/cli.py index 5eb0004..476207d 100644 --- a/beta_spectrum/cli.py +++ b/beta_spectrum/cli.py @@ -155,7 +155,7 @@ def _build_parser() -> argparse.ArgumentParser: log_group.add_argument( "--csv-separator", type=str, - default=',', + default=",", help="CSV column separator (default: ','). Use 'tab' for tab separation.", ) log_group.add_argument( @@ -183,7 +183,7 @@ def _run(args: argparse.Namespace) -> None: BranchConfig, SpectrumConfig, ) - from beta_spectrum.visualize import BetaSpectrumAnalyzer + from beta_spectrum.output import BetaSpectrumAnalyzer # Determine logging level if args.quiet: diff --git a/beta_spectrum/nuclear_data.py b/beta_spectrum/nuclear_data.py index 205ba53..fe69537 100644 --- a/beta_spectrum/nuclear_data.py +++ b/beta_spectrum/nuclear_data.py @@ -471,7 +471,7 @@ def decay_info_to_config( e_step_MeV: float = 0.001, intensity_cutoff: float = 0.0, neutrino_mass_MeV: float = 0.0, - csv_separator: str = ',', + csv_separator: str = ",", normalize_to_one: bool = False, ) -> SpectrumConfig: """ @@ -698,7 +698,7 @@ def create_config_from_source( e_step = kwargs.get("e_step_MeV", 0.001) intensity_cutoff = kwargs.get("intensity_cutoff", 0.0) neutrino_mass = kwargs.get("neutrino_mass_MeV", 0.0) - csv_separator = kwargs.get("csv_separator", ',') + csv_separator = kwargs.get("csv_separator", ",") normalize_to_one = kwargs.get("normalize_to_one", False) info = get_decay_info_from_paceENSDF(nuclide, decay_type, decay_index) return decay_info_to_config( diff --git a/beta_spectrum/output/__init__.py b/beta_spectrum/output/__init__.py new file mode 100644 index 0000000..0cc56d8 --- /dev/null +++ b/beta_spectrum/output/__init__.py @@ -0,0 +1,23 @@ +""" +Output module — BetaSpectrumAnalyzer class and related functionality. + +Introspection, analysis, visualization and export 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 + +from .core import BetaSpectrumAnalyzer, _UnifiedBranch +from .csv_export import CSVEmitter +from .visualization import Visualizer + +# Re-export the main class for backward compatibility +BetaSpectrumAnalyzer = BetaSpectrumAnalyzer + +__all__ = [ + "BetaSpectrumAnalyzer", + "CSVEmitter", + "Visualizer", + "_UnifiedBranch", +] diff --git a/beta_spectrum/output/core.py b/beta_spectrum/output/core.py new file mode 100644 index 0000000..ea011ab --- /dev/null +++ b/beta_spectrum/output/core.py @@ -0,0 +1,285 @@ +""" +Core analyzer module — BetaSpectrumAnalyzer base class. + +Core introspection and analysis tools for BetaSpectrum. +All nuclides (single- and multi-branch) are treated uniformly as one or more +branches — single-branch is just a special case with one branch. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import numpy as np + +from beta_spectrum.utils import T_to_W + +if TYPE_CHECKING: + from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig + + +class BetaSpectrumAnalyzer: + """ + Introspection, analysis and debugging tools for BetaSpectrum. + + Parameters + ---------- + spectrum : BetaSpectrum + The spectrum to analyze. + config : SpectrumConfig + Configuration used to create the spectrum. + logger : logging.Logger, optional + Logger for progress output. + """ + + def __init__( + self, + spectrum: BetaSpectrum, + config: SpectrumConfig, + logger: Optional[logging.Logger] = None, + ): + """ + Initialize analyzer with a spectrum and configuration. + """ + self.spectrum = spectrum + self.config = config + self.W, self.energies_MeV = spectrum.get_energy_grid(config) + self._components_cache: Optional[Dict[str, np.ndarray]] = None + self._branch_spectra_cache: Optional[List[np.ndarray]] = None + self._logger = logger + self._is_multi = spectrum._is_multi + self._timestamp: Optional[str] = None + + # Unified branch view: single-branch treated as one-branch multi-branch + if not self._is_multi and self.spectrum.branch_spectra == []: + self._unified_branches: List[_UnifiedBranch] = [ + _UnifiedBranch( + endpoint_MeV=config.endpoint_MeV, + transition_type=config.transition_type, + intensity=1.0, + ) + ] + else: + self._unified_branches = [ + _UnifiedBranch( + endpoint_MeV=b.endpoint_MeV, + transition_type=b.transition_type, + intensity=b.intensity, + ) + for b in self.spectrum.branches + ] + + @property + def components(self) -> Dict[str, np.ndarray]: + if self._components_cache is None: + raw = self.spectrum.calculate_components(self.W) + # Normalize keys: single-branch has bare names, + # multi-branch has branch_0., branch_1., etc. + # We keep the raw keys for export_to_csv compatibility. + self._components_cache = raw + assert self._components_cache is not None + return self._components_cache + + @property + def branch_spectra(self) -> List[np.ndarray]: + """Intensity-weighted branch spectra.""" + if self._branch_spectra_cache is None: + self._branch_spectra_cache = self.spectrum.get_branch_normalized_spectra( + self.W + ) + assert self._branch_spectra_cache is not None + return self._branch_spectra_cache + + def total_spectrum(self, normalize: bool = True) -> np.ndarray: + """Calculate the total spectrum.""" + total = self.spectrum(self.W) + + if normalize: + integral = np.trapezoid(total, self.energies_MeV) + total = total / integral + + return total + + # --------------------------------------------------------------------------- + # Unified branch view + # --------------------------------------------------------------------------- + + @property + def _n_branches(self) -> int: + """Number of branches (unified view).""" + return len(self._unified_branches) + + def _get_branch_spectrum(self, i: int) -> np.ndarray: + """Get the spectrum array for branch i.""" + if i < len(self.spectrum.branch_spectra): + bs_obj = self.spectrum.branch_spectra[i] + # branch_spectra stores BetaSpectrum objects — call them to get arrays + if hasattr(bs_obj, "__call__"): + val = bs_obj(self.W) + # Clamp to zero beyond endpoint (phase space goes negative) + W0 = T_to_W(self._unified_branches[i].endpoint_MeV) + return np.where(self.W <= W0, val, 0.0) + return np.asarray(bs_obj) + # Single-branch: the spectrum itself is the only branch + return self.spectrum(self.W) + + def _get_component_key(self, i: int, name: str) -> str: + """Get the canonical key for a per-branch component.""" + if self._is_multi: + return f"branch_{i}.{name}" + return name + + def _has_component(self, i: int, name: str) -> bool: + """Check if branch i has a given component.""" + key = self._get_component_key(i, name) + return key in self.components + + def _get_component(self, i: int, name: str) -> np.ndarray: + """Get a per-branch component array.""" + key = self._get_component_key(i, name) + return self.components[key] + + def _element_symbol(self, Z: int) -> str: + """Return chemical symbol for atomic number Z.""" + symbols = [ + "", + "H", + "He", + "Li", + "Be", + "B", + "C", + "N", + "O", + "F", + "Ne", + "Na", + "Mg", + "Al", + "Si", + "P", + "S", + "Cl", + "Ar", + "K", + "Ca", + "Sc", + "Ti", + "V", + "Cr", + "Mn", + "Fe", + "Co", + "Ni", + "Cu", + "Zn", + "Ga", + "Ge", + "As", + "Se", + "Br", + "Kr", + "Rb", + "Sr", + "Y", + "Zr", + "Nb", + "Mo", + "Tc", + "Ru", + "Rh", + "Pd", + "Ag", + "Cd", + "In", + "Sn", + "Sb", + "Te", + "I", + "Xe", + "Cs", + "Ba", + "La", + "Ce", + "Pr", + "Nd", + "Pm", + "Sm", + "Eu", + "Gd", + "Tb", + "Dy", + "Ho", + "Er", + "Tm", + "Yb", + "Lu", + "Hf", + "Ta", + "W", + "Re", + "Os", + "Ir", + "Pt", + "Au", + "Hg", + "Tl", + "Pb", + "Bi", + "Po", + "At", + "Rn", + "Fr", + "Ra", + "Ac", + "Th", + "Pa", + "U", + "Np", + "Pu", + "Am", + "Cm", + "Bk", + "Cf", + "Es", + ] + return symbols[Z] if Z < len(symbols) else f"Z{Z}" + + def get_data(self) -> Dict[str, Any]: + """Get all numerical data for custom analysis.""" + return { + "energies_MeV": self.energies_MeV, + "energies_W": self.W, + "spectrum": self.total_spectrum(normalize=True), + "components": self.components, + "config": self.config, + } + + def export_to_csv( + self, filename: str, source_type: str = "unknown", csv_verbosity: str = "full" + ) -> None: + """Export spectrum data to CSV with metadata header.""" + from .csv_export import CSVEmitter + + CSVEmitter.export_to_csv(self, filename, source_type, csv_verbosity) + + def plot_analysis( + self, + save_path: Optional[str] = None, + show_components: bool = True, + ) -> None: + """Create visualization of the spectrum and all correction factors.""" + from .visualization import Visualizer + + Visualizer.plot_analysis(self, save_path, show_components) + + +@dataclass +class _UnifiedBranch: + """Unified branch view for both single- and multi-branch modes.""" + + endpoint_MeV: float + transition_type: str + intensity: float diff --git a/beta_spectrum/output/csv_export.py b/beta_spectrum/output/csv_export.py new file mode 100644 index 0000000..6860dfb --- /dev/null +++ b/beta_spectrum/output/csv_export.py @@ -0,0 +1,219 @@ +""" +CSV export module — CSV export functionality for BetaSpectrumAnalyzer. +""" + +from __future__ import annotations + +import pandas as pd +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Dict + +import numpy as np + +from beta_spectrum.logging_utils import get_git_short_hash + +if TYPE_CHECKING: + from beta_spectrum.output.core import BetaSpectrumAnalyzer + + +class CSVEmitter: + """CSV export functionality for BetaSpectrumAnalyzer.""" + + @staticmethod + def export_to_csv( + analyzer: BetaSpectrumAnalyzer, + filename: str, + source_type: str = "unknown", + csv_verbosity: str = "full", + ) -> None: + """Export spectrum data to CSV with metadata header. + + Parameters + ---------- + analyzer : BetaSpectrumAnalyzer + The analyzer instance to export data from. + filename : str + Output CSV file path. + source_type : str + Source type (paceENSDF, json, etc.). + csv_verbosity : str + Verbosity level: 'full' (all components) or 'concise' (total spectrum and normalized branch components only). + """ + total = analyzer.total_spectrum(normalize=True) + components = analyzer.components + + # Only include enabled corrections that are not basic components + # Basic components (phase_space and fermi) are always used and should not be listed + enabled = [] + if analyzer.config.use_screening: + enabled.append("screening") + if analyzer.config.use_finite_size: + enabled.append("finite_size") + if analyzer.config.use_charge_dist: + enabled.append("charge_dist") + if analyzer.config.use_radiative: + enabled.append("radiative") + if analyzer.config.use_exchange: + enabled.append("exchange") + + parent_symbol = analyzer._element_symbol(analyzer.config.Z_parent) + daughter_symbol = analyzer._element_symbol(analyzer.config.Z_daughter) + + # Get git commit hash + git_commit = get_git_short_hash() + + # Format source line with paceENSDF URL if applicable + if source_type == "paceENSDF": + source_line = "# source: paceENSDF (https://pypi.org/project/paceENSDF/)" + else: + source_line = f"# source: {source_type}" + + header_lines = [ + f"# beta-spectrum v{__import__('beta_spectrum').__version__}(build.{git_commit})", + f"# timestamp: {datetime.now(timezone.utc).isoformat()}", + source_line, + f"# nuclide: {parent_symbol}{analyzer.config.A_number} -> {daughter_symbol}{analyzer.config.A_number} (Z={analyzer.config.Z_parent}->{analyzer.config.Z_daughter})", + f"# basic_endpoint: {analyzer.config.endpoint_MeV * 1000:.1f} keV", + ] + + # Add neutrino mass info if non-zero + if analyzer.config.neutrino_mass_MeV > 0: + header_lines.append( + f"# neutrino_mass_MeV: {analyzer.config.neutrino_mass_MeV:.6f} (non-zero neutrino mass effects included)" + ) + else: + header_lines.append( + f"# neutrino_mass_MeV: {analyzer.config.neutrino_mass_MeV:.6f}" + ) + + header_lines.append(f"# corrections: {', '.join(enabled)}") + header_lines.append(f"# e_step: {analyzer.config.e_step_MeV:.4f} MeV") + header_lines.append(f"# n_points: {len(analyzer.energies_MeV)}") + + # Add normalization mode + if analyzer.config.normalize_to_one: + header_lines.append("# branch normalization: to_one (ignoring intensity)") + else: + header_lines.append("# branch normalization: intensity") + + # Separate branches into active and excluded lists + active_branches = [] + excluded_branches = [] + + for i, branch in enumerate(analyzer._unified_branches): + branch_endpoint_keV = branch.endpoint_MeV * 1000 + neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000 + + # Check if branch is excluded due to neutrino mass + if branch_endpoint_keV <= neutrino_mass_keV: + excluded_branches.append( + f" branch_{i + 1}: E₀={branch_endpoint_keV:.1f} keV, transition={branch.transition_type}, intensity={branch.intensity:.4f} - excluded: endpoint < m_nu" + ) + # Check if branch is excluded due to intensity cutoff + elif branch.intensity < analyzer.config.intensity_cutoff: + excluded_branches.append( + f" branch_{i + 1}: E₀={branch_endpoint_keV:.1f} keV, transition={branch.transition_type}, intensity={branch.intensity:.4f} - excluded: intensity < cutoff" + ) + else: + active_branches.append( + f" branch_{i + 1}: E₀={branch_endpoint_keV:.1f} keV, transition={branch.transition_type}, intensity={branch.intensity:.4f}" + ) + + # Add active branches list + if active_branches: + header_lines.append("# active branches:") + for branch_info in active_branches: + header_lines.append(f"# {branch_info}") + + # Add excluded branches list if there are any + if excluded_branches: + header_lines.append("# excluded branches:") + for branch_info in excluded_branches: + header_lines.append(f"# {branch_info}") + + # Create data dictionary with concise column headers + data: Dict[str, np.ndarray] = { + "energy_MeV": analyzer.energies_MeV, + } + + # Only add total_spectrum if not normalizing to one + if not analyzer.config.normalize_to_one: + data["total_spectrum"] = total + + if csv_verbosity == "full": + # Per-branch spectra and components with concise branch naming + for i in range(analyzer._n_branches): + # Get branch spectrum + branch_spectrum = analyzer._get_branch_spectrum(i) + + if analyzer.config.normalize_to_one: + # Normalize to 1, ignoring intensity + sum_spectrum = np.sum(branch_spectrum) + if sum_spectrum > 0: + normalized_spectrum = branch_spectrum / sum_spectrum + else: + normalized_spectrum = branch_spectrum + else: + # Use original spectrum with intensity + normalized_spectrum = ( + branch_spectrum * analyzer._unified_branches[i].intensity + ) + + data[f"branch_{i + 1}_spectrum"] = normalized_spectrum + + for comp_name in ["PhaseSpace", "Radiative"]: + key = analyzer._get_component_key(i, comp_name) + if key in components: + data[f"branch_{i + 1}_{comp_name}"] = components[key] + + # Universal components (no prefix) + for name, values in components.items(): + if not name.startswith("branch_"): + data[name] = values + elif csv_verbosity == "concise": + # Only total spectrum and normalized branch components with concise naming + for i in range(analyzer._n_branches): + branch_spectrum = analyzer._get_branch_spectrum(i) + # Check if branch is degenerate (neutrino mass >= branch endpoint) + branch_endpoint_keV = analyzer._unified_branches[i].endpoint_MeV * 1000 + neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000 + if branch_endpoint_keV <= neutrino_mass_keV: + # Branch is degenerate, spectrum should be all zeros + data[f"branch_{i + 1}"] = np.zeros_like(branch_spectrum) + else: + # Branch is valid, normalize as usual + if analyzer.config.normalize_to_one: + # Normalize to 1, ignoring intensity + sum_spectrum = np.sum(branch_spectrum) + if sum_spectrum > 0: + data[f"branch_{i + 1}"] = branch_spectrum / sum_spectrum + else: + data[f"branch_{i + 1}"] = branch_spectrum + else: + # Normalize as usual with intensity + sum_spectrum = np.sum(branch_spectrum) + if sum_spectrum > 0: + data[f"branch_{i + 1}"] = ( + branch_spectrum + / sum_spectrum + * analyzer._unified_branches[i].intensity + ) + else: + data[f"branch_{i + 1}"] = branch_spectrum + + df = pd.DataFrame(data) + # Determine separator + separator = ( + "\t" + if analyzer.config.csv_separator == "tab" + else analyzer.config.csv_separator + ) + df.to_csv(filename, index=False, float_format="%.4e", sep=separator) + + with open(filename, "r") as f: + content = f.read() + with open(filename, "w") as f: + f.write("\n".join(header_lines) + "\n" + content) + + if analyzer._logger: + analyzer._logger.info("CSV exported to %s", filename) diff --git a/beta_spectrum/output/visualization.py b/beta_spectrum/output/visualization.py new file mode 100644 index 0000000..a349c6f --- /dev/null +++ b/beta_spectrum/output/visualization.py @@ -0,0 +1,444 @@ +""" +Visualization module — Plotting functionality for BetaSpectrumAnalyzer. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Dict, 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.figure import Figure + + from beta_spectrum.output.core import BetaSpectrumAnalyzer + + +class Visualizer: + """Visualization functionality for BetaSpectrumAnalyzer.""" + + @staticmethod + def plot_analysis( + analyzer: BetaSpectrumAnalyzer, + save_path: Optional[str] = None, + show_components: bool = True, + ) -> None: + """ + Create visualization of the spectrum and all correction factors. + + Parameters + ---------- + analyzer : BetaSpectrumAnalyzer + The analyzer instance to visualize. + 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. + """ + total = analyzer.total_spectrum(normalize=True) + components = analyzer.components + analyzer._timestamp = datetime.now(timezone.utc).strftime( + "%Y-%m-%d %H:%M:%S UTC" + ) + + if show_components: + Visualizer._plot_debug_view(analyzer, total, components, save_path) + else: + Visualizer._plot_spectra_view(analyzer, total, components, save_path) + + @staticmethod + def _build_figure_header(analyzer: BetaSpectrumAnalyzer, fig: "Figure") -> None: + """Add a compact header with nuclide info, e_step, commit and timestamp outside axes.""" + parent = analyzer._element_symbol(analyzer.config.Z_parent) + daughter = analyzer._element_symbol(analyzer.config.Z_daughter) + e_step_keV = analyzer.config.e_step_MeV * 1000 + header = ( + f"{parent}{analyzer.config.A_number} -> {daughter}{analyzer.config.A_number} " + f"| E₀={analyzer.config.endpoint_MeV * 1000:.1f} keV " + f"| ΔE={e_step_keV:.1f} keV" + ) + fig.text( + 0.01, + 0.98, + header, + transform=fig.transFigure, + fontsize=9, + verticalalignment="top", + horizontalalignment="left", + fontfamily="monospace", + ) + + # Add neutrino mass info to header if non-zero + if analyzer.config.neutrino_mass_MeV > 0: + nu_mass_text = f"| m_ν={analyzer.config.neutrino_mass_MeV:.6f} MeV" + fig.text( + 0.01, + 0.96, + nu_mass_text, + transform=fig.transFigure, + fontsize=8, + verticalalignment="top", + horizontalalignment="left", + fontfamily="monospace", + color="red", + ) + + id_text = f"{get_git_short_hash(6)} | {analyzer._timestamp}" + fig.text( + 0.01, + 0.01, + id_text, + transform=fig.transFigure, + fontsize=8, + verticalalignment="bottom", + horizontalalignment="left", + fontfamily="monospace", + color="gray", + ) + + @staticmethod + def _build_figure_title(analyzer: BetaSpectrumAnalyzer, fig: "Figure") -> str: + """Build a unified figure title.""" + parent = analyzer._element_symbol(analyzer.config.Z_parent) + daughter = analyzer._element_symbol(analyzer.config.Z_daughter) + n = analyzer.config.A_number + title = f"Beta-decay: {n}{parent} -> {n}{daughter}" + if analyzer._n_branches > 1: + # Count active branches (not excluded by neutrino mass or intensity cutoff) + active_branches = 0 + for i, branch in enumerate(analyzer._unified_branches): + branch_endpoint_keV = branch.endpoint_MeV * 1000 + neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000 + is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV + is_excluded_by_intensity = ( + branch.intensity < analyzer.config.intensity_cutoff + ) + if not is_excluded_by_mass and not is_excluded_by_intensity: + active_branches += 1 + + if active_branches < analyzer._n_branches: + title += f" ({active_branches} of {analyzer._n_branches} branches)" + else: + title += f" ({analyzer._n_branches} branches)" + return title + + @staticmethod + def _format_branch_label(analyzer: BetaSpectrumAnalyzer, i: int) -> str: + """Abbreviated branch label for legend.""" + branch = analyzer._unified_branches[i] + e0 = branch.endpoint_MeV * 1000 + intensity_pct = branch.intensity * 100 + + # Check if branch is excluded + branch_endpoint_keV = branch.endpoint_MeV * 1000 + neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000 + is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV + is_excluded_by_intensity = branch.intensity < analyzer.config.intensity_cutoff + + if is_excluded_by_mass: + return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%) - excluded (endpoint < m_nu)" + elif is_excluded_by_intensity: + return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%) - excluded (intensity < cutoff)" + else: + return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%)" + + @staticmethod + def _plot_spectra_view( + analyzer: BetaSpectrumAnalyzer, + total: np.ndarray, + components: Dict[str, np.ndarray], + save_path: Optional[str], + ) -> None: + """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", + ] + + e_step_keV = analyzer.config.e_step_MeV * 1000 + + # Compute branch contributions (universal * intensity * branch_spectrum) + branch_contribs = [] + # Universal components are shared (no per-branch key prefix) + universal_names = [ + "Fermi", + "Screening", + "FiniteSizeL0", + "ChargeDistributionU", + "Exchange", + ] + for i in range(analyzer._n_branches): + universal_product = np.ones_like(analyzer.W) + for name in universal_names: + key = analyzer._get_component_key(i, name) + if key in components: + universal_product *= components[key] + elif name in components: # fallback: universal components + universal_product *= components[name] + bs = analyzer._get_branch_spectrum(i) + contrib = universal_product * analyzer._unified_branches[i].intensity * bs + branch_contribs.append(contrib) + + total_raw = sum(branch_contribs) + + ax.plot( + analyzer.energies_MeV, + total_raw, + "k-", + lw=2.5, + label="Total spectrum", + zorder=10, + ) + + # Only plot active branches (not excluded by neutrino mass or intensity cutoff) + active_branch_indices = [] + for i in range(analyzer._n_branches): + branch_endpoint_keV = analyzer._unified_branches[i].endpoint_MeV * 1000 + neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000 + is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV + is_excluded_by_intensity = ( + analyzer._unified_branches[i].intensity + < analyzer.config.intensity_cutoff + ) + if not is_excluded_by_mass and not is_excluded_by_intensity: + active_branch_indices.append(i) + + for i in active_branch_indices: + contrib = branch_contribs[i] + color = branch_colors[i % len(branch_colors)] + ax.plot( + analyzer.energies_MeV, + contrib, + color=color, + lw=1.5, + alpha=0.7, + label=Visualizer._format_branch_label(analyzer, i), + ) + + ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=12) + ax.set_ylabel(f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=12) + fig.suptitle( + Visualizer._build_figure_title(analyzer, fig), + fontsize=14, + fontweight="bold", + ) + ax.set_yscale("log") + ax.grid(True, alpha=0.3) + ax.legend( + loc="best", + fontsize=9, + ncol=min(2, len(active_branch_indices) + 1), + ) + + Visualizer._build_figure_header(analyzer, fig) + + plt.tight_layout() + + if save_path: + fig.savefig(save_path, dpi=150, bbox_inches="tight") + print(f"Figure saved to {save_path}") + + @staticmethod + def _plot_debug_view( + analyzer: BetaSpectrumAnalyzer, + total: np.ndarray, + components: Dict[str, np.ndarray], + save_path: Optional[str], + ) -> None: + """Debug view: vertical layout with 4 panels.""" + n_branches = analyzer._n_branches + e_step_keV = analyzer.config.e_step_MeV * 1000 + + fig, axes = plt.subplots(4, 1, figsize=(12, 16), sharex=True) + fig.suptitle( + Visualizer._build_figure_title(analyzer, fig), + fontsize=14, + fontweight="bold", + ) + + branch_colors = [ + "#1f77b4", + "#ff7f0e", + "#2ca02c", + "#d62728", + "#9467bd", + "#8c564b", + "#e377c2", + "#7f7f7f", + "#bcbd22", + "#17becf", + ] + + # Panel 1: Total spectrum + per-branch spectra + branch_contribs = [] + universal_names = [ + "Fermi", + "Screening", + "FiniteSizeL0", + "ChargeDistributionU", + "Exchange", + ] + for i in range(n_branches): + universal_product = np.ones_like(analyzer.W) + for name in universal_names: + key = analyzer._get_component_key(i, name) + if key in components: + universal_product *= components[key] + elif name in components: + universal_product *= components[name] + bs = analyzer._get_branch_spectrum(i) + contrib = universal_product * analyzer._unified_branches[i].intensity * bs + branch_contribs.append(contrib) + + total_raw = sum(branch_contribs) + + ax = axes[0] + ax.plot( + analyzer.energies_MeV, + total_raw, + "k-", + lw=2, + label="Total spectrum", + zorder=10, + ) + + # Only plot active branches (not excluded by neutrino mass or intensity cutoff) + active_branch_indices = [] + for i in range(n_branches): + branch_endpoint_keV = analyzer._unified_branches[i].endpoint_MeV * 1000 + neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000 + is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV + is_excluded_by_intensity = ( + analyzer._unified_branches[i].intensity + < analyzer.config.intensity_cutoff + ) + if not is_excluded_by_mass and not is_excluded_by_intensity: + active_branch_indices.append(i) + + for i in active_branch_indices: + contrib = branch_contribs[i] + color = branch_colors[i % len(branch_colors)] + ax.plot( + analyzer.energies_MeV, + contrib, + color=color, + lw=1.5, + alpha=0.7, + label=Visualizer._format_branch_label(analyzer, i), + ) + ax.set_yscale("log") + 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, len(active_branch_indices) + 1)) + + # 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( + analyzer.energies_MeV, + fermi_scaled, + "r-", + lw=1.5, + label="Fermi", + alpha=0.8, + ) + for i in active_branch_indices: + color = branch_colors[i % len(branch_colors)] + branch = analyzer._unified_branches[i] + key = analyzer._get_component_key(i, "PhaseSpace") + if key in components: + ps_vals = components[key] + W0_branch = T_to_W(branch.endpoint_MeV) + mask = analyzer.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( + analyzer.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 & 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", + ] + 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( + analyzer.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 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 in active_branch_indices: + color = branch_colors[i % len(branch_colors)] + key = analyzer._get_component_key(i, "Radiative") + if key in components: + rad_vals = components[key] + ax.plot( + analyzer.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("Radiative Corrections", fontsize=11) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8, loc="best", ncol=min(2, len(active_branch_indices) + 1)) + + Visualizer._build_figure_header(analyzer, fig) + + plt.tight_layout() + + if save_path: + fig.savefig(save_path, dpi=150, bbox_inches="tight") + print(f"Figure saved to {save_path}") diff --git a/beta_spectrum/spectrum.py b/beta_spectrum/spectrum.py index 9336779..6c107ca 100644 --- a/beta_spectrum/spectrum.py +++ b/beta_spectrum/spectrum.py @@ -87,10 +87,12 @@ class SpectrumConfig: intensity_cutoff: float = ( 0.0 # Ignore branches below this fraction (default: 0.0 = all) ) - + # CSV output options - csv_separator: str = ',' # Default separator for CSV output - normalize_to_one: bool = False # If True, normalize each branch to 1, ignoring intensity + csv_separator: str = "," # Default separator for CSV output + normalize_to_one: bool = ( + False # If True, normalize each branch to 1, ignoring intensity + ) class BetaSpectrum: diff --git a/beta_spectrum/visualize/__init__.py b/beta_spectrum/visualize/__init__.py deleted file mode 100644 index 27fb73f..0000000 --- a/beta_spectrum/visualize/__init__.py +++ /dev/null @@ -1,868 +0,0 @@ -""" -Visualize module — BetaSpectrumAnalyzer class. - -Visualization and plotting tools for BetaSpectrum. -All nuclides (single- and multi-branch) are treated uniformly as one or more -branches — single-branch is just a special case with one branch. -""" - -from __future__ import annotations - -import logging -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 -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.figure import Figure - - from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig - - -class BetaSpectrumAnalyzer: - """ - Introspection, analysis and debugging tools for BetaSpectrum. - - Parameters - ---------- - spectrum : BetaSpectrum - The spectrum to analyze. - config : SpectrumConfig - Configuration used to create the spectrum. - logger : logging.Logger, optional - Logger for progress output. - """ - - def __init__( - self, - spectrum: BetaSpectrum, - config: SpectrumConfig, - logger: Optional[logging.Logger] = None, - ): - """ - Initialize analyzer with a spectrum and configuration. - """ - self.spectrum = spectrum - self.config = config - self.W, self.energies_MeV = spectrum.get_energy_grid(config) - self._components_cache: Optional[Dict[str, np.ndarray]] = None - self._branch_spectra_cache: Optional[List[np.ndarray]] = None - self._logger = logger - self._is_multi = spectrum._is_multi - - # Unified branch view: single-branch treated as one-branch multi-branch - if not self._is_multi and self.spectrum.branch_spectra == []: - self._unified_branches: List[_UnifiedBranch] = [ - _UnifiedBranch( - endpoint_MeV=config.endpoint_MeV, - transition_type=config.transition_type, - intensity=1.0, - ) - ] - else: - self._unified_branches = [ - _UnifiedBranch( - endpoint_MeV=b.endpoint_MeV, - transition_type=b.transition_type, - intensity=b.intensity, - ) - for b in self.spectrum.branches - ] - - @property - def components(self) -> Dict[str, np.ndarray]: - if self._components_cache is None: - raw = self.spectrum.calculate_components(self.W) - # Normalize keys: single-branch has bare names, - # multi-branch has branch_0., branch_1., etc. - # We keep the raw keys for export_to_csv compatibility. - self._components_cache = raw - assert self._components_cache is not None - return self._components_cache - - @property - def branch_spectra(self) -> List[np.ndarray]: - """Intensity-weighted branch spectra.""" - if self._branch_spectra_cache is None: - self._branch_spectra_cache = self.spectrum.get_branch_normalized_spectra( - self.W - ) - assert self._branch_spectra_cache is not None - return self._branch_spectra_cache - - def total_spectrum(self, normalize: bool = True) -> np.ndarray: - """Calculate the total spectrum.""" - total = self.spectrum(self.W) - - if normalize: - integral = np.trapezoid(total, self.energies_MeV) - total = total / integral - - return total - - # --------------------------------------------------------------------------- - # Unified branch view - # --------------------------------------------------------------------------- - - @property - def _n_branches(self) -> int: - """Number of branches (unified view).""" - return len(self._unified_branches) - - def _get_branch_spectrum(self, i: int) -> np.ndarray: - """Get the spectrum array for branch i.""" - if i < len(self.spectrum.branch_spectra): - bs_obj = self.spectrum.branch_spectra[i] - # branch_spectra stores BetaSpectrum objects — call them to get arrays - if hasattr(bs_obj, "__call__"): - val = bs_obj(self.W) - # Clamp to zero beyond endpoint (phase space goes negative) - W0 = T_to_W(self._unified_branches[i].endpoint_MeV) - return np.where(self.W <= W0, val, 0.0) - return np.asarray(bs_obj) - # Single-branch: the spectrum itself is the only branch - return self.spectrum(self.W) - - def _get_component_key(self, i: int, name: str) -> str: - """Get the canonical key for a per-branch component.""" - if self._is_multi: - return f"branch_{i}.{name}" - return name - - def _has_component(self, i: int, name: str) -> bool: - """Check if branch i has a given component.""" - key = self._get_component_key(i, name) - return key in self.components - - def _get_component(self, i: int, name: str) -> np.ndarray: - """Get a per-branch component array.""" - key = self._get_component_key(i, name) - return self.components[key] - - # --------------------------------------------------------------------------- - # 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) - 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" - ) - fig.text( - 0.01, - 0.98, - header, - transform=fig.transFigure, - fontsize=9, - verticalalignment="top", - horizontalalignment="left", - fontfamily="monospace", - ) - - # Add neutrino mass info to header if non-zero - if self.config.neutrino_mass_MeV > 0: - nu_mass_text = f"| m_ν={self.config.neutrino_mass_MeV:.6f} MeV" - fig.text( - 0.01, - 0.96, - nu_mass_text, - transform=fig.transFigure, - fontsize=8, - verticalalignment="top", - horizontalalignment="left", - fontfamily="monospace", - color="red", - ) - - 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.""" - 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._n_branches > 1: - # Count active branches (not excluded by neutrino mass or intensity cutoff) - active_branches = 0 - for i, branch in enumerate(self._unified_branches): - branch_endpoint_keV = branch.endpoint_MeV * 1000 - neutrino_mass_keV = self.config.neutrino_mass_MeV * 1000 - is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV - is_excluded_by_intensity = ( - branch.intensity < self.config.intensity_cutoff - ) - if not is_excluded_by_mass and not is_excluded_by_intensity: - active_branches += 1 - - if active_branches < self._n_branches: - title += f" ({active_branches} of {self._n_branches} branches)" - else: - title += f" ({self._n_branches} branches)" - return title - - def _format_branch_label(self, i: int) -> str: - """Abbreviated branch label for legend.""" - branch = self._unified_branches[i] - e0 = branch.endpoint_MeV * 1000 - intensity_pct = branch.intensity * 100 - - # Check if branch is excluded - branch_endpoint_keV = branch.endpoint_MeV * 1000 - neutrino_mass_keV = self.config.neutrino_mass_MeV * 1000 - is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV - is_excluded_by_intensity = branch.intensity < self.config.intensity_cutoff - - if is_excluded_by_mass: - return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%) - excluded (endpoint < m_nu)" - elif is_excluded_by_intensity: - return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%) - excluded (intensity < cutoff)" - else: - 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, - show_components: bool = True, - ) -> None: - """ - Create visualization of the spectrum and all correction factors. - - 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. - """ - total = self.total_spectrum(normalize=True) - components = self.components - self._timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") - - if show_components: - self._plot_debug_view(total, components, save_path) - else: - self._plot_spectra_view(total, components, save_path) - - def _plot_spectra_view( - self, - total: np.ndarray, - components: Dict[str, np.ndarray], - save_path: Optional[str], - ) -> None: - """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", - ] - - e_step_keV = self.config.e_step_MeV * 1000 - - # Compute branch contributions (universal * intensity * branch_spectrum) - branch_contribs = [] - # Universal components are shared (no per-branch key prefix) - universal_names = [ - "Fermi", - "Screening", - "FiniteSizeL0", - "ChargeDistributionU", - "Exchange", - ] - for i in range(self._n_branches): - universal_product = np.ones_like(self.W) - for name in universal_names: - key = self._get_component_key(i, name) - if key in components: - universal_product *= components[key] - elif name in components: # fallback: universal components - universal_product *= components[name] - bs = self._get_branch_spectrum(i) - contrib = universal_product * self._unified_branches[i].intensity * bs - branch_contribs.append(contrib) - - total_raw = sum(branch_contribs) - - ax.plot( - self.energies_MeV, - total_raw, - "k-", - lw=2.5, - label="Total spectrum", - zorder=10, - ) - - # Only plot active branches (not excluded by neutrino mass or intensity cutoff) - active_branch_indices = [] - for i in range(self._n_branches): - branch_endpoint_keV = self._unified_branches[i].endpoint_MeV * 1000 - neutrino_mass_keV = self.config.neutrino_mass_MeV * 1000 - is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV - is_excluded_by_intensity = ( - self._unified_branches[i].intensity < self.config.intensity_cutoff - ) - if not is_excluded_by_mass and not is_excluded_by_intensity: - active_branch_indices.append(i) - - for i in active_branch_indices: - contrib = branch_contribs[i] - color = branch_colors[i % len(branch_colors)] - ax.plot( - 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(f"Normalized Counts / {e_step_keV:.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=9, - ncol=min(2, len(active_branch_indices) + 1), - ) - - 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}") - - def _plot_debug_view( - self, - total: np.ndarray, - components: Dict[str, np.ndarray], - save_path: Optional[str], - ) -> None: - """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) - fig.suptitle(self._build_figure_title(fig), fontsize=14, fontweight="bold") - - branch_colors = [ - "#1f77b4", - "#ff7f0e", - "#2ca02c", - "#d62728", - "#9467bd", - "#8c564b", - "#e377c2", - "#7f7f7f", - "#bcbd22", - "#17becf", - ] - - # Panel 1: Total spectrum + per-branch spectra - branch_contribs = [] - universal_names = [ - "Fermi", - "Screening", - "FiniteSizeL0", - "ChargeDistributionU", - "Exchange", - ] - for i in range(n_branches): - universal_product = np.ones_like(self.W) - for name in universal_names: - key = self._get_component_key(i, name) - if key in components: - universal_product *= components[key] - elif name in components: - universal_product *= components[name] - bs = self._get_branch_spectrum(i) - contrib = universal_product * self._unified_branches[i].intensity * bs - 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, - ) - - # Only plot active branches (not excluded by neutrino mass or intensity cutoff) - active_branch_indices = [] - for i in range(n_branches): - branch_endpoint_keV = self._unified_branches[i].endpoint_MeV * 1000 - neutrino_mass_keV = self.config.neutrino_mass_MeV * 1000 - is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV - is_excluded_by_intensity = ( - self._unified_branches[i].intensity < self.config.intensity_cutoff - ) - if not is_excluded_by_mass and not is_excluded_by_intensity: - active_branch_indices.append(i) - - for i in active_branch_indices: - contrib = branch_contribs[i] - color = branch_colors[i % len(branch_colors)] - ax.plot( - 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(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, len(active_branch_indices) + 1)) - - # 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", - alpha=0.8, - ) - for i in active_branch_indices: - color = branch_colors[i % len(branch_colors)] - 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) - 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 Br. {i + 1}", - ) - ax.set_ylabel("Normalized Factor", fontsize=10) - 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", - ] - 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=name, - alpha=0.8, - ) - ax.set_ylabel("Normalized Factor", fontsize=10) - 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 in active_branch_indices: - color = branch_colors[i % len(branch_colors)] - 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}", - ) - ax.set_ylabel("Radiative Correction", fontsize=10) - ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=10) - ax.set_title("Radiative Corrections", fontsize=11) - ax.grid(True, alpha=0.3) - ax.legend(fontsize=8, loc="best", ncol=min(2, len(active_branch_indices) + 1)) - - 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}") - - 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}" - - # --------------------------------------------------------------------------- - # Export - # --------------------------------------------------------------------------- - - def export_to_csv( - self, filename: str, source_type: str = "unknown", csv_verbosity: str = "full" - ) -> None: - """Export spectrum data to CSV with metadata header. - - Parameters - ---------- - filename : str - Output CSV file path. - source_type : str - Source type (paceENSDF, json, etc.). - csv_verbosity : str - Verbosity level: 'full' (all components) or 'concise' (total spectrum and normalized branch components only). - """ - total = self.total_spectrum(normalize=True) - components = self.components - - # Only include enabled corrections that are not basic components - # Basic components (phase_space and fermi) are always used and should not be listed - enabled = [] - 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") - - parent_symbol = self._element_symbol(self.config.Z_parent) - daughter_symbol = self._element_symbol(self.config.Z_daughter) - - # Get git commit hash - git_commit = get_git_short_hash() - - # Format source line with paceENSDF URL if applicable - if source_type == "paceENSDF": - source_line = "# source: paceENSDF (https://pypi.org/project/paceENSDF/)" - else: - source_line = f"# source: {source_type}" - - header_lines = [ - f"# beta-spectrum v{__import__('beta_spectrum').__version__}(build.{git_commit})", - f"# timestamp: {datetime.now(timezone.utc).isoformat()}", - source_line, - f"# nuclide: {parent_symbol}{self.config.A_number} -> {daughter_symbol}{self.config.A_number} (Z={self.config.Z_parent}->{self.config.Z_daughter})", - f"# basic_endpoint: {self.config.endpoint_MeV * 1000:.1f} keV", - ] - - # Add neutrino mass info if non-zero - if self.config.neutrino_mass_MeV > 0: - header_lines.append( - f"# neutrino_mass_MeV: {self.config.neutrino_mass_MeV:.6f} (non-zero neutrino mass effects included)" - ) - else: - header_lines.append( - f"# neutrino_mass_MeV: {self.config.neutrino_mass_MeV:.6f}" - ) - - header_lines.append(f"# corrections: {', '.join(enabled)}") - header_lines.append(f"# e_step: {self.config.e_step_MeV:.4f} MeV") - header_lines.append(f"# n_points: {len(self.energies_MeV)}") - - # Add normalization mode - if self.config.normalize_to_one: - header_lines.append("# branch normalization: to_one (ignoring intensity)") - else: - header_lines.append("# branch normalization: intensity") - - # Separate branches into active and excluded lists - active_branches = [] - excluded_branches = [] - - for i, branch in enumerate(self._unified_branches): - branch_endpoint_keV = branch.endpoint_MeV * 1000 - neutrino_mass_keV = self.config.neutrino_mass_MeV * 1000 - - # Check if branch is excluded due to neutrino mass - if branch_endpoint_keV <= neutrino_mass_keV: - excluded_branches.append( - f" branch_{i + 1}: E₀={branch_endpoint_keV:.1f} keV, transition={branch.transition_type}, intensity={branch.intensity:.4f} - excluded: endpoint < m_nu" - ) - # Check if branch is excluded due to intensity cutoff - elif branch.intensity < self.config.intensity_cutoff: - excluded_branches.append( - f" branch_{i + 1}: E₀={branch_endpoint_keV:.1f} keV, transition={branch.transition_type}, intensity={branch.intensity:.4f} - excluded: intensity < cutoff" - ) - else: - active_branches.append( - f" branch_{i + 1}: E₀={branch_endpoint_keV:.1f} keV, transition={branch.transition_type}, intensity={branch.intensity:.4f}" - ) - - # Add active branches list - if active_branches: - header_lines.append("# active branches:") - for branch_info in active_branches: - header_lines.append(f"# {branch_info}") - - # Add excluded branches list if there are any - if excluded_branches: - header_lines.append("# excluded branches:") - for branch_info in excluded_branches: - header_lines.append(f"# {branch_info}") - - # Create data dictionary with concise column headers - data: Dict[str, np.ndarray] = { - "energy_MeV": self.energies_MeV, - } - - # Only add total_spectrum if not normalizing to one - if not self.config.normalize_to_one: - data["total_spectrum"] = total - - if csv_verbosity == "full": - # Per-branch spectra and components with concise branch naming - for i in range(self._n_branches): - # Get branch spectrum - branch_spectrum = self._get_branch_spectrum(i) - - if self.config.normalize_to_one: - # Normalize to 1, ignoring intensity - sum_spectrum = np.sum(branch_spectrum) - if sum_spectrum > 0: - normalized_spectrum = branch_spectrum / sum_spectrum - else: - normalized_spectrum = branch_spectrum - else: - # Use original spectrum with intensity - normalized_spectrum = branch_spectrum * self._unified_branches[i].intensity - - data[f"branch_{i + 1}_spectrum"] = normalized_spectrum - - 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] - - # Universal components (no prefix) - for name, values in components.items(): - if not name.startswith("branch_"): - data[name] = values - elif csv_verbosity == "concise": - # Only total spectrum and normalized branch components with concise naming - for i in range(self._n_branches): - branch_spectrum = self._get_branch_spectrum(i) - # Check if branch is degenerate (neutrino mass >= branch endpoint) - branch_endpoint_keV = self._unified_branches[i].endpoint_MeV * 1000 - neutrino_mass_keV = self.config.neutrino_mass_MeV * 1000 - if branch_endpoint_keV <= neutrino_mass_keV: - # Branch is degenerate, spectrum should be all zeros - data[f"branch_{i + 1}"] = np.zeros_like(branch_spectrum) - else: - # Branch is valid, normalize as usual - if self.config.normalize_to_one: - # Normalize to 1, ignoring intensity - sum_spectrum = np.sum(branch_spectrum) - if sum_spectrum > 0: - data[f"branch_{i + 1}"] = branch_spectrum / sum_spectrum - else: - data[f"branch_{i + 1}"] = branch_spectrum - else: - # Normalize as usual with intensity - sum_spectrum = np.sum(branch_spectrum) - if sum_spectrum > 0: - data[f"branch_{i + 1}"] = branch_spectrum / sum_spectrum * self._unified_branches[i].intensity - else: - data[f"branch_{i + 1}"] = branch_spectrum - - df = pd.DataFrame(data) - # Determine separator - separator = '\t' if self.config.csv_separator == 'tab' else self.config.csv_separator - df.to_csv(filename, index=False, float_format="%.4e", sep=separator) - - 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: - self._logger.info("CSV exported to %s", filename) - - 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, - } - - -@dataclass -class _UnifiedBranch: - """Unified branch view for both single- and multi-branch modes.""" - - endpoint_MeV: float - transition_type: str - intensity: float diff --git a/tests/physics/test_multi_branch.py b/tests/physics/test_multi_branch.py index 1036158..d3981f3 100644 --- a/tests/physics/test_multi_branch.py +++ b/tests/physics/test_multi_branch.py @@ -33,7 +33,7 @@ from beta_spectrum import ( BranchConfig, SpectrumConfig, ) -from beta_spectrum.visualize import BetaSpectrumAnalyzer +from beta_spectrum.output import BetaSpectrumAnalyzer from beta_spectrum.utils import T_to_W # --------------------------------------------------------------------------- @@ -461,7 +461,7 @@ class TestMultiBranchCSVExport: import pandas as pd # Read CSV (skip comment headers, use comma delimiter by default) - df = pd.read_csv(csv_path, comment="#", sep=',') + df = pd.read_csv(csv_path, comment="#", sep=",") # Should have total spectrum column assert "total_spectrum" in df.columns @@ -500,7 +500,7 @@ class TestMultiBranchCSVExport: import pandas as pd - df = pd.read_csv(csv_path, comment="#", sep=',') + df = pd.read_csv(csv_path, comment="#", sep=",") assert "total_spectrum" in df.columns assert "branch_1_spectrum" in df.columns diff --git a/tests/quality/test_logging_utils.py b/tests/quality/test_logging_utils.py index e84038e..f17ea84 100644 --- a/tests/quality/test_logging_utils.py +++ b/tests/quality/test_logging_utils.py @@ -131,7 +131,7 @@ class TestCSVMetadataHeader: def test_header_contains_version(self): from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig - from beta_spectrum.visualize import BetaSpectrumAnalyzer + from beta_spectrum.output import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, @@ -159,7 +159,7 @@ class TestCSVMetadataHeader: def test_header_contains_timestamp(self): from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig - from beta_spectrum.visualize import BetaSpectrumAnalyzer + from beta_spectrum.output import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, @@ -186,7 +186,7 @@ class TestCSVMetadataHeader: def test_header_contains_nuclide_info(self): from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig - from beta_spectrum.visualize import BetaSpectrumAnalyzer + from beta_spectrum.output import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, @@ -214,7 +214,7 @@ class TestCSVMetadataHeader: def test_header_contains_git_commit(self): from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig - from beta_spectrum.visualize import BetaSpectrumAnalyzer + from beta_spectrum.output import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43, @@ -243,7 +243,7 @@ class TestCSVMetadataHeader: def test_header_contains_corrections(self): from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig - from beta_spectrum.visualize import BetaSpectrumAnalyzer + from beta_spectrum.output import BetaSpectrumAnalyzer config = SpectrumConfig( Z_parent=43,