beta_shape_pnpi/beta_spectrum/spectrum.py

# spectrum.py
from __future__ import annotations

import logging
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass

import numpy as np

from beta_spectrum.base import SpectrumComponent
from beta_spectrum.components.phase_space import PhaseSpace
from beta_spectrum.components.fermi import FermiFunction
from beta_spectrum.components.finite_size import FiniteSizeL0, ChargeDistributionU
from beta_spectrum.components.screening import ScreeningCorrection
from beta_spectrum.components.exchange import ExchangeCorrection
from beta_spectrum.components.radiative import RadiativeCorrection
from beta_spectrum.utils import T_to_W


@dataclass
class BranchConfig:
    """
    Configuration for a single decay branch.

    Each branch corresponds to a decay to a specific daughter nuclear state,
    with its own endpoint energy (Q_value - excitation_energy) and transition type.

    Attributes
    ----------
    endpoint_MeV : float
        Effective endpoint energy for this branch (Q_value - excitation_energy).
    transition_type : str
        Forbiddenness classification for this branch (A, F1, F1U, F2, ...).
    intensity : float
        Branch intensity as a fraction of total decays (0.0–1.0).
        Intensities across all branches should sum to 1.0 (100%).
    """

    endpoint_MeV: float
    transition_type: str = "A"
    intensity: float = 1.0

    def __post_init__(self) -> None:
        if self.intensity < 0 or self.intensity > 1.0:
            raise ValueError(
                f"Branch intensity must be in [0, 1], got {self.intensity}"
            )
        if self.endpoint_MeV <= 0:
            raise ValueError(
                f"Branch endpoint_MeV must be positive, got {self.endpoint_MeV}"
            )


@dataclass
class SpectrumConfig:
    """
    Configuration for beta spectrum calculation.

    Supports declarative detector response specification:
    to enable automatic detector smearing via convolve_with_detector().

    Supports multi-branch decay via the `branches` field. When branches are
    provided, the total spectrum is the intensity-weighted sum of individual
    branch spectra. Each branch has its own endpoint energy and transition type.
    """

    Z_parent: int
    Z_daughter: int
    A_number: int
    endpoint_MeV: float
    transition_type: str = "A"  # should be [A, F1, F1U, F2, F2U, F3, F3U, F4]
    e_step_MeV: float = 0.001
    neutrino_mass_MeV: float = 0.0

    # Toggle components and corrections
    use_phase_space: bool = True
    use_fermi: bool = True
    use_screening: bool = True
    use_finite_size: bool = True
    use_charge_dist: bool = True
    use_radiative: bool = True
    use_exchange: bool = True

    # Multi-branch support
    branches: Optional[List[BranchConfig]] = None  # None → single-branch mode
    intensity_cutoff: float = (
        0.0  # Ignore branches below this fraction (default: 0.0 = all)
    )


