refactor: rename visualize module to output module and separate concerns into core, csv_export, and visualization modules
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user