py-SP(k) Model

.gitignore

@@ -42,6 +42,7 @@ testfile*
 *.ipynb
 !CAMBdemo.ipynb
 !ScalEqs.ipynb
+!docs/SPk_demo.ipynb
 *.dll
 *.a
 *.o
@@ -75,6 +76,7 @@ lensing_cgrads.f90
 !inifiles/params_21cm.ini
 !inifiles/params_counts.ini
 !inifiles/params_lensing.ini
+!inifiles/params_spk.ini
 !inifiles/planck_2018.ini
 !inifiles/planck_2018_acc.ini
 !inifiles/sources_defaults.ini

camb/__init__.py

@@ -43,6 +43,6 @@
 from .initialpower import InitialPowerLaw, SplinedInitialPower
 from .mathutils import threej
 from .model import CAMBparams, TransferParams
-from .nonlinear import Halofit
+from .nonlinear import Halofit, SPkNonLinear
 from .reionization import ExpReionization, TanhReionization
 from .results import CAMBdata, ClTransferData, MatterTransferData

camb/nonlinear.py

@@ -1,9 +1,10 @@
-from ctypes import POINTER, byref, c_double, c_int
+import math
+from ctypes import POINTER, byref, c_bool, c_double, c_int
 
 import numpy as np
 from numpy.ctypeslib import ndpointer
 
-from .baseconfig import F2003Class, fortran_class, numpy_1d
+from .baseconfig import AllocatableObject, CAMBValueError, F2003Class, fortran_class, numpy_1d
 
 
 class NonLinearModel(F2003Class):
