Container for MG params (#1119)

* Container for mg params

* Update doc string

* pacify flake8

* add test

* Another test

* added doc

* pacify flake8

* switch to AE spelling

* fix test

* More spelling changes

benchmarks/test_correlation_MG.py

@@ -1,6 +1,8 @@
 import os
 import numpy as np
 import pyccl as ccl
+from pyccl.modified_gravity import MuSigmaMG
+
 from scipy.interpolate import interp1d
 import pytest
 
@@ -22,7 +24,7 @@ def set_up(request):
     cosmo = ccl.Cosmology(Omega_c=0.12/h0**2, Omega_b=0.0221/h0**2, Omega_k=0,
                           h=h0, A_s=np.exp(logA)/10**10, n_s=0.96, Neff=3.046,
                           m_nu=0.0, w0=-1, wa=0, T_CMB=2.7255,
-                          mu_0=0.1, sigma_0=0.1,
+                          mg_parametrization=MuSigmaMG(mu_0=0.1, sigma_0=0.1),
                           transfer_function='boltzmann_class',
                           matter_power_spectrum='linear')
 

benchmarks/test_correlation_MG2.py

@@ -1,6 +1,8 @@
 import os
 import numpy as np
 import pyccl as ccl
+from pyccl.modified_gravity import MuSigmaMG
+
 from scipy.interpolate import interp1d
 import pytest
 
@@ -22,7 +24,7 @@ def set_up(request):
     cosmo = ccl.Cosmology(Omega_c=0.12/h0**2, Omega_b=0.0221/h0**2, Omega_k=0,
                           h=h0, A_s=np.exp(logA)/10**10, n_s=0.96, Neff=3.046,
                           m_nu=0.0, w0=-1, wa=0, T_CMB=2.7255,
-                          mu_0=0.1, sigma_0=0.1,
+                          mg_parametrization=MuSigmaMG(mu_0=0.1, sigma_0=0.1),
                           transfer_function='boltzmann_isitgr',
                           matter_power_spectrum='linear')
 

benchmarks/test_correlation_MG3_SD.py

@@ -1,6 +1,8 @@
 import os
 import numpy as np
 import pyccl as ccl
+from pyccl.modified_gravity import MuSigmaMG
+
 from scipy.interpolate import interp1d
 import pytest
 
@@ -23,8 +25,9 @@ def set_up(request):
     cosmo = ccl.Cosmology(Omega_c=0.12/h0**2, Omega_b=0.0221/h0**2, Omega_k=0,
                           h=h0, A_s=np.exp(logA)/10**10, n_s=0.96, Neff=3.046,
                           m_nu=0.0, w0=-1, wa=0, T_CMB=2.7255,
-                          mu_0=0.1, sigma_0=0.1,
-                          c1_mg=1.1, c2_mg=1.1, lambda_mg=1,
+                          mg_parametrization=MuSigmaMG(
+                              mu_0=0.1, sigma_0=0.1,
+                              c1_mg=1.1, c2_mg=1.1, lambda_mg=1),
                           transfer_function='boltzmann_isitgr',
                           matter_power_spectrum='linear')
 

benchmarks/test_distances.py

@@ -1,5 +1,7 @@
 import numpy as np
 import pyccl as ccl
+from pyccl.modified_gravity import MuSigmaMG
+
 import pytest
 
 # Set tolerances
@@ -463,7 +465,8 @@ def compare_distances_muSig(z, chi_bench, dm_bench, Omega_v, w0, wa):
     cosmo = ccl.Cosmology(
         Omega_c=Omega_c, Omega_b=Omega_b, Neff=Neff,
         h=h, A_s=A_s, n_s=n_s, Omega_k=Omega_k,
-        w0=w0, wa=wa, mu_0=mu_0, sigma_0=sigma_0, Omega_g=0.)
+        w0=w0, wa=wa, Omega_g=0.,
+        mg_parametrization=MuSigmaMG(mu_0=mu_0, sigma_0=sigma_0))
 
     # Calculate distance using pyccl
     a = 1. / (1. + z)
