From c2795bf86e45d528cfff92708c1d259b580914a6 Mon Sep 17 00:00:00 2001 From: Evgeniy Unzhakov Date: Fri, 3 Jul 2026 15:10:10 +0300 Subject: [PATCH] fix: use underscores in branch names for CSV headers --- beta_spectrum/visualize/__init__.py | 10 +++++----- tests/physics/test_multi_branch.py | 23 ++++++++++++----------- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/beta_spectrum/visualize/__init__.py b/beta_spectrum/visualize/__init__.py index 7bbe793..30763e2 100644 --- a/beta_spectrum/visualize/__init__.py +++ b/beta_spectrum/visualize/__init__.py @@ -723,11 +723,11 @@ class BetaSpectrumAnalyzer: if csv_verbosity == "full": # Per-branch spectra and components with concise branch naming for i in range(self._n_branches): - data[f"branch {i + 1}"] = self._get_branch_spectrum(i) + data[f"branch_{i + 1}_spectrum"] = self._get_branch_spectrum(i) for comp_name in ["PhaseSpace", "Radiative"]: key = self._get_component_key(i, comp_name) if key in components: - data[f"branch {i + 1}_{comp_name}"] = components[key] + data[f"branch_{i + 1}_{comp_name}"] = components[key] # Universal components (no prefix) for name, values in components.items(): @@ -742,14 +742,14 @@ class BetaSpectrumAnalyzer: neutrino_mass_keV = self.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) + data[f"branch_{i + 1}_spectrum_normalized"] = np.zeros_like(branch_spectrum) else: # Branch is valid, normalize as usual sum_spectrum = np.sum(branch_spectrum) if sum_spectrum > 0: - data[f"branch {i + 1}"] = branch_spectrum / sum_spectrum + data[f"branch_{i + 1}_spectrum_normalized"] = branch_spectrum / sum_spectrum else: - data[f"branch {i + 1}"] = branch_spectrum + data[f"branch_{i + 1}_spectrum_normalized"] = branch_spectrum df = pd.DataFrame(data) df.to_csv(filename, index=False, float_format="%.4e") diff --git a/tests/physics/test_multi_branch.py b/tests/physics/test_multi_branch.py index 9618648..4711fbd 100644 --- a/tests/physics/test_multi_branch.py +++ b/tests/physics/test_multi_branch.py @@ -467,8 +467,8 @@ class TestMultiBranchCSVExport: assert "total_spectrum" in df.columns # Should have per-branch spectrum columns (intensity/endpoint/transition are in the header) - assert "branch 1" in df.columns - assert "branch 2" in df.columns + assert "branch_1_spectrum" in df.columns + assert "branch_2_spectrum" in df.columns # Should NOT have repeated per-branch metadata columns assert "branch_1_intensity" not in df.columns @@ -503,6 +503,7 @@ class TestMultiBranchCSVExport: df = pd.read_csv(csv_path, comment="#") assert "total_spectrum" in df.columns + assert "branch_1_spectrum" in df.columns assert "Fermi" in df.columns @@ -659,9 +660,9 @@ class TestMultiBranchEdgeCases: W, _ = spectrum.get_energy_grid(config) values = spectrum(W) - assert np.all( - np.isfinite(values) - ), "High-Z multi-branch must not produce NaN/inf" + assert np.all(np.isfinite(values)), ( + "High-Z multi-branch must not produce NaN/inf" + ) def test_no_branches_defaults_to_single(self): """No branches list should default to single-branch mode.""" @@ -726,9 +727,9 @@ class TestMultiBranchEdgeCases: contrib = universal_product * branch.intensity * bs[i] # Branch contribution should not exceed total - assert np.all( - contrib <= total + 1e-10 - ), f"Branch {i} contribution exceeds total at some points" + assert np.all(contrib <= total + 1e-10), ( + f"Branch {i} contribution exceeds total at some points" + ) # Branch contribution should be non-negative assert np.all(contrib >= 0), f"Branch {i} contribution has negative values" @@ -778,6 +779,6 @@ class TestMultiBranchEdgeCases: assert integral > 0, f"Branch {i} has zero integral" # Sum of branch integrals should equal total - assert np.isclose( - sum(integrals), total_integral, rtol=1e-6 - ), "Sum of branch integrals should equal total integral" + assert np.isclose(sum(integrals), total_integral, rtol=1e-6), ( + "Sum of branch integrals should equal total integral" + )