Improve error handling for the low-k approximation

Author: AgenttiX

pttools/ssmtools/low_k/join.py

@@ -1,8 +1,12 @@
+import logging
+
 import numpy as np
 from scipy.special import erf, erfc
 
 from pttools.ssmtools.low_k import integration, intersection
 
+logger = logging.getLogger(__name__)
+
 
 def pow_gw_junction(
         z: np.ndarray[tuple[int], np.float64],
@@ -12,7 +16,8 @@ def pow_gw_junction(
         cs: float,
         nu: float,
         tau_star: float,
-        tau_end: float):
+        tau_end: float,
+        HLf: float):
     r"""
     Create the junction of the gravitational wave power spectrum between different regimes
     starting from the profiles in each regime.
@@ -22,19 +27,30 @@ def pow_gw_junction(
     :param Pgw_int: array of gravitational wave power spectrum values in the intermediate-frequency regime
     :param Pgw_high: array of gravitational wave power spectrum values in the high-frequency regime
     :param cs: sound speed, $0 < c_s < \frac{1}{\sqrt{3}}$
+    :param nu: $\nu_\text{gdh2024}$
     :param tau_star: $\tau_* = \frac{\eta_*}{L_f}$
     :param tau_end: $\tau_{end} = \frac{\eta_{end}}{L_f}$
+    :param HLf: $H L_f = r_*$
     :return: gravitational wave power spectrum values at the given momentum
     """
-    # z_star = 4*cs*np.pi * (1+nu) / HLf
+    # if not (z.size and Pgw_low.size and Pgw_int.size and Pgw_high.size):
+    #     raise ValueError(
+    #         "Input arrays must not be empty. Got sizes: "
+    #         f"z={z.size}, Pgw_low={Pgw_low.size}, Pgw_int={Pgw_int.size}, Pgw_high={Pgw_high.size}"
+    #     )
+
     difference = Pgw_high - Pgw_int
     index = np.where(difference > 0)[0]
-    z_star = z[index[0]]  # if len(index) > 0 else z_star
+    if index.size:
+        z_star = z[index[0]]
+    else:
+        z_star = 4*cs*np.pi * (1+nu) / HLf
+        logger.warning("Using fallback z_star=%s for low-k junction.", z_star)
     z_cross = intersection.cross_z_junction(cs=cs, nu=nu, tau_star=tau_star, tau_end=tau_end)
 
     term_low = 0.5 * erfc(2 * np.pi * tau_star * (z - z_cross)) * Pgw_low
-    term_int = 0.5 * (1 + erf(2 * np.pi * tau_star * (z - z_cross))) * Pgw_int * 0.5 * erfc(
-        2 * np.pi * tau_star * (z - z_star))
+    term_int = 0.5 * (1 + erf(2 * np.pi * tau_star * (z - z_cross))) * Pgw_int * \
+               0.5 * erfc(2 * np.pi * tau_star * (z - z_star))
     term_high = 0.5 * (1 + erf(2 * np.pi * tau_star * (z - z_star))) * Pgw_high
 
     return term_low + term_int + term_high

pttools/ssmtools/spectrum.py

@@ -67,15 +67,16 @@ def __init__(
         self.label_latex = self.bubble.label_latex[:-1] + label_suffix_latex if label_latex is None else label_latex
         self.label_unicode = self.bubble.label_unicode + label_suffix_unicode if label_unicode is None else label_unicode
 
-        if self.r_star <= 0:
-            raise ValueError("r_star must be positive. Got r_star={r_star}.")
-        if self.r_star >= 1:
-            # Todo: Find a better reference for this.
-            logger.warning(
-                "r_star < 1 is required for the phase transition to complete. "
-                "Got r_star=%s. See Hindmarsh & Hijazi 2019, p. 6.",
-                self.r_star
-            )
+        if not (self.r_star is None or np.isnan(r_star)):
+            if self.r_star <= 0:
+                raise ValueError("r_star must be positive. Got r_star={r_star}.")
+            if self.r_star >= 1:
+                # Todo: Find a better reference for this.
+                logger.warning(
+                    "r_star < 1 is required for the phase transition to complete. "
+                    "Got r_star=%s. See Hindmarsh & Hijazi 2019, p. 6.",
+                    self.r_star
+                )
 
         # Values generated by compute()
         # $|A(z)|^2$
@@ -92,6 +93,8 @@ def __init__(
         self.pow_v: tp.Optional[np.ndarray] = None
         #: $\mathcal{P}_{\text{gw}}(k)$
         self.pow_gw: tp.Optional[np.ndarray] = None
+        #: $\mathcal{P}_{\text{gw}}(k)$, expanded to low frequencies
+        self.pow_gw_expanded: tp.Optional[np.ndarray] = None
         #: $\mathcal{P}_{\text{gw}}(k)$ for intermediate frequencies
         self.pow_gw_int: tp.Optional[np.ndarray] = None
         #: $\mathcal{P}_{\text{gw}}(k)$ for low frequencies
@@ -163,7 +166,8 @@ def compute(self):
                 Pgw_low=self.pow_gw_low * factor,
                 Pgw_int=self.pow_gw_int * factor,
                 Pgw_high=self.pow_gw_ssm * factor,
-                cs=self.cs, nu=self.bubble.nu_gdh2024, tau_star=tau_star, tau_end=tau_end
+                cs=self.cs, nu=self.bubble.nu_gdh2024, tau_star=tau_star, tau_end=tau_end,
+                HLf=self.r_star
             ) / factor
             if self.low_k:
                 self.pow_gw = self.pow_gw_expanded