@@ -493,7 +496,8 @@ def compare_distances_hiz_muSig(z, chi_bench, Omega_v, w0, wa):
     cosmo = ccl.Cosmology(
         Omega_c=Omega_c, Omega_b=Omega_b, Neff=Neff,
         h=h, A_s=A_s, n_s=n_s, Omega_k=Omega_k,
-        w0=w0, wa=wa, mu_0=mu_0, sigma_0=sigma_0, Omega_g=0.)
+        w0=w0, wa=wa, Omega_g=0.,
+        mg_parametrization=MuSigmaMG(mu_0=mu_0, sigma_0=sigma_0))
 
     # Calculate distance using pyccl
     a = 1. / (1. + z)

benchmarks/test_growth.py

@@ -2,6 +2,8 @@
 import pytest
 
 import pyccl as ccl
+from pyccl.modified_gravity import MuSigmaMG
+
 
 GROWTH_HIZ_TOLERANCE = 6.0e-6
 GROWTH_TOLERANCE = 1e-4
@@ -87,7 +89,8 @@ def compare_growth(
     cosmo = ccl.Cosmology(
         Omega_c=Omega_c, Omega_b=Omega_b, Neff=Neff, m_nu=m_nu,
         h=h, A_s=A_s, n_s=n_s, Omega_k=Omega_k, Omega_g=0,
-        w0=w0, wa=wa, mu_0=mu_0, sigma_0=sigma_0)
+        w0=w0, wa=wa,
+        mg_parametrization=MuSigmaMG(mu_0=mu_0, sigma_0=sigma_0))
 
     # Calculate distance using pyccl
     a = 1. / (1. + z)

benchmarks/test_halomod_numbercounts.py

@@ -15,7 +15,6 @@ def test_hmcalculator_number_counts_numcosmo():
         w0=-1.0,
         wa=0.0,
         T_CMB=2.7245,
-        mu_0=0.0,
         transfer_function='eisenstein_hu',
         matter_power_spectrum='linear'
     )
@@ -74,7 +73,6 @@ def test_hmcalculator_number_counts_numcosmo_highacc():
         w0=-1.0,
         wa=0.0,
         T_CMB=2.7245,
-        mu_0=0.0,
         transfer_function='eisenstein_hu',
         matter_power_spectrum='linear'
     )

benchmarks/test_power_mg.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 import pyccl as ccl
+from pyccl.modified_gravity import MuSigmaMG
+
 
 POWER_MG_TOL = 1e-2
 
@@ -17,9 +19,10 @@ def test_power_mg(model):
         A_s=2.1e-9,
         n_s=0.96,
         Neff=3.046,
-        mu_0=mu_0[model],
-        sigma_0=sigma_0[model],
         Omega_k=0,
+        mg_parametrization=MuSigmaMG(
+            mu_0=mu_0[model],
+            sigma_0=sigma_0[model]),
         matter_power_spectrum='linear',
         transfer_function='boltzmann_class')
 

benchmarks/test_power_mg2.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 import pyccl as ccl
+from pyccl.modified_gravity import MuSigmaMG
+
 
 POWER_MG_TOL = 1e-4
 
@@ -17,11 +19,12 @@ def test_power_mg(model):
         A_s=2.1e-9,
         n_s=0.96,
         Neff=3.046,
-        mu_0=mu_0[model],
-        sigma_0=sigma_0[model],
         Omega_k=0,
         m_nu=0,
         T_CMB=2.7255,
+        mg_parametrization=MuSigmaMG(
+            mu_0=mu_0[model],
+            sigma_0=sigma_0[model]),
         matter_power_spectrum='linear',
         transfer_function='boltzmann_isitgr')
 

pyccl/boltzmann.py

@@ -273,11 +273,11 @@ def get_isitgr_pk_lin(cosmo):
     cp.omk = cosmo['Omega_k']
     cp.GR = 1  # means GR modified!
     cp.ISiTGR_muSigma = True