@@ -93,6 +94,187 @@ def set_params(
         self.HMCode_logT_AGN = HMCode_logT_AGN
 
 
+@fortran_class
+class SPkNonLinear(NonLinearModel):
+    """SP(k) baryon suppression model applied on top of a base non-linear model."""
+
+    _fields_ = (
+        ("BaseModel", AllocatableObject(NonLinearModel)),
+        ("SPk_feedback", c_bool, "Enable SP(k) suppression"),
+        ("SPk_SO", c_int, "SP(k) spherical overdensity (200 or 500)"),
+        (
+            "SPk_relation_kind",
+            c_int,
+            "SP(k) relation kind: 1=power_law, 2=cosmo_power_law, 3=double_power_law",
+        ),
+        ("SPk_fb_a", c_double, "Power-law relation normalization"),
+        ("SPk_fb_pow", c_double, "Power-law relation exponent"),
+        ("SPk_fb_pivot", c_double, "Power-law relation pivot mass [M_sun]"),
+        ("SPk_alpha", c_double, "Relation alpha parameter (kinds 2/3)"),
+        ("SPk_beta", c_double, "Relation beta parameter (kinds 2/3)"),
+        ("SPk_gamma", c_double, "Relation gamma parameter (kinds 2/3)"),
+        ("SPk_epsilon", c_double, "Relation epsilon parameter (kind 3)"),
+        ("SPk_m_pivot", c_double, "Relation pivot mass [M_sun] (kind 3)"),
+    )
+
+    _fortran_class_module_ = "SPkNonLinear"
+    _fortran_class_name_ = "TSPkNonLinear"
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.BaseModel = Halofit()
+        self.set_params(**kwargs)
+
+    def _validate(self):
+        if self.SPk_SO not in (200, 500):
+            raise CAMBValueError("SPk_SO must be 200 or 500")
+        if self.SPk_relation_kind not in (1, 2, 3):
+            raise CAMBValueError(
+                "SPk_relation_kind must be 1 (power_law), 2 (cosmo_power_law), or 3 (double_power_law)"
+            )
+        if self.SPk_relation_kind == 1 and self.SPk_fb_pivot <= 0:
+            raise CAMBValueError("SPk_fb_pivot must be > 0 for power_law relation")
+        if self.SPk_relation_kind == 3 and self.SPk_m_pivot <= 0:
+            raise CAMBValueError("SPk_m_pivot must be > 0 for double_power_law relation")
+
+        if self.SPk_feedback and isinstance(self.BaseModel, Halofit):
+            if isinstance(self.BaseModel.halofit_version, str):
+                halofit_version_int = halofit_version_names[self.BaseModel.halofit_version]
+            else:
+                halofit_version_int = int(self.BaseModel.halofit_version)
+            if halofit_version_int == halofit_version_names[halofit_mead2020_feedback]:
+                raise CAMBValueError(
+                    "SP(k) is not compatible with halofit_version='mead2020_feedback'. "
+                    "Use halofit_version='mead2020' (or another non-feedback option) when enabling SPk_feedback."
+                )
+
+            hmcode_2015_2016_versions = {
+                halofit_version_names[halofit_mead],
+                halofit_version_names[halofit_mead2015],
+                halofit_version_names[halofit_mead2016],
+            }
+            if halofit_version_int in hmcode_2015_2016_versions and (
+                (not math.isclose(self.BaseModel.HMCode_A_baryon, 3.13, rel_tol=0.0, abs_tol=1e-12))
+                or (not math.isclose(self.BaseModel.HMCode_eta_baryon, 0.603, rel_tol=0.0, abs_tol=1e-12))
+            ):
+                raise CAMBValueError(
+                    "SP(k) cannot be combined with HMCode_A_baryon/HMCode_eta_baryon baryonic corrections in HMCode 2015/2016"
+                )
+
+    def set_params(
+        self,
+        SPk_feedback=False,
+        SPk_SO=200,
+        SPk_relation_kind=1,
+        SPk_fb_a=1.0,
+        SPk_fb_pow=0.0,
+        SPk_fb_pivot=1.0,
+        SPk_alpha=0.0,
+        SPk_beta=0.0,
+        SPk_gamma=0.0,
+        SPk_epsilon=0.0,
+        SPk_m_pivot=1.0,
+        halofit_version=halofit_default,
+    ):
+        """
+        Configure the SP(k) baryon suppression model.
+
+        References:
+          - SP(k) model: `MNRAS 523, 2247 (2023) <https://doi.org/10.1093/mnras/stad1474>`_
+          - pyspk: https://github.com/jemme07/pyspk
+
+        The base model is evaluated first (Halofit by default), then SP(k)
+        suppression is applied to CAMB's non-linear ratio as:
+
+        ``sqrt(P_NL/P_L) -> sqrt(P_NL/P_L) * sqrt(SPk_suppression)``
+
+        **SP(k) relation kinds:**
+
+        - **kind=1** (power_law):
+          ``f_b / (Omega_b/Omega_m) = SPk_fb_a * (M / SPk_fb_pivot)^SPk_fb_pow``
+
+        - **kind=2** (cosmo_power_law):
+          ``f_b / (Omega_b/Omega_m) = (exp(SPk_alpha)/100) * (M_500c/1e14)^(SPk_beta - 1) * (E(z)/E(0.3))^SPk_gamma``
+
+        - **kind=3** (double_power_law):
+          ``f_b / (Omega_b/Omega_m) = 0.5 * SPk_epsilon * ((M/SPk_m_pivot)^SPk_alpha + (M/SPk_m_pivot)^SPk_beta) * (E(z)/E(0.3))^SPk_gamma``
+
+        :param SPk_feedback: If True, apply SP(k) suppression on top of the base model.
+        :param SPk_SO: Spherical overdensity calibration (200 or 500).
+        :param SPk_relation_kind: Relation type: 1 (power_law), 2 (cosmo_power_law), 3 (double_power_law).
+        :param SPk_fb_a: Power-law normalization (kind=1).
+        :param SPk_fb_pow: Power-law exponent (kind=1).
+        :param SPk_fb_pivot: Power-law pivot mass in M_sun (kind=1).
+        :param SPk_alpha: Alpha parameter (kinds 2, 3).
+        :param SPk_beta: Beta parameter (kinds 2, 3).
+        :param SPk_gamma: Gamma parameter (kinds 2, 3).
+        :param SPk_epsilon: Epsilon parameter (kind=3).
+        :param SPk_m_pivot: Pivot mass in M_sun (kind=3).
+        :param halofit_version: Base Halofit version for the wrapped non-linear model.
+        :return: Self, for fluent configuration.
+        :raises CAMBValueError: If parameters are invalid or incompatible with the base model.
+
+        **Cobaya usage:**
+
+        Cobaya passes keys from ``extra_args`` directly to ``set_params()``.
+        Parameters under the theory ``params:`` block are also forwarded and can be sampled.
+
+        - **extra_args** (fixed): ``non_linear_model``, ``halofit_version``, ``SPk_feedback``,
+          ``SPk_SO``, ``SPk_relation_kind``, and pivot masses.
+        - **params** (sampled): continuous relation parameters (e.g. ``SPk_fb_a``, ``SPk_fb_pow``).
+
+        Example YAML (kind=3, double_power_law)::
+
+            params:
+              SPk_epsilon:
+                prior: {min: 0.24, max: 0.35}
+                ref: {dist: norm, loc: 0.30, scale: 0.02}
+              SPk_alpha:
+                prior: {min: -0.12, max: 0.34}
+              SPk_beta:
+                prior: {min: -0.74, max: 0.77}
+              SPk_gamma:
+                prior: {min: -0.5, max: 1.20}
+              log10_SPk_m_pivot:
+                prior: {min: 13, max: 14}
+                drop: true
+              SPk_m_pivot:
+                value: "lambda log10_SPk_m_pivot: 10**log10_SPk_m_pivot"
+
+            theory:
+              camb:
+                extra_args:
+                  non_linear_model: SPkNonLinear
+                  halofit_version: mead2020
+                  SPk_feedback: true
+                  SPk_SO: 200
+                  SPk_relation_kind: 3
+
+        **Notes:**
+
+        - Calibrated for ``0 <= z <= 3`` and ``k <= 12 h/Mpc``.
+        - Cannot be combined with ``halofit_version='mead2020_feedback'``.
+        """
+        if self.BaseModel is None:
+            self.BaseModel = Halofit()
+        if isinstance(self.BaseModel, Halofit):
+            self.BaseModel.set_params(halofit_version=halofit_version)
+
+        self.SPk_feedback = SPk_feedback
+        self.SPk_SO = SPk_SO
+        self.SPk_relation_kind = SPk_relation_kind
+        self.SPk_fb_a = SPk_fb_a
+        self.SPk_fb_pow = SPk_fb_pow
+        self.SPk_fb_pivot = SPk_fb_pivot
+        self.SPk_alpha = SPk_alpha
+        self.SPk_beta = SPk_beta
+        self.SPk_gamma = SPk_gamma
+        self.SPk_epsilon = SPk_epsilon
+        self.SPk_m_pivot = SPk_m_pivot
+        self._validate()
+        return self
+
+
 @fortran_class
 class SecondOrderPK(NonLinearModel):
     """

camb/tests/spk_test.py

@@ -0,0 +1,131 @@
+import inspect
+import os
+import sys
+import unittest
+
+import numpy as np
+
+try:
+    import camb
+except ImportError:
+    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
+    import camb
+
+from camb.nonlinear import Halofit, SPkNonLinear  # type: ignore[attr-defined]
+
+
+class SPkTest(unittest.TestCase):
+    def _get_pk(self, model_obj, z=0.5, kmax=5.0):
+        pars = camb.CAMBparams()
+        pars.set_cosmology(H0=67.5, ombh2=0.02237, omch2=0.12, mnu=0.06)
+        pars.InitPower.set_params(As=2.1e-9, ns=0.965)
+        pars.set_matter_power(redshifts=[z], kmax=kmax, k_per_logint=100)
+        pars.NonLinear = camb.model.NonLinear_both
+        pars.NonLinearModel = model_obj
+        results = camb.get_results(pars)
+        kh, _z, pk = results.get_nonlinear_matter_power_spectrum()
+        return kh, pk[0], results
+
+    def test_spk_invalid_params(self):
+        model = SPkNonLinear()
+        with self.assertRaises(camb.CAMBValueError):
+            model.set_params(SPk_SO=300)
+        with self.assertRaises(camb.CAMBValueError):
+            model.set_params(SPk_relation_kind=99)
+        with self.assertRaises(camb.CAMBValueError):
+            model.set_params(SPk_relation_kind=1, SPk_fb_pivot=0.0)
+
+    def test_spk_accepts_halofit_version(self):
+        model = SPkNonLinear()
+        model.set_params(
+            halofit_version="mead2016",
+            SPk_feedback=True,
+            SPk_SO=200,
+            SPk_relation_kind=1,
+            SPk_fb_a=0.4,
+            SPk_fb_pow=0.2,
+            SPk_fb_pivot=1e14,
+        )
+        k, pk, _ = self._get_pk(model, z=0.5, kmax=3.0)
+        self.assertTrue(np.all(np.isfinite(k)))
+        self.assertTrue(np.all(np.isfinite(pk)))
+
+    def test_spk_rejects_mead2020_feedback_via_set_params(self):
+        model = SPkNonLinear()
+        with self.assertRaises(camb.CAMBValueError):
+            model.set_params(SPk_feedback=True, halofit_version="mead2020_feedback")
+
+    def test_spk_cobaya_friendly_set_params_signature(self):
+        signature = inspect.signature(SPkNonLinear.set_params)
+        self.assertIn("halofit_version", signature.parameters)
+
+    def test_spk_disabled_matches_base(self):
+        base = Halofit()
+        base.set_params(halofit_version="mead2020")
+        k_base, pk_base, _ = self._get_pk(base)
+
+        spk = SPkNonLinear()
+        spk.set_params(halofit_version="mead2020", SPk_feedback=False)
+        k_spk, pk_spk, _ = self._get_pk(spk)
+
+        self.assertTrue(np.allclose(k_base, k_spk, rtol=0, atol=0))
+        self.assertTrue(np.allclose(pk_base, pk_spk, rtol=2e-12, atol=1e-14))
+
+    def test_spk_out_of_range_behaviour(self):
+        """Verify suppression is skipped for z outside calibrated range and k is clamped."""
+        # z=4 is beyond calibrated range [0, 3]: suppression should not be applied.
+        base = Halofit()
+        base.set_params(halofit_version="mead2020")
+        k_base, pk_base_z4, _ = self._get_pk(base, z=4.0, kmax=20.0)
+
+        spk = SPkNonLinear()
+        spk.set_params(
+            halofit_version="mead2020",
+            SPk_feedback=True,
+            SPk_SO=200,
+            SPk_relation_kind=1,
+            SPk_fb_a=0.4,
+            SPk_fb_pow=0.3,
+            SPk_fb_pivot=1e14,
+        )
+        k_spk, pk_spk_z4, _ = self._get_pk(spk, z=4.0, kmax=20.0)
+
+        # At z=4 (out of range), SPk should be identity — P(k) unchanged.
+        np.testing.assert_allclose(pk_spk_z4, pk_base_z4, rtol=1e-10)
+
+        # At z=0.5 (in range) with k up to 20, k > 12 is clamped — suppression still applied.
+        _, pk_base_z05, _ = self._get_pk(base, z=0.5, kmax=20.0)
+        _, pk_spk_z05, _ = self._get_pk(spk, z=0.5, kmax=20.0)
+        sup = pk_spk_z05 / pk_base_z05
+        # Suppression should differ from 1 for k in calibrated range.
+        mask = (k_spk > 0.1) & (k_spk <= 12.0)
+        self.assertFalse(np.allclose(sup[mask], 1.0, atol=1e-4))
+
+    def test_spk_class_selection_via_set_classes(self):
+        pars = camb.CAMBparams()
+        pars.set_classes(non_linear_model="SPkNonLinear")
+        self.assertEqual(pars.NonLinearModel.__class__.__name__, "SPkNonLinear")
+
+        pars.set_cosmology(H0=67.5, ombh2=0.02237, omch2=0.12, mnu=0.06)
+        pars.InitPower.set_params(As=2.1e-9, ns=0.965)
+        pars.set_matter_power(redshifts=[0.5], kmax=3.0)
+        pars.NonLinear = camb.model.NonLinear_both
+        pars.NonLinearModel.set_params(
+            halofit_version="mead2020",
+            SPk_feedback=True,
+            SPk_SO=200,
+            SPk_relation_kind=1,
+            SPk_fb_a=0.4,
+            SPk_fb_pow=0.2,
+            SPk_fb_pivot=1e14,
+        )
+
+        data = camb.get_results(pars)
+        k, z, pk = data.get_matter_power_spectrum(minkh=1e-2, maxkh=1.0, npoints=8)
+        self.assertEqual(len(z), 1)
+        self.assertTrue(np.all(np.isfinite(k)))
+        self.assertTrue(np.all(np.isfinite(pk)))
+
+
+if __name__ == "__main__":
+    unittest.main()