refactor: rename visualize module to output module and separate concerns into core, csv_export, and visualization modules
This commit is contained in:
@@ -11,7 +11,7 @@ from .components.exchange import ExchangeCorrection
|
||||
from .components.radiative import RadiativeCorrection
|
||||
|
||||
from .spectrum import BetaSpectrum, SpectrumConfig, BranchConfig
|
||||
from .visualize import BetaSpectrumAnalyzer
|
||||
from .output import BetaSpectrumAnalyzer
|
||||
from .nuclear_data import (
|
||||
DecayInfo,
|
||||
BranchInfo,
|
||||
|
||||
@@ -155,7 +155,7 @@ def _build_parser() -> argparse.ArgumentParser:
|
||||
log_group.add_argument(
|
||||
"--csv-separator",
|
||||
type=str,
|
||||
default=',',
|
||||
default=",",
|
||||
help="CSV column separator (default: ','). Use 'tab' for tab separation.",
|
||||
)
|
||||
log_group.add_argument(
|
||||
@@ -183,7 +183,7 @@ def _run(args: argparse.Namespace) -> None:
|
||||
BranchConfig,
|
||||
SpectrumConfig,
|
||||
)
|
||||
from beta_spectrum.visualize import BetaSpectrumAnalyzer
|
||||
from beta_spectrum.output import BetaSpectrumAnalyzer
|
||||
|
||||
# Determine logging level
|
||||
if args.quiet:
|
||||
|
||||
@@ -471,7 +471,7 @@ def decay_info_to_config(
|
||||
e_step_MeV: float = 0.001,
|
||||
intensity_cutoff: float = 0.0,
|
||||
neutrino_mass_MeV: float = 0.0,
|
||||
csv_separator: str = ',',
|
||||
csv_separator: str = ",",
|
||||
normalize_to_one: bool = False,
|
||||
) -> SpectrumConfig:
|
||||
"""
|
||||
@@ -698,7 +698,7 @@ def create_config_from_source(
|
||||
e_step = kwargs.get("e_step_MeV", 0.001)
|
||||
intensity_cutoff = kwargs.get("intensity_cutoff", 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)
|
||||
info = get_decay_info_from_paceENSDF(nuclide, decay_type, decay_index)
|
||||
return decay_info_to_config(
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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}")
|
||||
@@ -87,10 +87,12 @@ class SpectrumConfig:
|
||||
intensity_cutoff: float = (
|
||||
0.0 # Ignore branches below this fraction (default: 0.0 = all)
|
||||
)
|
||||
|
||||
|
||||
# CSV output options
|
||||
csv_separator: str = ',' # Default separator for CSV output
|
||||
normalize_to_one: bool = False # If True, normalize each branch to 1, ignoring intensity
|
||||
csv_separator: str = "," # Default separator for CSV output
|
||||
normalize_to_one: bool = (
|
||||
False # If True, normalize each branch to 1, ignoring intensity
|
||||
)
|
||||
|
||||
|
||||
class BetaSpectrum:
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user