-    cp.mu0 = cosmo['mu_0']
-    cp.Sigma0 = cosmo['sigma_0']
-    cp.c1 = cosmo['c1_mg']
-    cp.c2 = cosmo['c2_mg']
-    cp.Lambda = cosmo['lambda_mg']
+    cp.mu0 = cosmo.mg_parametrization.mu_0
+    cp.Sigma0 = cosmo.mg_parametrization.sigma_0
+    cp.c1 = cosmo.mg_parametrization.c1_mg
+    cp.c2 = cosmo.mg_parametrization.c2_mg
+    cp.Lambda = cosmo.mg_parametrization.lambda_mg
 
     # "constants"
     cp.TCMB = cosmo['T_CMB']

pyccl/cosmology.py

@@ -17,7 +17,7 @@
 from . import (
     CCLError, CCLObject, CCLParameters, CosmologyParams,
     DEFAULT_POWER_SPECTRUM, DefaultParams, Pk2D, check, lib,
-    unlock_instance, emulators, baryons)
+    unlock_instance, emulators, baryons, modified_gravity)
 from . import physical_constants as const
 
 
@@ -148,35 +148,18 @@ class Cosmology(CCLObject):
             of state. Defaults to 0.
         T_CMB (:obj:`float`): The CMB temperature today. The default of
             is 2.725.
-        mu_0 (:obj:`float`): One of the parameters of the mu-Sigma
-            modified gravity model. Defaults to 0.0
-        sigma_0 (:obj:`float`): One of the parameters of the mu-Sigma
-            modified gravity model. Defaults to 0.0
-        c1_mg (:obj:`float`): MG parameter that enters in the scale
-            dependence of mu affecting its large scale behavior. Default to 1.
-            See, e.g., Eqs. (46) in Ade et al. 2015, arXiv:1502.01590
-            where their f1 and f2 functions are set equal to the commonly used
-            ratio of dark energy density parameter at scale factor a over
-            the dark energy density parameter today
-        c2_mg (:obj:`float`): MG parameter that enters in the scale
-            dependence of Sigma affecting its large scale behavior. Default 1.
-            See, e.g., Eqs. (47) in Ade et al. 2015, arXiv:1502.01590
-            where their f1 and f2 functions are set equal to the commonly used
-            ratio of dark energy density parameter at scale factor a over
-            the dark energy density parameter today
-        lambda_mg (:obj:`float`): MG parameter that sets the start
-            of dependance on c1 and c2 MG parameters. Defaults to 0.0
-            See, e.g., Eqs. (46) & (47) in Ade et al. 2015, arXiv:1502.01590
-            where their f1 and f2 functions are set equal to the commonly used
-            ratio of dark energy density parameter at scale factor a over
-            the dark energy density parameter today
         transfer_function (:obj:`str` or :class:`~pyccl.emulators.emu_base.EmulatorPk`):
             The transfer function to use. Defaults to 'boltzmann_camb'.
         matter_power_spectrum (:obj:`str` or :class:`~pyccl.emulators.emu_base.EmulatorPk`):
             The matter power spectrum to use. Defaults to 'halofit'.
         baryonic_effects (:class:`~pyccl.baryons.baryons_base.Baryons` or `None`):
             The baryonic effects model to use. Options are `None` (no baryonic effects), or
             a :class:`~pyccl.baryons.baryons_base.Baryons` object.
+        mg_parametrization (:class:`~pyccl.modified_gravity.modified_gravity_base.ModifiedGravity`
+            or `None`):
+            The modified gravity parametrization to use. Options are `None` (no MG), or
+            a :class:`~pyccl.modified_gravity.modified_gravity_base.ModifiedGravity` object. 
+            Currently, only :class:`~pyccl.modified_gravity.MuSigmaMG` is supported.
         extra_parameters (:obj:`dict`): Dictionary holding extra
             parameters. Currently supports extra parameters for CAMB.
             Details described below. Defaults to None.
