Classy transfer_model

CONTRIBUTORS.rst

@@ -5,7 +5,7 @@ Steven Murray: `@steven-murray <https://github.com/steven-murray>`_
 Contributors
 ------------
 * Jordan Mirocha (UCLA): `@mirochaj <https://github.com/mirochaj>`_
-* `@jlashner <https://github.com/jlashner>`_
+* Jack Lashner (USC): `@jlashner <https://github.com/jlashner>`_
 * `@liweitianux <https://github.com/liweitianux>`_
 
 Comments, corrections and suggestions

hmf/density_field/transfer_models.py

@@ -20,7 +20,14 @@
 except ImportError:
     HAVE_CAMB = False
 
-_allfits = ["CAMB", "FromFile", "EH_BAO", "EH_NoBAO", "BBKS", "BondEfs"]
+try:
+    import classy
+
+    HAVE_CLASS = True
+except ImportError:
+    HAVE_CLASS = False
+
+_allfits = ["CAMB", "FromFile", "EH_BAO", "EH_NoBAO", "BBKS", "BondEfs", "CLASS"]
 
 
 class TransferComponent(Component):
@@ -713,3 +720,76 @@ class EH(EH_BAO):
     """Alias of :class:`EH_BAO`."""
 
     pass
+
+
+class CLASS(FromFile):
+    r"""
+    Transfer function computed by CLASS.
+
+    Parameters
+    ----------
+    cosmo : :class:`astropy.cosmology.FLRW` instance
+        The cosmology used in the calculation
+
+    model_parameters : unpack-dict
+        Parameters specific to this model.
+
+        **class_params:**
+            Dict of params to set for CLASS object
+        **class_object:**
+            Custom class object. If not set, a new Class object will be
+            created.
+    """
+
+    _defaults = {"class_params": None, "class_obj": None}
+
+    def __init__(self, *args, **kwargs):
+        if not HAVE_CLASS:
+            raise ImportError(
+                "classy failed to import! Cannot use the Class " "transfer model."
+            )
+        super(CLASS, self).__init__(*args, **kwargs)
+        self.spline_fn = None
+
+        if self.params["class_obj"] is not None:
+            self.class_obj = self.params["class_obj"]
+            if not isinstance(self.class_obj, classy.Class):
+                raise TypeError("Parameter 'class_obj' must be of type classy.Class.")
+        else:
+            h = self.cosmo.h
+            class_params = {
+                "omega_b": self.cosmo.Ob0 * h ** 2,
+                "omega_cdm": (self.cosmo.Om0 - self.cosmo.Ob0) * h ** 2,
+                "Omega_k": self.cosmo.Ok0,
+                "N_ncdm": 1,  # Camb/class defaults differ here so set manually
+                "N_ur": self.cosmo.Neff - 1,
+                "m_ncdm": sum(self.cosmo.m_nu).value,
+                "T_cmb": self.cosmo.Tcmb0.value,
+                "h": h,
+            }
+            if self.params["class_params"] is not None:
+                class_params.update(self.params["class_params"])
+            class_params["output"] = "dTk"
+
+            self.class_obj = classy.Class()
+            self.class_obj.set(class_params)
+
+    def lnt(self, lnk):
+        if self.spline_fn is None:
+            if self.class_obj.state == 0:
+                self.class_obj.compute()
+            trans = self.class_obj.get_transfer(output_format="camb")
+
+            _k, _t = trans["k (h/Mpc)"], trans["-T_tot/k2"]
+            if np.any(_t < 0):
+                print(
+                    "Warning! The transfer function dips below zero! "
+                    "Returning ln(|T|)."
+                )
+                _t = np.abs(_t)
+            _lnk, _lnt = np.log(_k), np.log(_t)
+
+            if lnk[0] < _lnk[0]:
+                _lnk, _lnt = self._check_low_k(_lnk, _lnt, lnk[0])
+            self.spline_fn = spline(_lnk, _lnt, k=1)
+        return self.spline_fn(lnk)

tests/test_transfer.py

@@ -52,6 +52,14 @@ def test_bondefs():
     assert np.isclose(np.exp(t._unnormalised_lnT[0]), 1, rtol=1e-5)
 
 
+def test_class_camb_defaults():
+    # Sets a small lnk_max to capture main features but avoid extrap error
+    tclass = Transfer(transfer_model="CLASS", lnk_max=0)._unnormalised_lnT
+    tcamb = Transfer(transfer_model="CAMB", lnk_max=0)._unnormalised_lnT
+    percent_err = 200 * (tclass - tcamb) / (tclass + tcamb)
+    assert np.all(percent_err < 1)
+
+
 @pytest.mark.skip("Too slow and needs to be constantly updated.")
 def test_data(datadir):
     import camb