gh-492: Adding running of the spectral index and curvature emulators

cloelib/cosmology/camb_cosmology.py

@@ -33,6 +33,7 @@ def __init__(
         gamma_MG: float,
         N_mnu: int,
         N_ur: Optional[float] = None,
+        alpha_s: float = 0.0,
     ) -> None:
         """
         Initialize the CAMBBackground instance with cosmological parameters.
@@ -44,7 +45,8 @@ def __init__(
             Omega_k0(float): Curvature density parameter.
             As (float): Scalar amplitude of primordial fluctuations.
             ns (float): Scalar spectral index.
-            mnu (Union[float, Sequence[float], np.ndarray]): Total neutrino mass in eV.
+            alpha_s (float): Running of the scalar spectral index (d ns / d ln k).
+            mnu (Union[float, Sequence[float]], np.ndarray]): Total neutrino mass in eV.
                 Can be a single float for degenerate masses, an array (or a sequence of floats) for individual species.
             w0 (float): Equation of state parameter for dark energy.
             wa (float): Time evolution of the dark energy equation of state.
@@ -60,6 +62,7 @@ def __init__(
         self.Omega_k0 = Omega_k0
         self.As = As
         self.ns = ns
+        self.alpha_s = alpha_s
         self.w0 = w0
         self.wa = wa
         self.gamma_MG = gamma_MG
@@ -100,7 +103,9 @@ def __init__(
         self.interface_args["CAMBparams"].set_dark_energy(
             w=self.w0, wa=self.wa, dark_energy_model="ppf"
         )
-        self.interface_args["CAMBparams"].InitPower.set_params(As=self.As, ns=self.ns)
+        self.interface_args["CAMBparams"].InitPower.set_params(
+            As=self.As, ns=self.ns, nrun=self.alpha_s
+        )
 
         # Call CAMB to compute the background
         self.results = camb.get_background(self.interface_args["CAMBparams"])

cloelib/cosmology/class_cosmology.py

@@ -36,6 +36,8 @@ def __init__(
         gamma_MG: float,
         N_mnu: int,
         N_ur: Optional[float] = None,
+        alpha_s: float = 0.0,
+        **kwargs,
     ) -> None:
         """
         Initialize the CLASSBackground instance with cosmological parameters.
