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
+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}")