220 lines
9.2 KiB
Python
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)
|