class BetaSpectrum:
    """
    Precise beta spectrum calculator.

    Supports both single-branch and multi-branch decay modes.

    In single-branch mode (branches=None), the spectrum is the product of
    all enabled correction factors.

    In multi-branch mode (branches provided), each branch is calculated
    independently with its own endpoint energy and transition type, then
    summed with intensity weighting.

    Parameters
    ----------
    components : List[SpectrumComponent], optional
        List of enabled universal spectral components (Fermi, screening, etc.).
        Only used in multi-branch mode or for universal components.
    branch_spectra : List[BetaSpectrum], optional
        List of per-branch spectrum calculators (each with its own PhaseSpace, etc.).
        Only used in multi-branch mode.
    branches : List[BranchConfig], optional
        Branch configuration metadata (intensities, endpoint energies).
    logger : logging.Logger, optional
        Logger for debug/info output.
    """

    def __init__(
        self,
        components: Optional[List[SpectrumComponent]] = None,
        branch_spectra: Optional[List[BetaSpectrum]] = None,
        branches: Optional[List[BranchConfig]] = None,
        logger: Optional[logging.Logger] = None,
    ):
        self.components = components or []
        self._component_names = [self._get_component_name(c) for c in self.components]
        self.branch_spectra = branch_spectra or []
        self.branches = branches or []
        self._logger = logger
        self._is_multi = len(self.branch_spectra) > 0

    def __call__(self, W: np.ndarray) -> np.ndarray:
        """
        Calculate spectrum.

        In single-branch mode: product of all components.
        In multi-branch mode: intensity-weighted sum of branch spectra.
        """
        if self._is_multi:
            return self._multi_branch_call(W)
        return self._single_branch_call(W)

    def _single_branch_call(self, W: np.ndarray) -> np.ndarray:
        """Calculate spectrum as product of components (single branch)."""
        result: np.ndarray = np.ones_like(W, dtype=float)
        for comp in self.components:
            if self._logger:
                self._logger.debug(
                    "Evaluating %s at %d energy points", comp.__class__.__name__, len(W)
                )
            result = result * comp(W)
        return result.astype(np.float64)

    def _multi_branch_call(self, W: np.ndarray) -> np.ndarray:
        """Calculate total spectrum as intensity-weighted sum of branches.

        The total is: (universal_components_product) * sum(intensity_i * branch_spectrum_i)
        """
        # Sum of intensity-weighted branch spectra
        branch_sum = np.zeros_like(W, dtype=float)
        for i, branch_spec in enumerate(self.branch_spectra):
            if self._logger:
                self._logger.debug(
                    "Evaluating branch %d/%d (endpoint=%.3f MeV, intensity=%.4f)",
                    i + 1,
                    len(self.branch_spectra),
                    self.branches[i].endpoint_MeV,
                    self.branches[i].intensity,
                )
            branch_spectrum = branch_spec(W)
            # Mask out energies beyond this branch's endpoint
            W0_branch = T_to_W(self.branches[i].endpoint_MeV)
            mask = W <= W0_branch
            branch_sum += self.branches[i].intensity * np.where(
                mask, branch_spectrum, 0.0
            )

        # Multiply by universal components
        result: np.ndarray = np.ones_like(W, dtype=float)
        for comp in self.components:
            result = result * comp(W)

        return np.asarray((result * branch_sum).astype(np.float64), dtype=np.float64)

    def _get_component_name(self, comp: SpectrumComponent) -> str:
        """Get human-readable component name."""
        class_name = comp.__class__.__name__
        return class_name.replace("Correction", "").replace("Function", "")

    def calculate_components(self, W: np.ndarray) -> Dict[str, np.ndarray]:
        """
        Calculate individual components.

        In single-branch mode: returns per-component values.
        In multi-branch mode: returns universal components + per-branch
        branch-dependent components (PhaseSpace, Radiative).
        """
        if self._is_multi:
            return self._multi_branch_components(W)

        result: Dict[str, np.ndarray] = {}
        for name, comp in zip(self._component_names, self.components):
            result[name] = comp(W)
        return result

    def _multi_branch_components(self, W: np.ndarray) -> Dict[str, np.ndarray]:
        """Return universal + per-branch components for multi-branch mode."""
        result: Dict[str, np.ndarray] = {}

        # Universal components (same for all branches)
        for name, comp in zip(self._component_names, self.components):
            result[name] = comp(W)

        # Per-branch components (branch-dependent)
        for i, branch_spec in enumerate(self.branch_spectra):
            branch_name = f"branch_{i}"
            branch_components = branch_spec.calculate_components(W)
            for comp_name, values in branch_components.items():
                result[f"{branch_name}.{comp_name}"] = values

        return result

    @classmethod
    def from_config(
        cls, config: SpectrumConfig, logger: Optional[logging.Logger] = None
    ) -> BetaSpectrum:
        """
        Create a BetaSpectrum from configuration.

        In single-branch mode: creates a single spectrum calculator.
        In multi-branch mode: creates per-branch calculators + universal components.

        Parameters
        ----------
        config : SpectrumConfig
            Configuration with decay parameters and correction toggles.
        logger : logging.Logger, optional
            Logger for progress output.
        """
        if config.branches is not None and len(config.branches) > 1:
            return cls._from_config_multi(config, logger)
        return cls._from_config_single(config, logger)

    @classmethod
    def _from_config_single(
        cls, config: SpectrumConfig, logger: Optional[logging.Logger] = None
    ) -> BetaSpectrum:
        """Create single-branch spectrum calculator."""
        components: List[SpectrumComponent] = []
        W0 = float(T_to_W(config.endpoint_MeV))

        if config.use_phase_space:
            components.append(
                PhaseSpace(
                    W0=W0,
                    transition_type=config.transition_type,
                    m_nu=config.neutrino_mass_MeV,
                    logger=logger,
                )
            )

        if config.use_fermi:
            components.append(
                FermiFunction(Z=config.Z_daughter, A=config.A_number, logger=logger)
            )

        if config.use_finite_size:
            components.append(
                FiniteSizeL0(Z=config.Z_daughter, A=config.A_number, logger=logger)
            )

        if config.use_charge_dist:
            components.append(
                ChargeDistributionU(
                    Z=config.Z_daughter, A=config.A_number, logger=logger
                )
            )

        if config.use_screening:
            components.append(
                ScreeningCorrection(
                    FermiFunction(Z=config.Z_parent, A=config.A_number, logger=logger),
                    logger=logger,
                )
            )

        if config.use_exchange:
            components.append(ExchangeCorrection(Z=config.Z_parent, logger=logger))

        if config.use_radiative:
            components.append(
                RadiativeCorrection(
                    W0=W0,
                    Z=config.Z_daughter,
                    A=config.A_number,
                    use_endpoint_resummation=True,
                    logger=logger,
                )
            )

        return cls(components, logger=logger)

    @classmethod
    def _from_config_multi(
        cls, config: SpectrumConfig, logger: Optional[logging.Logger] = None
    ) -> BetaSpectrum:
        """Create multi-branch spectrum calculator.

        Creates per-branch calculators (each with its own PhaseSpace, Radiative, etc.)
        and universal components (Fermi, screening, finite_size, etc.) shared across
        all branches.

        Note: Radiative correction is branch-dependent (depends on W0), so it goes
        into each branch's calculator, not the universal components.
        """
        if logger:
            logger.info("Multi-branch mode: %d branches", len(config.branches))

        # Create per-branch calculators
        branch_spectra: List[BetaSpectrum] = []
        for i, branch in enumerate(config.branches):
            if logger:
                logger.info(
                    "  Branch %d: endpoint=%.3f MeV, transition=%s, intensity=%.4f",
                    i + 1,
                    branch.endpoint_MeV,
                    branch.transition_type,
                    branch.intensity,
                )

            # Create a per-branch config with the branch's endpoint and transition type
            branch_config = SpectrumConfig(
                Z_parent=config.Z_parent,
                Z_daughter=config.Z_daughter,
                A_number=config.A_number,
                endpoint_MeV=branch.endpoint_MeV,
                transition_type=branch.transition_type,
                e_step_MeV=config.e_step_MeV,
                neutrino_mass_MeV=config.neutrino_mass_MeV,
                # Branch-dependent components
                use_phase_space=config.use_phase_space,
                use_radiative=config.use_radiative,
                # Universal components disabled (handled separately)
                use_fermi=False,
                use_screening=False,
                use_finite_size=False,
                use_charge_dist=False,
                use_exchange=False,
            )

            branch_spectrum = cls._from_config_single(branch_config, logger)
            branch_spectra.append(branch_spectrum)

        # Create universal components (shared across all branches)
        # Note: Radiative is NOT included here because it depends on W0
        universal_components: List[SpectrumComponent] = []

        if config.use_fermi:
            universal_components.append(
                FermiFunction(Z=config.Z_daughter, A=config.A_number, logger=logger)
            )

        if config.use_finite_size:
            universal_components.append(
                FiniteSizeL0(Z=config.Z_daughter, A=config.A_number, logger=logger)
            )

        if config.use_charge_dist:
            universal_components.append(
                ChargeDistributionU(
                    Z=config.Z_daughter, A=config.A_number, logger=logger
                )
            )

        if config.use_screening:
            universal_components.append(
                ScreeningCorrection(
                    FermiFunction(Z=config.Z_parent, A=config.A_number, logger=logger),
                    logger=logger,
                )
            )

        if config.use_exchange:
            universal_components.append(
                ExchangeCorrection(Z=config.Z_parent, logger=logger)
            )

        return cls(
            components=universal_components,
            branch_spectra=branch_spectra,
            branches=list(config.branches),
            logger=logger,
        )

    def get_energy_grid(self, config: SpectrumConfig) -> Tuple[np.ndarray, np.ndarray]:
        """
        Generate array of energy values for spectrum calculation.

        In multi-branch mode, uses the maximum endpoint across all branches.
        """
        if self._is_multi and self.branches:
            max_endpoint = max(b.endpoint_MeV for b in self.branches)
        else:
            max_endpoint = config.endpoint_MeV

        kinetic_MeV = np.arange(config.e_step_MeV, max_endpoint, config.e_step_MeV)
        W = T_to_W(kinetic_MeV)

        return W, kinetic_MeV.astype(np.float64)

    def get_branch_spectra(self, W: np.ndarray) -> List[np.ndarray]:
        """
        Get individual branch spectra for multi-branch mode.

        Each branch spectrum is masked to zero beyond its endpoint.

        Parameters
        ----------
        W : np.ndarray
            Energy grid in m_e units.

        Returns
        -------
        List[np.ndarray]
            Per-branch spectrum arrays (unweighted). Empty in single-branch mode.
        """
        if not self._is_multi:
            return []
        result = []
        for i, branch_spec in enumerate(self.branch_spectra):
            branch_spectrum = branch_spec(W)
            # Use np.maximum to smoothly approach zero at endpoint,
            # and clamp negative values beyond endpoint to zero
            result.append(np.maximum(0.0, branch_spectrum))
        return result

    def get_branch_normalized_spectra(self, W: np.ndarray) -> List[np.ndarray]:
        """
        Get intensity-weighted branch spectra for multi-branch mode.

        Each branch spectrum is masked to zero beyond its endpoint.

        Parameters
        ----------
        W : np.ndarray
            Energy grid in m_e units.

        Returns
        -------
        List[np.ndarray]
            Per-branch spectrum arrays weighted by intensity. Empty in single-branch mode.
        """
        if not self._is_multi:
            return []
        result = []
        for i, branch_spec in enumerate(self.branch_spectra):
            branch_spectrum = branch_spec(W)
            result.append(self.branches[i].intensity * np.maximum(0.0, branch_spectrum))
        return result