py-SP(k) Model

.gitignore

@@ -42,6 +42,7 @@ testfile*
 *.ipynb
 !CAMBdemo.ipynb
 !ScalEqs.ipynb
+!docs/SPk_demo.ipynb
 *.dll
 *.a
 *.o

.vscode/settings.json

@@ -116,4 +116,8 @@
     },
     "python-envs.defaultEnvManager": "ms-python.python:venv",
     "python-envs.defaultPackageManager": "ms-python.python:pip",
+    "fortran.linter.includePaths": [
+        "${workspaceFolder}/fortran/Releaselib",
+        "${workspaceFolder}/forutils/Releaselib"
+    ]
 }

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,171 @@ 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.
+
+        References:
+        - SP(k) model: `arXiv:2305.09710 <https://arxiv.org/abs/2305.09710>`_,
+            `MNRAS 523, 2247 (2023) <https://doi.org/10.1093/mnras/stad1474>`_
+        - pyspk implementation details and mode definitions:
+            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)``.
+
+    Notes:
+    - SP(k) calibration is defined for ``0 <= z <= 3`` and ``k <= 12 h/Mpc``.
+        Outside these ranges, behavior follows the Fortran implementation.
+    - SP(k) cannot be combined with HMCode baryon-feedback modes
+        (for example ``halofit_version='mead2020_feedback'``).
+    """
+
+    _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,
+        base_model=None,
+        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,
+    ):
+        """
+        Set SP(k) model and relation parameters.
+
+                References:
+                - SP(k) model: `arXiv:2305.09710 <https://arxiv.org/abs/2305.09710>`_,
+                    `MNRAS 523, 2247 (2023) <https://doi.org/10.1093/mnras/stad1474>`_
+                - For mode details and examples, see https://github.com/jemme07/pyspk
+
+                SP(k) modes:
+
+                - ``1`` (``power_law``):
+                    ``f_b/(Omega_b/Omega_m) = a * (M_SO/M_pivot)^b``
+                    where ``a=SPk_fb_a``, ``b=SPk_fb_pow``, ``M_pivot=SPk_fb_pivot``.
+
+                - ``2`` (``cosmo_power_law``; redshift-dependent power law):
+                    ``f_b/(Omega_b/Omega_m) = (exp(alpha)/100) * (M_500c/1e14 M_sun)^(beta-1) * (E(z)/E(0.3))^gamma``
+                    where ``alpha=SPk_alpha``, ``beta=SPk_beta``, ``gamma=SPk_gamma``.
+
+                - ``3`` (``double_power_law``; redshift-dependent double power law):
+                    ``f_b/(Omega_b/Omega_m) = 0.5 * epsilon * ((M_500c/M_pivot)^alpha + (M_500c/M_pivot)^beta) * (E(z)/E(0.3))^gamma``
+                    where ``epsilon=SPk_epsilon``, ``alpha=SPk_alpha``, ``beta=SPk_beta``,
+                    ``gamma=SPk_gamma``, ``M_pivot=SPk_m_pivot``.
+
+        Parameter definitions:
+
+        :param base_model: Base non-linear model instance to wrap.
+            If None, keeps current base model (default Halofit).
+        :param SPk_feedback: If True, apply SP(k) suppression on top of the base model.
+        :param SPk_SO: Spherical overdensity calibration. Allowed values: 200 or 500.
+        :param SPk_relation_kind: Relation type.
+            Allowed values: 1 (power_law), 2 (cosmo_power_law), 3 (double_power_law).
+        :param SPk_fb_a: Power-law normalization.
+        :param SPk_fb_pow: Power-law exponent.
+        :param SPk_fb_pivot: Power-law pivot mass in solar masses.
+        :param SPk_alpha: Alpha parameter.
+        :param SPk_beta: Beta parameter.
+        :param SPk_gamma: Gamma parameter.
+        :param SPk_epsilon: Epsilon parameter.
+        :param SPk_m_pivot: Pivot mass in solar masses.
+        :return: Self, for fluent configuration.
+        :raises CAMBValueError: If relation or pivot constraints are invalid,
+            or if configuration is incompatible with the selected Halofit
+            baryon-feedback mode.
+        """
+        if base_model is not None:
+            self.BaseModel = base_model
+        elif self.BaseModel is None:
+            self.BaseModel = Halofit()
+
+        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/camb_test.py

@@ -248,21 +248,21 @@ def testBackground(self):
 
         # Test BBN consistency, base_plikHM_TT_lowTEB best fit model
         pars.set_cosmology(H0=67.31, ombh2=0.022242, omch2=0.11977, mnu=0.06, omk=0)
-        self.assertAlmostEqual(pars.YHe, 0.2458, 5)
+        self.assertAlmostEqual(float(pars.YHe), 0.2458, 5)
         data.calc_background(pars)
         self.assertAlmostEqual(data.cosmomc_theta(), 0.0104090741, 7)
         self.assertAlmostEqual(data.get_derived_params()["kd"], 0.14055, 4)
 
         pars.set_cosmology(
             H0=67.31, ombh2=0.022242, omch2=0.11977, mnu=0.06, omk=0, bbn_predictor=bbn.BBN_table_interpolator()
         )
-        self.assertAlmostEqual(pars.YHe, 0.2458, 5)
+        self.assertAlmostEqual(float(pars.YHe), 0.2458, 5)
         self.assertAlmostEqual(pars.get_Y_p(), bbn.BBN_table_interpolator().Y_p(0.022242, 0), 5)
 
         # test massive sterile models as in Planck papers
         pars.set_cosmology(H0=68.0, ombh2=0.022305, omch2=0.11873, mnu=0.06, nnu=3.073, omk=0, meffsterile=0.013)
         self.assertAlmostEqual(pars.omnuh2, 0.00078, 5)
-        self.assertAlmostEqual(pars.YHe, 0.246218, 5)
+        self.assertAlmostEqual(float(pars.YHe), 0.246218, 5)
         self.assertAlmostEqual(pars.N_eff, 3.073, 4)
 
         data.calc_background(pars)
@@ -959,3 +959,28 @@ def test_quintessence(self):
         camb.get_background(pars)
         results = camb.get_results(pars)
         self.assertAlmostEqual(results.get_derived_params()["thetastar"], 1.044341764253, delta=1e-5)
+
+    def testSPkNonLinearClassSelection(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 = model.NonLinear_both
+        pars.NonLinearModel.BaseModel.set_params(halofit_version="mead2020")
+        pars.NonLinearModel.set_params(
+            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)))