feat(output): physical branch spectra, universal factor, CSV-driven plotting
This commit is contained in:
@@ -1,99 +1,228 @@
|
||||
"""
|
||||
Visualization module — Plotting functionality for BetaSpectrumAnalyzer.
|
||||
Visualization module — plotting built ON TOP OF the CSV export.
|
||||
|
||||
All graphical representations are produced from the exported CSV file
|
||||
(data columns + metadata header) and never recomputed from the analyzer in
|
||||
parallel. The CSV file is the single source of truth: a plot is produced
|
||||
exactly from the bytes a researcher would download, which turns every plot
|
||||
into an additional quality check of the export routine.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Dict, Optional
|
||||
import re
|
||||
import shutil
|
||||
import tempfile
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import pandas as pd
|
||||
|
||||
from beta_spectrum.logging_utils import get_git_short_hash
|
||||
from beta_spectrum.utils import T_to_W
|
||||
import matplotlib.pyplot as plt
|
||||
from matplotlib.figure import Figure
|
||||
|
||||
from .csv_export import CSVEmitter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from matplotlib.figure import Figure
|
||||
|
||||
from beta_spectrum.output.core import BetaSpectrumAnalyzer
|
||||
|
||||
_BRANCH_LINE_RE = re.compile(
|
||||
r"^\s*branch_(\d+):\s*E₀=([0-9.]+)\s*keV,\s*transition=(\w+),\s*"
|
||||
r"intensity=([0-9.]+)(.*)$"
|
||||
)
|
||||
_EXCLUDED_SUFFIX_RE = re.compile(r"-\s*excluded:\s*(.*)$")
|
||||
|
||||
class Visualizer:
|
||||
"""Visualization functionality for BetaSpectrumAnalyzer."""
|
||||
_BRANCH_COLORS = [
|
||||
"#1f77b4",
|
||||
"#ff7f0e",
|
||||
"#2ca02c",
|
||||
"#d62728",
|
||||
"#9467bd",
|
||||
"#8c564b",
|
||||
"#e377c2",
|
||||
"#7f7f7f",
|
||||
"#bcbd22",
|
||||
"#17becf",
|
||||
]
|
||||
|
||||
@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"
|
||||
)
|
||||
def parse_spectrum_csv(
|
||||
csv_path: str,
|
||||
) -> Tuple[pd.DataFrame, Dict[str, Any]]:
|
||||
"""Read an exported spectrum CSV and parse its metadata header.
|
||||
|
||||
if show_components:
|
||||
Visualizer._plot_debug_view(analyzer, total, components, save_path)
|
||||
Parameters
|
||||
----------
|
||||
csv_path : str
|
||||
Path to a CSV file produced by CSVEmitter.export_to_csv.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tuple[pandas.DataFrame, dict]
|
||||
The data frame (comment lines skipped, separator auto-detected)
|
||||
and a metadata dict with keys: version, build, timestamp, source,
|
||||
nuclide, endpoint_keV, neutrino_mass_MeV, corrections, e_step_MeV,
|
||||
n_points, normalization, branches.
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
If the file cannot be read or parsed.
|
||||
"""
|
||||
try:
|
||||
df = pd.read_csv(csv_path, sep=None, engine="python", comment="#")
|
||||
raw_lines = []
|
||||
with open(csv_path, "r") as f:
|
||||
for line in f:
|
||||
if line.startswith("#"):
|
||||
raw_lines.append(line.rstrip("\n"))
|
||||
else:
|
||||
break # header block ended
|
||||
except OSError as exc:
|
||||
raise ValueError(f"Cannot read spectrum CSV {csv_path}: {exc}") from exc
|
||||
|
||||
meta: Dict[str, Any] = {
|
||||
"version": None,
|
||||
"build": None,
|
||||
"timestamp": None,
|
||||
"source": None,
|
||||
"nuclide": None,
|
||||
"endpoint_keV": None,
|
||||
"neutrino_mass_MeV": None,
|
||||
"corrections": [],
|
||||
"e_step_MeV": None,
|
||||
"n_points": None,
|
||||
"normalization": None,
|
||||
"branches": [],
|
||||
}
|
||||
|
||||
for line in raw_lines:
|
||||
body = line.lstrip("#").strip()
|
||||
m = re.match(r"beta-spectrum v([\d.]+)\(build\.(\w+)\)", body)
|
||||
if m:
|
||||
meta["version"] = m.group(1)
|
||||
meta["build"] = m.group(2)
|
||||
continue
|
||||
m = re.match(r"timestamp:\s*(.+)", body)
|
||||
if m:
|
||||
meta["timestamp"] = m.group(1).strip()
|
||||
continue
|
||||
m = re.match(r"source:\s*(\S+)", body)
|
||||
if m:
|
||||
meta["source"] = m.group(1)
|
||||
continue
|
||||
m = re.match(r"nuclide:\s*(.+?)\s+\(Z=", body)
|
||||
if m:
|
||||
meta["nuclide"] = m.group(1).strip()
|
||||
continue
|
||||
m = re.match(r"basic_endpoint:\s*([0-9.]+)\s*keV", body)
|
||||
if m:
|
||||
meta["endpoint_keV"] = float(m.group(1))
|
||||
continue
|
||||
m = re.match(r"neutrino_mass_MeV:\s*([0-9.]+)", body)
|
||||
if m:
|
||||
meta["neutrino_mass_MeV"] = float(m.group(1))
|
||||
continue
|
||||
m = re.match(r"corrections:\s*(.*)", body)
|
||||
if m:
|
||||
meta["corrections"] = [
|
||||
c.strip() for c in m.group(1).split(",") if c.strip()
|
||||
]
|
||||
continue
|
||||
m = re.match(r"e_step:\s*([0-9.]+)\s*MeV", body)
|
||||
if m:
|
||||
meta["e_step_MeV"] = float(m.group(1))
|
||||
continue
|
||||
m = re.match(r"n_points:\s*(\d+)", body)
|
||||
if m:
|
||||
meta["n_points"] = int(m.group(1))
|
||||
continue
|
||||
m = re.match(r"branch normalization:\s*(\S+)", body)
|
||||
if m:
|
||||
meta["normalization"] = m.group(1)
|
||||
continue
|
||||
m = _BRANCH_LINE_RE.match(body)
|
||||
if m:
|
||||
excluded = _EXCLUDED_SUFFIX_RE.search(m.group(5))
|
||||
meta["branches"].append(
|
||||
{
|
||||
"index": int(m.group(1)),
|
||||
"endpoint_keV": float(m.group(2)),
|
||||
"transition": m.group(3),
|
||||
"intensity": float(m.group(4)),
|
||||
"excluded": None if excluded is None else excluded.group(1).strip(),
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
if meta["n_points"] is None:
|
||||
meta["n_points"] = len(df)
|
||||
|
||||
return df, meta
|
||||
|
||||
|
||||
def _branch_column(df: pd.DataFrame, index: int) -> Optional[str]:
|
||||
"""Find the branch spectrum column name for a 1-based branch index."""
|
||||
for candidate in (f"branch_{index}_spectrum", f"branch_{index}"):
|
||||
if candidate in df.columns:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def _active_branches(meta: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
return [b for b in meta["branches"] if b["excluded"] is None]
|
||||
|
||||
|
||||
def _spectrum_title(meta: Dict[str, Any]) -> str:
|
||||
n_branches = len(meta["branches"])
|
||||
active = len(_active_branches(meta))
|
||||
title = f"Beta-decay: {meta['nuclide']}"
|
||||
if n_branches > 1:
|
||||
if active < n_branches:
|
||||
title += f" ({active} of {n_branches} branches)"
|
||||
else:
|
||||
Visualizer._plot_spectra_view(analyzer, total, components, save_path)
|
||||
title += f" ({n_branches} branches)"
|
||||
return title
|
||||
|
||||
@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"
|
||||
)
|
||||
|
||||
def _apply_figure_header(fig: Figure, meta: Dict[str, Any]) -> None:
|
||||
"""Write the figure header/footer lines from the CSV metadata."""
|
||||
e_step_keV = (meta["e_step_MeV"] or 0.0) * 1000.0
|
||||
endpoint = meta["endpoint_keV"]
|
||||
header = (
|
||||
f"{meta['nuclide']} "
|
||||
f"| E₀={endpoint:.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",
|
||||
)
|
||||
|
||||
if meta["neutrino_mass_MeV"]:
|
||||
fig.text(
|
||||
0.01,
|
||||
0.98,
|
||||
header,
|
||||
0.96,
|
||||
f"| m_ν={meta['neutrino_mass_MeV']:.6f} MeV",
|
||||
transform=fig.transFigure,
|
||||
fontsize=9,
|
||||
fontsize=8,
|
||||
verticalalignment="top",
|
||||
horizontalalignment="left",
|
||||
fontfamily="monospace",
|
||||
color="red",
|
||||
)
|
||||
|
||||
# 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}"
|
||||
if meta["build"]:
|
||||
fig.text(
|
||||
0.01,
|
||||
0.01,
|
||||
id_text,
|
||||
f"{meta['build']} | {meta['timestamp']}",
|
||||
transform=fig.transFigure,
|
||||
fontsize=8,
|
||||
verticalalignment="bottom",
|
||||
@@ -102,343 +231,238 @@ class Visualizer:
|
||||
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
|
||||
def _plot_spectra_view_from_df(
|
||||
df: pd.DataFrame,
|
||||
meta: Dict[str, Any],
|
||||
save_path: Optional[str],
|
||||
) -> Figure:
|
||||
"""Plot total spectrum and per-branch spectra from CSV columns."""
|
||||
energies = df["energy_MeV"].to_numpy()
|
||||
e_step_keV = (meta["e_step_MeV"] or 0.0) * 1000.0
|
||||
|
||||
@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)
|
||||
fig, ax = plt.subplots(figsize=(12, 7))
|
||||
|
||||
if "total_spectrum" in df.columns:
|
||||
ax.plot(
|
||||
analyzer.energies_MeV,
|
||||
total_raw,
|
||||
energies,
|
||||
df["total_spectrum"].to_numpy(),
|
||||
"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]
|
||||
for color_i, b in enumerate(_active_branches(meta)):
|
||||
col = _branch_column(df, int(b["index"]))
|
||||
if col is None:
|
||||
continue
|
||||
ax.plot(
|
||||
analyzer.energies_MeV,
|
||||
total_raw,
|
||||
energies,
|
||||
df[col].to_numpy(),
|
||||
color=_BRANCH_COLORS[color_i % len(_BRANCH_COLORS)],
|
||||
lw=1.5,
|
||||
alpha=0.7,
|
||||
label=(
|
||||
f"Br. {b['index']}: {b['endpoint_keV']:.1f} keV "
|
||||
f"({b['intensity'] * 100:.1f}%)"
|
||||
),
|
||||
)
|
||||
|
||||
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(_spectrum_title(meta), fontsize=14, fontweight="bold")
|
||||
ax.set_yscale("log")
|
||||
ax.grid(True, alpha=0.3)
|
||||
ax.legend(loc="best", fontsize=9, ncol=2)
|
||||
|
||||
_apply_figure_header(fig, meta)
|
||||
fig.tight_layout()
|
||||
|
||||
if save_path:
|
||||
fig.savefig(save_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Figure saved to {save_path}")
|
||||
return fig
|
||||
|
||||
|
||||
def _plot_debug_view_from_df(
|
||||
df: pd.DataFrame,
|
||||
meta: Dict[str, Any],
|
||||
save_path: Optional[str],
|
||||
) -> Figure:
|
||||
"""Debug view (4 panels) built entirely from CSV columns."""
|
||||
energies = df["energy_MeV"].to_numpy()
|
||||
e_step_keV = (meta["e_step_MeV"] or 0.0) * 1000.0
|
||||
active = _active_branches(meta)
|
||||
|
||||
fig, axes = plt.subplots(4, 1, figsize=(12, 16), sharex=True)
|
||||
fig.suptitle(_spectrum_title(meta), fontsize=14, fontweight="bold")
|
||||
|
||||
# Panel 1: total spectrum + per-branch spectra
|
||||
ax = axes[0]
|
||||
if "total_spectrum" in df.columns:
|
||||
ax.plot(
|
||||
energies,
|
||||
df["total_spectrum"].to_numpy(),
|
||||
"k-",
|
||||
lw=2,
|
||||
label="Total spectrum",
|
||||
zorder=10,
|
||||
)
|
||||
for color_i, b in enumerate(active):
|
||||
col = _branch_column(df, int(b["index"]))
|
||||
if col is None:
|
||||
continue
|
||||
ax.plot(
|
||||
energies,
|
||||
df[col].to_numpy(),
|
||||
color=_BRANCH_COLORS[color_i % len(_BRANCH_COLORS)],
|
||||
lw=1.5,
|
||||
alpha=0.7,
|
||||
label=f"Br. {b['index']} ({b['intensity'] * 100:.1f}%)",
|
||||
)
|
||||
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=2)
|
||||
|
||||
# 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)
|
||||
# Panel 2: Fermi function + all phase spaces
|
||||
ax = axes[1]
|
||||
if "Fermi" in df.columns:
|
||||
fermi = df["Fermi"].to_numpy()
|
||||
fermi_scaled = fermi / np.max(fermi)
|
||||
ax.plot(energies, fermi_scaled, "r-", lw=1.5, label="Fermi", alpha=0.8)
|
||||
for color_i, b in enumerate(active):
|
||||
col = f"branch_{b['index']}_PhaseSpace"
|
||||
if col not in df.columns:
|
||||
continue
|
||||
ps_vals = df[col].to_numpy()
|
||||
ps_max = np.max(ps_vals) if np.any(ps_vals > 0) else 1.0
|
||||
ax.plot(
|
||||
energies,
|
||||
ps_vals / ps_max if ps_max > 0 else ps_vals,
|
||||
color=_BRANCH_COLORS[color_i % len(_BRANCH_COLORS)],
|
||||
lw=1.5,
|
||||
alpha=0.7,
|
||||
label=f"PS Br. {b['index']}",
|
||||
)
|
||||
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")
|
||||
|
||||
for i in active_branch_indices:
|
||||
contrib = branch_contribs[i]
|
||||
color = branch_colors[i % len(branch_colors)]
|
||||
# Panel 3: universal higher-order corrections
|
||||
ax = axes[2]
|
||||
correction_colors = ["green", "orange", "purple", "brown"]
|
||||
for corr_i, name in enumerate(
|
||||
["Screening", "FiniteSizeL0", "ChargeDistributionU", "Exchange"]
|
||||
):
|
||||
if name in df.columns:
|
||||
values = df[name].to_numpy()
|
||||
values_scaled = values / np.max(values)
|
||||
ax.plot(
|
||||
analyzer.energies_MeV,
|
||||
contrib,
|
||||
color=color,
|
||||
energies,
|
||||
values_scaled,
|
||||
color=correction_colors[corr_i],
|
||||
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",
|
||||
label=name,
|
||||
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")
|
||||
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 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 per branch)
|
||||
ax = axes[3]
|
||||
for color_i, b in enumerate(active):
|
||||
col = f"branch_{b['index']}_Radiative"
|
||||
if col not in df.columns:
|
||||
continue
|
||||
ax.plot(
|
||||
energies,
|
||||
df[col].to_numpy(),
|
||||
color=_BRANCH_COLORS[color_i % len(_BRANCH_COLORS)],
|
||||
lw=1.5,
|
||||
alpha=0.7,
|
||||
label=f"Br. {b['index']}",
|
||||
)
|
||||
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=2)
|
||||
|
||||
# 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))
|
||||
_apply_figure_header(fig, meta)
|
||||
fig.tight_layout()
|
||||
|
||||
Visualizer._build_figure_header(analyzer, fig)
|
||||
if save_path:
|
||||
fig.savefig(save_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Figure saved to {save_path}")
|
||||
return fig
|
||||
|
||||
plt.tight_layout()
|
||||
|
||||
if save_path:
|
||||
fig.savefig(save_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Figure saved to {save_path}")
|
||||
class Visualizer:
|
||||
"""Visualization built from the exported CSV data (single source of truth)."""
|
||||
|
||||
@staticmethod
|
||||
def plot_from_csv(
|
||||
csv_path: str,
|
||||
save_path: Optional[str] = None,
|
||||
show_components: bool = True,
|
||||
) -> Figure:
|
||||
"""Create a visualization from an exported CSV spectrum file.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
csv_path : str
|
||||
Path to a CSV exported by CSVEmitter.export_to_csv. The plots
|
||||
are built exclusively from the file contents (data columns and
|
||||
metadata header) — nothing is recomputed from the model.
|
||||
save_path : str, optional
|
||||
Path to save the figure.
|
||||
show_components : bool
|
||||
If True, produce the full 4-panel debug view; requires the
|
||||
'full' CSV verbosity. If the file lacks component columns, a
|
||||
warning is printed and the simple spectrum view is produced.
|
||||
"""
|
||||
df, meta = parse_spectrum_csv(csv_path)
|
||||
if show_components and "Fermi" in df.columns:
|
||||
return _plot_debug_view_from_df(df, meta, save_path)
|
||||
if show_components:
|
||||
print(
|
||||
"Warning: requested component panels but CSV lacks component "
|
||||
"columns (export with --csv-verbosity full); falling back to "
|
||||
"simple spectrum view."
|
||||
)
|
||||
return _plot_spectra_view_from_df(df, meta, save_path)
|
||||
|
||||
@staticmethod
|
||||
def plot_analysis(
|
||||
analyzer: BetaSpectrumAnalyzer,
|
||||
save_path: Optional[str] = None,
|
||||
show_components: bool = True,
|
||||
) -> Figure:
|
||||
"""Create a visualization by exporting CSV and plotting from it.
|
||||
|
||||
The export goes to a temporary file using the 'full' verbosity so
|
||||
that the plots exercise the same data pipeline (header + columns)
|
||||
a researcher receives.
|
||||
"""
|
||||
tmp_dir = tempfile.mkdtemp(prefix="bs_plot_")
|
||||
try:
|
||||
csv_path = f"{tmp_dir}/spectrum.csv"
|
||||
CSVEmitter.export_to_csv(
|
||||
analyzer,
|
||||
csv_path,
|
||||
source_type="plot",
|
||||
csv_verbosity="full",
|
||||
)
|
||||
return Visualizer.plot_from_csv(
|
||||
csv_path,
|
||||
save_path=save_path,
|
||||
show_components=show_components,
|
||||
)
|
||||
finally:
|
||||
shutil.rmtree(tmp_dir, ignore_errors=True)
|
||||
|
||||
Reference in New Issue
Block a user