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