From f04f52b7478a8113614acda03274a4d829ecb1e3 Mon Sep 17 00:00:00 2001
From: Evgeniy Unzhakov <unzhakov_ev@pnpi.nrcki.ru>
Date: Mon, 27 Apr 2026 15:36:36 +0300
Subject: [PATCH] Update Phase Space calculation

Add support for proper calculation of phase space component based on the
transition type (allowed/Nth order forbidden).
---
 beta_spectrum/components/exchange.py    |  2 +-
 beta_spectrum/components/phase_space.py | 59 +++++++++++++++++++++----
 beta_spectrum/spectrum.py               |  5 ++-
 beta_spectrum/utils.py                  | 14 +++++-
 4 files changed, 67 insertions(+), 13 deletions(-)

diff --git a/beta_spectrum/components/exchange.py b/beta_spectrum/components/exchange.py
index 998be24..f94a1ea 100644
--- a/beta_spectrum/components/exchange.py
+++ b/beta_spectrum/components/exchange.py
@@ -42,7 +42,7 @@ class ExchangeCorrection(SpectrumComponent):
         W = np.asarray(W)
 
         # Define physical cutoff
-        W_cut = 1.005
+        W_cut = 1.003
 
         # Compute full expression ONLY above cutoff
         W_safe = np.maximum(W, W_cut)
diff --git a/beta_spectrum/components/phase_space.py b/beta_spectrum/components/phase_space.py
index 1b690c4..b448423 100644
--- a/beta_spectrum/components/phase_space.py
+++ b/beta_spectrum/components/phase_space.py
@@ -1,19 +1,60 @@
 # components/phase_space.py
 import numpy as np
-from ..base import SpectrumComponent
+from beta_spectrum.base import SpectrumComponent
+from beta_spectrum.utils import momentum
 
 
 class PhaseSpace(SpectrumComponent):
-    def __init__(self, W0: float, m_nu: float = 0.0):
+    """
+    3-body phase space calculation (assuming heavy nucleus and no recoil)
+    """
+
+    def __init__(
+        self,
+        W0: float,
+        transition_type: str = "A",
+        m_e: float = 1.0,
+        m_nu: float = 0.0,
+    ):
+        """
+        Create instance of phase space with given decay parameters.
+
+        Args:
+
+            W0: Endpoint energy in natural units m_e=c=h_bar=1
+            transition_type: allowed or N-order forbidden, should be [A, F1, F1U, F2, F2U, F3, F3U, F4]
+            m_e: 1st body mass (electron by default: m_e = 1)
+            m_nu:2nd body mass (massless neutrino by default: m_nu = 0)
+        """
+        self.tr_type = transition_type
+        if self.tr_type not in ["A", "F1", "F1U", "F2", "F2U", "F3", "F3U", "F4"]:
+            raise ValueError(
+                f"Cannot calculate phase space for forbiddance '{self.tr_type}'.\nShould be [A, F1, F1U, F2, F2U, F3, F3U, F4].\n"
+            )
         self.W0 = W0
+        self.m_e = m_e
         self.m_nu = m_nu
 
     def __call__(self, W: np.ndarray) -> np.ndarray:
-        p = np.sqrt(W**2 - 1.0)
-        Wv = self.W0 - W
+        w_e = W
+        p_e = momentum(W, self.m_e)
+        w_nu = self.W0 - W
+        p_nu = momentum(w_nu, self.m_nu)
 
-        # neutrino mass generalization
-        if self.m_nu == 0:
-            return p * W * Wv**2
-        else:
-            return p * W * Wv * np.sqrt(Wv**2 - self.m_nu**2)
+        phase_space = p_e * w_e * p_nu * w_nu
+
+        if self.tr_type in ["A", "F1"]:
+            forbid_factor = 1
+
+        elif self.tr_type in ["F1U", "F2"]:
+            forbid_factor = p_nu**2 + p_e**2
+
+        elif self.tr_type in ["F2U", "F3"]:
+            forbid_factor = p_nu**4 + 3 / 10 * p_nu**2 * p_e**2 + p_e**4
+
+        elif self.tr_type in ["F3U", "F4"]:
+            forbid_factor = (
+                p_nu**6 + 7 * p_nu**4 * p_e**2 + 7 * p_nu**2 * p_e**4 + p_e**6
+            )
+
+        return phase_space * forbid_factor
diff --git a/beta_spectrum/spectrum.py b/beta_spectrum/spectrum.py
index eab4529..dd983ed 100644
--- a/beta_spectrum/spectrum.py
+++ b/beta_spectrum/spectrum.py
@@ -25,6 +25,7 @@ class SpectrumConfig:
     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
 
     # Toggle components and corrections
@@ -82,7 +83,9 @@ class BetaSpectrum:
         W0 = T_to_W(config.endpoint_MeV)
 
         if config.use_phase_space:
-            components.append(PhaseSpace(W0=W0))
+            components.append(
+                PhaseSpace(W0=W0, transistion_type=config.transition_type)
+            )
 
         if config.use_fermi:
             components.append(FermiFunction(Z=config.Z_daughter, A=config.A_number))
diff --git a/beta_spectrum/utils.py b/beta_spectrum/utils.py
index 324fe0b..a7f326e 100644
--- a/beta_spectrum/utils.py
+++ b/beta_spectrum/utils.py
@@ -11,8 +11,18 @@ def nuclear_radius(A: float, r0: float = 1.2) -> float:
     return R_fm / (HBAR_C_MEV_FM / ME_MEV)
 
 
-def momentum(W):
-    return np.sqrt(W**2 - 1)
+def momentum(W, rest_mass: float = 1):
+    """
+    Helper function to get momentum value from energy.
+
+    Args:
+        W np.ndarray: Array of full particle energy values in natural units
+        rest_mass float: rest mass of the particle in taural units (m_e = 1 by default for electron)
+
+    Returns:
+        np.ndarray: array of particle momentum values.
+    """
+    return np.sqrt(W**2 - rest_mass**2)
 
 
 def neutrino_energy(W, W0):