@@ -47,6 +49,7 @@ def __init__(
             Omega_k0 (float): Curvature density parameter.
             As (float): Scalar amplitude of primordial fluctuations.
             ns (float): Scalar spectral index.
+            alpha_s (float): Running of the scalar spectral index (d ns / d ln k).
             mnu (Union[float, Sequence[float], np.ndarray]): Total neutrino mass in eV.
                 Can be a single float for degenerate masses, an array (or a sequence of floats) for individual species.
             w0 (float): Equation of state parameter for dark energy.
@@ -63,6 +66,7 @@ def __init__(
         self.Omega_k0 = Omega_k0
         self.As = As
         self.ns = ns
+        self.alpha_s = alpha_s
         self.w0 = w0
         self.wa = wa
         self.gamma_MG = gamma_MG  # Kept for protocol, but CLASS doesn't directly use it
@@ -87,6 +91,7 @@ def __init__(
         )
         self.interface_args["CLASSparams"]["Omega_k"] = self.Omega_k0
         self.interface_args["CLASSparams"]["n_s"] = self.ns
+        self.interface_args["CLASSparams"]["alpha_s"] = self.alpha_s
         self.interface_args["CLASSparams"]["A_s"] = self.As
         self.interface_args["CLASSparams"]["w0_fld"] = self.w0  # or w0
         self.interface_args["CLASSparams"]["wa_fld"] = self.wa  # or wa

cloelib/cosmology/cosmology.py

@@ -77,6 +77,11 @@ def ns(self) -> float:
         """Scalar index of the primordial power spectrum."""
         ...
 
+    @property
+    def alpha_s(self) -> float:
+        """Running of the scalar spectral index (d ns / d ln k)."""
+        ...
+
     @property
     def w0(self) -> float:
         """Dark energy parameter."""

cloelib/cosmology/jax_cosmology.py

@@ -46,6 +46,8 @@ def __init__(
         gamma_MG: float,
         N_mnu: int,
         N_ur: Optional[float] = None,
+        alpha_s: float = 0.0,
+        **kwargs,
     ):
         """
         Initialize the JAXBackground class.
@@ -57,6 +59,8 @@ def __init__(
             Omega_k0(float): Curvature density parameter.
             As (float): Scalar amplitude of primordial fluctuations.
             ns (float): Scalar spectral index.
+            alpha_s (float): Running of the scalar spectral index (d ns / d ln k).
+                Note: JAXBackground does not use alpha_s (emulator not trained with it).
             mnu (Union[float, Sequence[float], jnp.ndarray]): Total neutrino mass in eV.
                 Can be a single float for degenerate masses, an array (or a sequence of floats) for individual species.
             w0 (float): Equation of state parameter for dark energy.
@@ -74,6 +78,7 @@ def __init__(
         self.Omega_k0 = Omega_k0
         self.As = As
         self.ns = ns
+        self.alpha_s = alpha_s
         self.w0 = w0
         self.wa = wa
         self.gamma_MG = gamma_MG

docs/code_structure/background.md

@@ -34,6 +34,7 @@ Every Background implementation must provide:
 - **`N_mnu`**: Number of massive neutrino species
 - **`As`**: Primordial power spectrum amplitude
 - **`ns`**: Primordial power spectrum spectral index
+- **`alpha_s`**: Running of the spectral index (d ns / d ln k)
 - **`w0`**: Dark energy equation of state parameter
 - **`wa`**: Dark energy evolution parameter
 - **`gamma_MG`**: Modified gravity parameter

docs/code_structure/perturbations.md

@@ -104,6 +104,7 @@ bg = CAMBBackground(
     Omega_cdm0=0.2650,
     As=2.1e-9,
     ns=0.965,
+    alpha_s=0.0,   # running of the spectral index
     # ... other parameters
 )
 
@@ -204,37 +205,47 @@ Fast JAX-based emulator for linear and nonlinear power spectra using [cosmopower
 
 #### Available classes
 
-| Class                                           | Cosmology | Spectrum                    |
-| ----------------------------------------------- | --------- | --------------------------- |
-| `CosmoPowerJAXLCDMPerturbations.Linear`         | ΛCDM      | P(k) linear                 |
-| `CosmoPowerJAXLCDMPerturbations.LinearCB`       | ΛCDM      | P_cb(k) linear              |
-| `CosmoPowerJAXLCDMPerturbations.NonLinear`      | ΛCDM      | P(k) nonlinear (HMcode2020) |
-| `CosmoPowerJAXLCDMPerturbations.NonLinearCB`    | ΛCDM      | P_cb(k) nonlinear           |
-| `CosmoPowerJAXwCDMPerturbations.Linear`         | wCDM      | P(k) linear                 |
-| `CosmoPowerJAXwCDMPerturbations.LinearCB`       | wCDM      | P_cb(k) linear              |
-| `CosmoPowerJAXwCDMPerturbations.NonLinear`      | wCDM      | P(k) nonlinear              |
-| `CosmoPowerJAXwCDMPerturbations.NonLinearCB`    | wCDM      | P_cb(k) nonlinear           |
-| `CosmoPowerJAXw0waCDMPerturbations.Linear`      | w0waCDM   | P(k) linear                 |
-| `CosmoPowerJAXw0waCDMPerturbations.LinearCB`    | w0waCDM   | P_cb(k) linear              |
-| `CosmoPowerJAXw0waCDMPerturbations.NonLinear`   | w0waCDM   | P(k) nonlinear              |
-| `CosmoPowerJAXw0waCDMPerturbations.NonLinearCB` | w0waCDM   | P_cb(k) nonlinear           |
-
-All classes support **N_mnu = 0, 1, 3** massive neutrinos.
+| Class                                                | Cosmology        | Spectrum                    |
+| ---------------------------------------------------- | ---------------- | --------------------------- |
+| `CosmoPowerJAXLCDMPerturbations.Linear`              | ΛCDM             | P(k) linear                 |
+| `CosmoPowerJAXLCDMPerturbations.LinearCB`            | ΛCDM             | P_cb(k) linear              |
+| `CosmoPowerJAXLCDMPerturbations.NonLinear`           | ΛCDM             | P(k) nonlinear (HMcode2020) |
+| `CosmoPowerJAXLCDMPerturbations.NonLinearCB`         | ΛCDM             | P_cb(k) nonlinear           |
+| `CosmoPowerJAXwCDMPerturbations.Linear`              | wCDM             | P(k) linear                 |
+| `CosmoPowerJAXwCDMPerturbations.LinearCB`            | wCDM             | P_cb(k) linear              |
+| `CosmoPowerJAXwCDMPerturbations.NonLinear`           | wCDM             | P(k) nonlinear              |
+| `CosmoPowerJAXwCDMPerturbations.NonLinearCB`         | wCDM             | P_cb(k) nonlinear           |
+| `CosmoPowerJAXw0waCDMPerturbations.Linear`           | w0waCDM          | P(k) linear                 |
+| `CosmoPowerJAXw0waCDMPerturbations.LinearCB`         | w0waCDM          | P_cb(k) linear              |
+| `CosmoPowerJAXw0waCDMPerturbations.NonLinear`        | w0waCDM          | P(k) nonlinear              |
+| `CosmoPowerJAXw0waCDMPerturbations.NonLinearCB`      | w0waCDM          | P_cb(k) nonlinear           |
+| `CosmoPowerJAXCurvaturePerturbations.Linear`         | ΛCDM + curvature | P(k) linear                 |
+| `CosmoPowerJAXCurvaturePerturbations.LinearCB`       | ΛCDM + curvature | P_cb(k) linear              |
+| `CosmoPowerJAXCurvaturePerturbations.NonLinear`      | ΛCDM + curvature | P(k) nonlinear              |
+| `CosmoPowerJAXCurvaturePerturbations.NonLinearCB`    | ΛCDM + curvature | P_cb(k) nonlinear           |
+| `CosmoPowerJAXRunningIndexPerturbations.Linear`      | ΛCDM + α_s       | P(k) linear                 |
+| `CosmoPowerJAXRunningIndexPerturbations.LinearCB`    | ΛCDM + α_s       | P_cb(k) linear              |
+| `CosmoPowerJAXRunningIndexPerturbations.NonLinear`   | ΛCDM + α_s       | P(k) nonlinear              |
+| `CosmoPowerJAXRunningIndexPerturbations.NonLinearCB` | ΛCDM + α_s       | P_cb(k) nonlinear           |
+
+ΛCDM, wCDM, and w0waCDM classes support **N_mnu = 0, 1, 2, 3** massive neutrinos. Curvature and running spectral index classes have neutrino mass fixed at mnu = 0.06 eV.
 
 #### Parameter ranges
 
-| Parameter | ΛCDM         | wCDM         | w0waCDM      |
-| --------- | ------------ | ------------ | ------------ |
-| ombh2     | [0.001, 0.1] | [0.001, 0.1] | [0.001, 0.1] |
-| omch2     | [0.05, 0.9]  | [0.05, 0.9]  | [0.05, 0.9]  |
-| H0        | [20, 100]    | [20, 100]    | [20, 100]    |
-| ns        | [0.6, 1.3]   | [0.6, 1.3]   | [0.6, 1.3]   |
-| lnAs      | [1.61, 5]    | [1.61, 5]    | [1.61, 5]    |
-| z         | [0, 5]       | [0, 5]       | [0, 5]       |
-| w0        | —            | [-3, -0.33]  | [-3, -0.33]  |
-| wa        | —            | —            | [-3, 3]      |
-| mnu       | [0, 1] eV    | [0, 1] eV    | [0, 1] eV    |
-| log10TAGN | [7.6, 8.5]   | [7.6, 8.5]   | [7.6, 8.5]   |
+| Parameter | ΛCDM / wCDM / w0waCDM   | ΛCDM + curvature | ΛCDM + α_s      |
+| --------- | ----------------------- | ---------------- | --------------- |
+| ombh2     | [0.001, 0.1]            | [0.019, 0.025]   | [0.019, 0.025]  |
+| omch2     | [0.05, 0.9]             | [0.09, 0.15]     | [0.09, 0.15]    |
+| H0        | [20, 100]               | [60, 80]         | [60, 80]        |
+| ns        | [0.6, 1.3]              | [0.8, 1.2]       | [0.8, 1.2]      |
+| lnAs      | [1.61, 5]               | [1.6, 4.0]       | [1.6, 4.0]      |
+| z         | [0, 5]                  | [0, 5]           | [0, 5]          |
+| w0        | [-3, -0.33] (wCDM/w0wa) | —                | —               |
+| wa        | [-3, 3] (w0wa only)     | —                | —               |
+| mnu       | [0, 1] eV               | 0.06 eV (fixed)  | 0.06 eV (fixed) |
+| Omega_k0  | 0 (fixed)               | [-0.1, 0.1]      | 0 (fixed)       |
+| alpha_s   | —                       | —                | [-0.1, 0.1]     |
+| log10TAGN | [7.6, 8.5]              | [7.6, 8.5]       | [7.6, 8.5]      |
 
 #### Example
 

docs/code_structure/summary_statistics.md

@@ -21,7 +21,11 @@ This module computes final statistical quantities for likelihood evaluation, inc
 
 These quantities are directly measurable and form the basis for cosmological parameter inference.
 
-!!! warning cloelib does not use internal interpolations. Keep redshift and wavenumber arrays to a maximum of 1500 elements for optimal performance. Otherwise, memory problems may arise. See [Performance Tips](#performance-tips) for details.
+!!! warning cloelib does not use internal interpolations. Keep redshift and wavenumber arrays to a maximum of 1500 elements for optimal performance. Otherwise, memory problems may arise.
+
+## Performance Tips
+
+`cloelib` does not use internal interpolations. Keep redshift and wavenumber arrays to a maximum of 1500 elements for optimal performance. Otherwise, memory problems may arise.
 
 ## Available Summary Statistics
 

tests/test_baccoemu.py

@@ -30,6 +30,7 @@ def camb_background_instance(scope="module"):
         wa=0.0,
         gamma_MG=0.0,
         N_mnu=0,
+        alpha_s=0.0,
     )
     return camb_instance
 

tests/test_bnt.py

@@ -16,6 +16,7 @@
     "w0": -1.0,
     "wa": 0.0,
     "ns": 0.96,
+    "alpha_s": 0.0,
     "As": 2.1e-9,
     "mnu": 0.06,
     "N_mnu": 3,
@@ -33,6 +34,7 @@
             "w0",
             "wa",
             "ns",
+            "alpha_s",
             "As",
             "mnu",
             "gamma_MG",

tests/test_camb_cosmology.py

@@ -26,6 +26,7 @@ def camb_background_instance(scope="module"):
         Omega_k0=0.0,
         As=2e-9,
         ns=0.96,
+        alpha_s=0.0,
         mnu=0.0,
         w0=-1.0,
         wa=0.0,
@@ -159,6 +160,7 @@ def test_camb_set_neutrino_parameters_degenerate():
         Omega_k0=0.0,
         As=2e-9,
         ns=0.96,
+        alpha_s=0.0,
         mnu=0.3,
         w0=-1,
         wa=0,
@@ -181,6 +183,7 @@ def test_camb_set_neutrino_parameters_non_degenerate():
         Omega_k0=0.0,
         As=2e-9,
         ns=0.96,
+        alpha_s=0.0,
         mnu=np.array([0.05, 0.03]),
         w0=-1,
         wa=0,
@@ -204,6 +207,7 @@ def test_camb_set_neutrino_parameters_zero():
         Omega_k0=0.0,
         As=2e-9,
         ns=0.96,
+        alpha_s=0.0,
         mnu=0.0,
         w0=-1,
         wa=0,
@@ -227,6 +231,7 @@ def test_camb_set_neutrino_parameters_wrong_length():
             Omega_k0=0.0,
             As=2e-9,
             ns=0.96,
+            alpha_s=0.0,
             mnu=[0.02, 0.04],
             w0=-1,
             wa=0,
@@ -245,6 +250,7 @@ def test_camb_set_neutrino_parameters_wrong_type():
             Omega_k0=0.0,
             As=2e-9,
             ns=0.96,
+            alpha_s=0.0,
             mnu=None,
             w0=-1,
             wa=0,
@@ -474,6 +480,7 @@ def test_matter_power_spectrum_cb():
         w0=-1.0,
         wa=0.0,
         ns=0.96,
+        alpha_s=0.0,
         mnu=0.1,
         As=2e-9,
         gamma_MG=0.0,
@@ -496,3 +503,49 @@ def test_matter_power_spectrum_cb():
     assert_allclose(
         perturbations_nl.matter_power_spectrum_cb(0, 1), 747.017036, rtol=1.0e-03
     )
+
+
+def test_camb_background_running_spectral_index():
+    """Test that CAMBBackground initializes and runs correctly with non-zero alpha_s."""
+    bg = CAMBBackground(
+        H0=67.7,
+        Omega_b0=0.022 / 0.677**2,
+        Omega_cdm0=0.12 / 0.677**2,
+        Omega_k0=0.0,
+        As=2e-9,
+        ns=0.96,
+        alpha_s=0.02,
+        mnu=0.0,
+        w0=-1.0,
+        wa=0.0,
+        gamma_MG=0.0,
+        N_mnu=0,
+    )
+
+    assert bg.alpha_s == 0.02
+
+    zs = np.linspace(0.0, 2.0, 20)
+    lin = CAMBLinearPerturbations(bg, zs)
+
+    # Power spectrum should be positive and finite
+    pk = lin.matter_power_spectrum(0.0, 1.0)
+    assert pk > 0
+    assert np.isfinite(pk)
+
+    # sigma8 should differ from alpha_s=0 case (running modifies P(k))
+    bg0 = CAMBBackground(
+        H0=67.7,
+        Omega_b0=0.022 / 0.677**2,
+        Omega_cdm0=0.12 / 0.677**2,
+        Omega_k0=0.0,
+        As=2e-9,
+        ns=0.96,
+        alpha_s=0.0,
+        mnu=0.0,
+        w0=-1.0,
+        wa=0.0,
+        gamma_MG=0.0,
+        N_mnu=0,
+    )
+    lin0 = CAMBLinearPerturbations(bg0, zs)
+    assert lin.sigma8_0() != lin0.sigma8_0()

tests/test_camb_photo_z0_compatibility.py

@@ -35,6 +35,7 @@ def camb_photo_setup():
         Omega_k0=0.0,
         As=2e-9,
         ns=0.96,
+        alpha_s=0.0,
         mnu=0.06,
         w0=-1.0,
         wa=0.0,

tests/test_class_cosmology.py

@@ -25,6 +25,7 @@ def class_background_instance(scope="module"):
         Omega_k0=0.0,
         As=2e-9,
         ns=0.96,
+        alpha_s=0.0,
         mnu=0.0,
         w0=-1.0,
         wa=0.0,

tests/test_cmb.py

@@ -32,6 +32,7 @@ def camb_cmb_setup():
         Omega_k0=0.0,
         As=2e-9,
         ns=0.96,
+        alpha_s=0.0,
         mnu=0.06,
         w0=-1.0,
         wa=0.0,

tests/test_cosmology_background.py

@@ -48,6 +48,7 @@ def test_background_required_attributes():
         "Omega_cdm0",
         "H0",
         "ns",
+        "alpha_s",
         "N_ur",
         "N_mnu",
         "N_eff",
@@ -73,6 +74,7 @@ def test_derived_cosmology():
         As=2e-9,
         gamma_MG=0.0,
         N_mnu=1,
+        alpha_s=0.0,
     )
 
     _z_test = np.zeros(1)
@@ -102,6 +104,7 @@ def test_Omega_cb():
         w0=-1.0,
         wa=0.0,
         ns=0.96,
+        alpha_s=0.0,
         mnu=0.1,
         As=2e-9,
         gamma_MG=0.0,