445 lines
16 KiB
Python
445 lines
16 KiB
Python
"""
|
||
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}")
|