@@ -206,17 +189,17 @@ class Cosmology(CCLObject):
     from ._core.repr_ import build_string_Cosmology as __repr__
     __eq_attrs__ = ("_params_init_kwargs", "_config_init_kwargs",
                     "_accuracy_params", "lin_pk_emu", 'nl_pk_emu',
-                    "baryons",)
+                    "baryons", "mg_parametrization")
 
     def __init__(
             self, *, Omega_c=None, Omega_b=None, h=None, n_s=None,
             sigma8=None, A_s=None, Omega_k=0., Omega_g=None,
             Neff=None, m_nu=0., mass_split='normal', w0=-1., wa=0.,
             T_CMB=DefaultParams.T_CMB,
-            mu_0=0, sigma_0=0, c1_mg=1, c2_mg=1, lambda_mg=0,
             transfer_function='boltzmann_camb',
             matter_power_spectrum='halofit',
             baryonic_effects=None,
+            mg_parametrization=None,
             extra_parameters=None,
             T_ncdm=DefaultParams.T_ncdm):
 
@@ -243,13 +226,29 @@ def __init__(
                 raise ValueError("`baryonic_effects` must be `None` "
                                  "or a `Baryons` instance.")
 
+        self.mg_parametrization = mg_parametrization
+        if self.mg_parametrization is not None and not isinstance(
+                self.mg_parametrization,
+                modified_gravity.ModifiedGravity):
+            raise ValueError("`mg_parametrization` must be `None` "
+                             "or a `ModifiedGravity` instance.")
+
+        if self.mg_parametrization is None:
+            # Internally, CCL still relies exclusively on the mu-Sigma
+            # parametrization, so we fill that in for now unless something
+            # else is provided.
+            self.mg_parametrization = modified_gravity.MuSigmaMG()
+        if not isinstance(
+                self.mg_parametrization,
+                modified_gravity.MuSigmaMG):
+            raise NotImplementedError("`mg_parametrization` only supports the "
+                                      "mu-Sigma parametrization at this point")
+
         # going to save these for later
         self._params_init_kwargs = dict(
             Omega_c=Omega_c, Omega_b=Omega_b, h=h, n_s=n_s, sigma8=sigma8,
             A_s=A_s, Omega_k=Omega_k, Omega_g=Omega_g, Neff=Neff, m_nu=m_nu,
             mass_split=mass_split, w0=w0, wa=wa, T_CMB=T_CMB, T_ncdm=T_ncdm,
-            mu_0=mu_0, sigma_0=sigma_0,
-            c1_mg=c1_mg, c2_mg=c2_mg, lambda_mg=lambda_mg,
             extra_parameters=extra_parameters)
 
         self._config_init_kwargs = dict(
@@ -349,7 +348,6 @@ def _build_parameters(
             self, Omega_c=None, Omega_b=None, h=None, n_s=None, sigma8=None,
             A_s=None, Omega_k=None, Neff=None, m_nu=None, mass_split=None,
             w0=None, wa=None, T_CMB=None, T_ncdm=None,
-            mu_0=None, sigma_0=None, c1_mg=None, c2_mg=None, lambda_mg=None,
             Omega_g=None, extra_parameters=None):
         """Build a ccl_parameters struct"""
         # Fill-in defaults (SWIG converts `numpy.nan` to `NAN`)
@@ -413,6 +411,15 @@ def _build_parameters(
             # Omega_g was passed - modify Omega_l
             Omega_l += rho_g/rho_crit - Omega_g
 
+        # Take the mu-Sigma parameters from the modified_gravity container
+        # object. This is the only supported MG parametrization at this time.
+        assert isinstance(self.mg_parametrization, modified_gravity.MuSigmaMG)
+        mu_0 = self.mg_parametrization.mu_0
+        sigma_0 = self.mg_parametrization.sigma_0
+        c1_mg = self.mg_parametrization.c1_mg
+        c2_mg = self.mg_parametrization.c2_mg
+        lambda_mg = self.mg_parametrization.lambda_mg
+
         self._fill_params(
             m_nu=nu_mass, sum_nu_masses=sum(nu_mass), N_nu_mass=N_nu_mass,
             N_nu_rel=N_nu_rel, Neff=Neff, Omega_nu_mass=Omega_nu_mass,
@@ -505,7 +512,7 @@ def _compute_linear_power(self):
             # For MG, the input sigma8 includes the effects of MG, while the
             # sigma8 that CAMB uses is the GR definition. So we need to rescale
             # sigma8 afterwards.
-            if self["mu_0"] != 0:
+            if self.mg_parametrization.mu_0 != 0:
                 rescale_s8 = True
         elif trf == 'boltzmann_class':
             pk = self.get_class_pk_lin()
@@ -528,7 +535,7 @@ def _compute_linear_power(self):
         pkl = None
         if self._config_init_kwargs["matter_power_spectrum"] == "camb":
             rescale_mg = False
-            if self["mu_0"] != 0:
+            if self.mg_parametrization.mu_0 != 0:
                 raise ValueError("Can't rescale non-linear power spectrum "
                                  "from CAMB for mu-Sigma MG.")
             name = "delta_matter:delta_matter"
@@ -743,10 +750,12 @@ class CosmologyCalculator(Cosmology):
             same as in the Cosmology base class.
         T_ncdm (:obj:`float`): Non-CDM temperature in units of photon
             temperature. The default is the same as in the base class
-        mu_0 (:obj:`float`): One of the parameters of the mu-Sigma
-            modified gravity model. Defaults to 0.0
-        sigma_0 (:obj:`float`): One of the parameters of the mu-Sigma
-            modified gravity model. Defaults to 0.0
+        mg_parametrization (:class:`~pyccl.modified_gravity.ModifiedGravity`
+            or `None`):
+            The modified gravity parametrization to use. Options are `None`
+            (no MG), or a :class:`~pyccl.modified_gravity.ModifiedGravity`
+            object. Currently, only :class:`~pyccl.modified_gravity.MuSigmaMG`
+            is supported.
         background (:obj:`dict`): a dictionary describing the background
             expansion. It must contain three mandatory entries: ``'a'``: an
             array of monotonically ascending scale-factor values. ``'chi'``:
@@ -808,14 +817,14 @@ def __init__(
             sigma8=None, A_s=None, Omega_k=0., Omega_g=None,
             Neff=None, m_nu=0., mass_split="normal", w0=-1., wa=0.,
             T_CMB=DefaultParams.T_CMB, T_ncdm=DefaultParams.T_ncdm,
-            mu_0=0., sigma_0=0., background=None, growth=None,
+            mg_parametrization=None, background=None, growth=None,
             pk_linear=None, pk_nonlin=None, nonlinear_model=None):
 
         super().__init__(
             Omega_c=Omega_c, Omega_b=Omega_b, h=h, n_s=n_s, sigma8=sigma8,
             A_s=A_s, Omega_k=Omega_k, Omega_g=Omega_g, Neff=Neff, m_nu=m_nu,
             mass_split=mass_split, w0=w0, wa=wa, T_CMB=T_CMB, T_ncdm=T_ncdm,
-            mu_0=mu_0, sigma_0=sigma_0,
+            mg_parametrization=mg_parametrization,
             transfer_function="calculator", matter_power_spectrum="calculator")
 
         self._input_arrays = {"background": background, "growth": growth,

pyccl/modified_gravity/__init__.py

@@ -0,0 +1,2 @@
+from .modified_gravity_base import *
+from .mu_Sigma import *

pyccl/modified_gravity/modified_gravity_base.py

@@ -0,0 +1,5 @@
+__all__ = ("ModifiedGravity",)
+
+
+class ModifiedGravity:
+    parametrization: str = None