refactor: rename visualize module to output module and separate concerns into core, csv_export, and visualization modules

This commit is contained in:
2026-07-07 12:59:11 +03:00
parent 6dc1e3924d
commit 9bf1e1333e
11 changed files with 989 additions and 884 deletions
+1 -1
View File
@@ -11,7 +11,7 @@ from .components.exchange import ExchangeCorrection
from .components.radiative import RadiativeCorrection from .components.radiative import RadiativeCorrection
from .spectrum import BetaSpectrum, SpectrumConfig, BranchConfig from .spectrum import BetaSpectrum, SpectrumConfig, BranchConfig
from .visualize import BetaSpectrumAnalyzer from .output import BetaSpectrumAnalyzer
from .nuclear_data import ( from .nuclear_data import (
DecayInfo, DecayInfo,
BranchInfo, BranchInfo,
+2 -2
View File
@@ -155,7 +155,7 @@ def _build_parser() -> argparse.ArgumentParser:
log_group.add_argument( log_group.add_argument(
"--csv-separator", "--csv-separator",
type=str, type=str,
default=',', default=",",
help="CSV column separator (default: ','). Use 'tab' for tab separation.", help="CSV column separator (default: ','). Use 'tab' for tab separation.",
) )
log_group.add_argument( log_group.add_argument(
@@ -183,7 +183,7 @@ def _run(args: argparse.Namespace) -> None:
BranchConfig, BranchConfig,
SpectrumConfig, SpectrumConfig,
) )
from beta_spectrum.visualize import BetaSpectrumAnalyzer from beta_spectrum.output import BetaSpectrumAnalyzer
# Determine logging level # Determine logging level
if args.quiet: if args.quiet:
+2 -2
View File
@@ -471,7 +471,7 @@ def decay_info_to_config(
e_step_MeV: float = 0.001, e_step_MeV: float = 0.001,
intensity_cutoff: float = 0.0, intensity_cutoff: float = 0.0,
neutrino_mass_MeV: float = 0.0, neutrino_mass_MeV: float = 0.0,
csv_separator: str = ',', csv_separator: str = ",",
normalize_to_one: bool = False, normalize_to_one: bool = False,
) -> SpectrumConfig: ) -> SpectrumConfig:
""" """
@@ -698,7 +698,7 @@ def create_config_from_source(
e_step = kwargs.get("e_step_MeV", 0.001) e_step = kwargs.get("e_step_MeV", 0.001)
intensity_cutoff = kwargs.get("intensity_cutoff", 0.0) intensity_cutoff = kwargs.get("intensity_cutoff", 0.0)
neutrino_mass = kwargs.get("neutrino_mass_MeV", 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) normalize_to_one = kwargs.get("normalize_to_one", False)
info = get_decay_info_from_paceENSDF(nuclide, decay_type, decay_index) info = get_decay_info_from_paceENSDF(nuclide, decay_type, decay_index)
return decay_info_to_config( return decay_info_to_config(
+23
View File
@@ -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",
]
+285
View File
@@ -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
+219
View File
@@ -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)
+444
View File
@@ -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}")
+4 -2
View File
@@ -89,8 +89,10 @@ class SpectrumConfig:
) )
# CSV output options # CSV output options
csv_separator: str = ',' # Default separator for CSV output csv_separator: str = "," # Default separator for CSV output
normalize_to_one: bool = False # If True, normalize each branch to 1, ignoring intensity normalize_to_one: bool = (
False # If True, normalize each branch to 1, ignoring intensity
)
class BetaSpectrum: class BetaSpectrum:
-868
View File
@@ -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
+3 -3
View File
@@ -33,7 +33,7 @@ from beta_spectrum import (
BranchConfig, BranchConfig,
SpectrumConfig, SpectrumConfig,
) )
from beta_spectrum.visualize import BetaSpectrumAnalyzer from beta_spectrum.output import BetaSpectrumAnalyzer
from beta_spectrum.utils import T_to_W from beta_spectrum.utils import T_to_W
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -461,7 +461,7 @@ class TestMultiBranchCSVExport:
import pandas as pd import pandas as pd
# Read CSV (skip comment headers, use comma delimiter by default) # 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 # Should have total spectrum column
assert "total_spectrum" in df.columns assert "total_spectrum" in df.columns
@@ -500,7 +500,7 @@ class TestMultiBranchCSVExport:
import pandas as pd 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 "total_spectrum" in df.columns
assert "branch_1_spectrum" in df.columns assert "branch_1_spectrum" in df.columns
+5 -5
View File
@@ -131,7 +131,7 @@ class TestCSVMetadataHeader:
def test_header_contains_version(self): def test_header_contains_version(self):
from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig
from beta_spectrum.visualize import BetaSpectrumAnalyzer from beta_spectrum.output import BetaSpectrumAnalyzer
config = SpectrumConfig( config = SpectrumConfig(
Z_parent=43, Z_parent=43,
@@ -159,7 +159,7 @@ class TestCSVMetadataHeader:
def test_header_contains_timestamp(self): def test_header_contains_timestamp(self):
from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig
from beta_spectrum.visualize import BetaSpectrumAnalyzer from beta_spectrum.output import BetaSpectrumAnalyzer
config = SpectrumConfig( config = SpectrumConfig(
Z_parent=43, Z_parent=43,
@@ -186,7 +186,7 @@ class TestCSVMetadataHeader:
def test_header_contains_nuclide_info(self): def test_header_contains_nuclide_info(self):
from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig
from beta_spectrum.visualize import BetaSpectrumAnalyzer from beta_spectrum.output import BetaSpectrumAnalyzer
config = SpectrumConfig( config = SpectrumConfig(
Z_parent=43, Z_parent=43,
@@ -214,7 +214,7 @@ class TestCSVMetadataHeader:
def test_header_contains_git_commit(self): def test_header_contains_git_commit(self):
from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig
from beta_spectrum.visualize import BetaSpectrumAnalyzer from beta_spectrum.output import BetaSpectrumAnalyzer
config = SpectrumConfig( config = SpectrumConfig(
Z_parent=43, Z_parent=43,
@@ -243,7 +243,7 @@ class TestCSVMetadataHeader:
def test_header_contains_corrections(self): def test_header_contains_corrections(self):
from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig from beta_spectrum.spectrum import BetaSpectrum, SpectrumConfig
from beta_spectrum.visualize import BetaSpectrumAnalyzer from beta_spectrum.output import BetaSpectrumAnalyzer
config = SpectrumConfig( config = SpectrumConfig(
Z_parent=43, Z_parent=43,