From 6dc1e3924dd3eb35c1063af855dd25f48989b18f Mon Sep 17 00:00:00 2001 From: Evgeniy Unzhakov Date: Mon, 6 Jul 2026 17:32:56 +0300 Subject: [PATCH] feat: add CLI options for custom CSV separator and normalize-to-one branch normalization --- beta_spectrum/cli.py | 17 ++++++++++ beta_spectrum/nuclear_data.py | 12 ++++++++ beta_spectrum/spectrum.py | 6 ++++ beta_spectrum/visualize/__init__.py | 48 ++++++++++++++++++++++++----- tests/physics/test_multi_branch.py | 6 ++-- 5 files changed, 78 insertions(+), 11 deletions(-) diff --git a/beta_spectrum/cli.py b/beta_spectrum/cli.py index da62731..5eb0004 100644 --- a/beta_spectrum/cli.py +++ b/beta_spectrum/cli.py @@ -152,6 +152,17 @@ def _build_parser() -> argparse.ArgumentParser: choices=["full", "concise"], help="CSV output verbosity: 'full' (all components) or 'concise' (total spectrum and normalized branch components only) (default: full)", ) + log_group.add_argument( + "--csv-separator", + type=str, + default=',', + help="CSV column separator (default: ','). Use 'tab' for tab separation.", + ) + log_group.add_argument( + "--normalize-to-one", + action="store_true", + help="Normalize each branch to 1, ignoring intensity. In this case, total spectrum will not be produced (only energy and branch-wise spectra).", + ) return parser @@ -253,6 +264,8 @@ def _run(args: argparse.Namespace) -> None: neutrino_mass_MeV=args.neutrino_mass, branches=branches, intensity_cutoff=args.intensity_cutoff, + csv_separator=args.csv_separator, + normalize_to_one=args.normalize_to_one, ) else: config = create_config_from_source( @@ -261,6 +274,8 @@ def _run(args: argparse.Namespace) -> None: e_step_MeV=args.e_step, intensity_cutoff=args.intensity_cutoff, neutrino_mass_MeV=args.neutrino_mass, + csv_separator=args.csv_separator, + normalize_to_one=args.normalize_to_one, ) else: config = create_config_from_source( @@ -269,6 +284,8 @@ def _run(args: argparse.Namespace) -> None: e_step_MeV=args.e_step, intensity_cutoff=args.intensity_cutoff, neutrino_mass_MeV=args.neutrino_mass, + csv_separator=args.csv_separator, + normalize_to_one=args.normalize_to_one, ) else: logger.info("Source: JSON file=%s", source_value) diff --git a/beta_spectrum/nuclear_data.py b/beta_spectrum/nuclear_data.py index 171c558..205ba53 100644 --- a/beta_spectrum/nuclear_data.py +++ b/beta_spectrum/nuclear_data.py @@ -471,6 +471,8 @@ def decay_info_to_config( e_step_MeV: float = 0.001, intensity_cutoff: float = 0.0, neutrino_mass_MeV: float = 0.0, + csv_separator: str = ',', + normalize_to_one: bool = False, ) -> SpectrumConfig: """ Convert DecayInfo to a SpectrumConfig. @@ -485,6 +487,10 @@ def decay_info_to_config( Minimum branch intensity fraction to include. neutrino_mass_MeV : float Neutrino mass in MeV. + csv_separator : str + CSV column separator. + normalize_to_one : bool + If True, normalize each branch to 1, ignoring intensity. Returns ------- @@ -516,6 +522,8 @@ def decay_info_to_config( neutrino_mass_MeV=neutrino_mass_MeV, branches=branches, intensity_cutoff=intensity_cutoff, + csv_separator=csv_separator, + normalize_to_one=normalize_to_one, ) @@ -690,12 +698,16 @@ def create_config_from_source( e_step = kwargs.get("e_step_MeV", 0.001) intensity_cutoff = kwargs.get("intensity_cutoff", 0.0) neutrino_mass = kwargs.get("neutrino_mass_MeV", 0.0) + csv_separator = kwargs.get("csv_separator", ',') + normalize_to_one = kwargs.get("normalize_to_one", False) info = get_decay_info_from_paceENSDF(nuclide, decay_type, decay_index) return decay_info_to_config( info, e_step_MeV=e_step, intensity_cutoff=intensity_cutoff, neutrino_mass_MeV=neutrino_mass, + csv_separator=csv_separator, + normalize_to_one=normalize_to_one, ) elif source == "json": if not json_path: diff --git a/beta_spectrum/spectrum.py b/beta_spectrum/spectrum.py index 62a99ea..9336779 100644 --- a/beta_spectrum/spectrum.py +++ b/beta_spectrum/spectrum.py @@ -87,6 +87,10 @@ class SpectrumConfig: intensity_cutoff: float = ( 0.0 # Ignore branches below this fraction (default: 0.0 = all) ) + + # CSV output options + csv_separator: str = ',' # Default separator for CSV output + normalize_to_one: bool = False # If True, normalize each branch to 1, ignoring intensity class BetaSpectrum: @@ -321,6 +325,8 @@ class BetaSpectrum: # Create per-branch calculators branch_spectra: List[BetaSpectrum] = [] + if config.branches is None: + raise ValueError("branches must be provided for multi-branch mode") for i, branch in enumerate(config.branches): if logger: logger.info( diff --git a/beta_spectrum/visualize/__init__.py b/beta_spectrum/visualize/__init__.py index 19df3d0..27fb73f 100644 --- a/beta_spectrum/visualize/__init__.py +++ b/beta_spectrum/visualize/__init__.py @@ -731,7 +731,10 @@ class BetaSpectrumAnalyzer: header_lines.append(f"# n_points: {len(self.energies_MeV)}") # Add normalization mode - header_lines.append("# branch normalization: intensity") + if self.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 = [] @@ -771,13 +774,31 @@ class BetaSpectrumAnalyzer: # Create data dictionary with concise column headers data: Dict[str, np.ndarray] = { "energy_MeV": self.energies_MeV, - "total_spectrum": total, } + + # Only add total_spectrum if not normalizing to one + if not self.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(self._n_branches): - data[f"branch_{i + 1}_spectrum"] = self._get_branch_spectrum(i) + # Get branch spectrum + branch_spectrum = self._get_branch_spectrum(i) + + if self.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 * self._unified_branches[i].intensity + + data[f"branch_{i + 1}_spectrum"] = normalized_spectrum + for comp_name in ["PhaseSpace", "Radiative"]: key = self._get_component_key(i, comp_name) if key in components: @@ -799,14 +820,25 @@ class BetaSpectrumAnalyzer: data[f"branch_{i + 1}"] = 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 + if self.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: - data[f"branch_{i + 1}"] = branch_spectrum + # Normalize as usual with intensity + sum_spectrum = np.sum(branch_spectrum) + if sum_spectrum > 0: + data[f"branch_{i + 1}"] = branch_spectrum / sum_spectrum * self._unified_branches[i].intensity + else: + data[f"branch_{i + 1}"] = branch_spectrum df = pd.DataFrame(data) - df.to_csv(filename, index=False, float_format="%.4e", sep='\t') + # Determine separator + separator = '\t' if self.config.csv_separator == 'tab' else self.config.csv_separator + df.to_csv(filename, index=False, float_format="%.4e", sep=separator) with open(filename, "r") as f: content = f.read() diff --git a/tests/physics/test_multi_branch.py b/tests/physics/test_multi_branch.py index bc4e84e..1036158 100644 --- a/tests/physics/test_multi_branch.py +++ b/tests/physics/test_multi_branch.py @@ -460,8 +460,8 @@ class TestMultiBranchCSVExport: import pandas as pd - # Read CSV (skip comment headers, use tab delimiter) - df = pd.read_csv(csv_path, comment="#", sep='\t') + # Read CSV (skip comment headers, use comma delimiter by default) + df = pd.read_csv(csv_path, comment="#", sep=',') # Should have total spectrum column assert "total_spectrum" in df.columns @@ -500,7 +500,7 @@ class TestMultiBranchCSVExport: import pandas as pd - df = pd.read_csv(csv_path, comment="#", sep='\t') + df = pd.read_csv(csv_path, comment="#", sep=',') assert "total_spectrum" in df.columns assert "branch_1_spectrum" in df.columns