Files
beta_shape_pnpi/beta_spectrum/output/core.py
T

286 lines
8.0 KiB
Python

"""
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