465 lines
14 KiB
Python
465 lines
14 KiB
Python
"""
|
||
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
|
||
|
||
import re
|
||
import shutil
|
||
import tempfile
|
||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
|
||
import matplotlib.pyplot as plt
|
||
from matplotlib.figure import Figure
|
||
|
||
from .csv_export import CSVEmitter
|
||
|
||
if TYPE_CHECKING:
|
||
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*(.*)$")
|
||
|
||
_BRANCH_COLORS = [
|
||
"#1f77b4",
|
||
"#ff7f0e",
|
||
"#2ca02c",
|
||
"#d62728",
|
||
"#9467bd",
|
||
"#8c564b",
|
||
"#e377c2",
|
||
"#7f7f7f",
|
||
"#bcbd22",
|
||
"#17becf",
|
||
]
|
||
|
||
|
||
def parse_spectrum_csv(
|
||
csv_path: str,
|
||
) -> Tuple[pd.DataFrame, Dict[str, Any]]:
|
||
"""Read an exported spectrum CSV and parse its metadata header.
|
||
|
||
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:
|
||
title += f" ({n_branches} branches)"
|
||
return title
|
||
|
||
|
||
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']} | E₀={endpoint:.1f} keV | Δ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.96,
|
||
f"| m_ν={meta['neutrino_mass_MeV']:.6f} MeV",
|
||
transform=fig.transFigure,
|
||
fontsize=8,
|
||
verticalalignment="top",
|
||
horizontalalignment="left",
|
||
fontfamily="monospace",
|
||
color="red",
|
||
)
|
||
|
||
if meta["build"]:
|
||
fig.text(
|
||
0.01,
|
||
0.01,
|
||
f"{meta['build']} | {meta['timestamp']}",
|
||
transform=fig.transFigure,
|
||
fontsize=8,
|
||
verticalalignment="bottom",
|
||
horizontalalignment="left",
|
||
fontfamily="monospace",
|
||
color="gray",
|
||
)
|
||
|
||
|
||
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
|
||
|
||
fig, ax = plt.subplots(figsize=(12, 7))
|
||
|
||
if "total_spectrum" in df.columns:
|
||
ax.plot(
|
||
energies,
|
||
df["total_spectrum"].to_numpy(),
|
||
"k-",
|
||
lw=2.5,
|
||
label="Total spectrum",
|
||
zorder=10,
|
||
)
|
||
|
||
for color_i, b in enumerate(_active_branches(meta)):
|
||
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['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)
|
||
|
||
# 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")
|
||
|
||
# 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(
|
||
energies,
|
||
values_scaled,
|
||
color=correction_colors[corr_i],
|
||
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)
|
||
|
||
_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
|
||
|
||
|
||
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)
|