Files
beta_shape_pnpi/beta_spectrum/output/csv_export.py
T

220 lines
9.2 KiB
Python

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