Bocquet 2020 mass function emulator

.github/environment.yml

@@ -17,3 +17,4 @@ dependencies:
       - isitgr
       - velocileptors @ git+https://github.com/sfschen/velocileptors
       - baccoemu @ git+https://bitbucket.org/rangulo/baccoemu.git@master
+      - MiraTitanHMFemulator

benchmarks/data/codes/bocquet20mf.py

@@ -0,0 +1,21 @@
+import numpy as np
+import MiraTitanHMFemulator
+
+
+emu = MiraTitanHMFemulator.Emulator()
+
+hmcosmo = {'Ommh2': 0.147,
+           'Ombh2': 0.02205,
+           'Omnuh2': 0.005,
+           'n_s': 0.96,
+           'h': 0.7,
+           'sigma_8': 0.81,
+           'w_0': -1.05,
+           'w_a': 0.1}
+
+Ms = np.geomspace(1E13, 1E15, 32)
+zs = np.array([0.0, 1.0])
+mfp = emu.predict(hmcosmo, zs, Ms, get_errors=False)[0]
+np.savez("../mf_bocquet20mf.py", m=Ms, z=zs, mf=mfp,
+         **hmcosmo)
+

benchmarks/test_bocquet20_hmf.py

@@ -0,0 +1,24 @@
+import numpy as np
+import pyccl as ccl
+
+
+def test_bocquet20_mf():
+    d = np.load("./benchmarks/data/mf_bocquet20mf.py.npz")
+    m_nu = ccl.nu_masses(Omega_nu_h2=float(d['Omnuh2']),
+                         mass_split='equal')
+    h = float(d['h'])
+    cosmo = ccl.Cosmology(Omega_c=(float(d['Ommh2']) -
+                                   float(d['Ombh2']) -
+                                   float(d['Omnuh2']))/h**2,
+                          Omega_b=float(d['Ombh2'])/h**2,
+                          h=h, n_s=float(d['n_s']),
+                          sigma8=float(d['sigma_8']),
+                          m_nu=m_nu,
+                          w0=float(d['w_0']),
+                          wa=float(d['w_a']))
+    mf = ccl.halos.MassFuncBocquet20()
+    Ms = d['m']/h
+    for z, mfp in zip(d['z'], d['mf']):
+        mf_here = mf(cosmo, Ms, 1/(1+z))
+        assert np.allclose(mfp*h**3*np.log(10), mf_here,
+                           atol=0, rtol=1E-5)

pyccl/halos/hmfunc/__init__.py

@@ -8,3 +8,4 @@
 from .tinker08 import *
 from .tinker10 import *
 from .watson13 import *
+from .bocquet20 import *

pyccl/halos/hmfunc/bocquet20.py

@@ -0,0 +1,79 @@
+__all__ = ("MassFuncBocquet20",)
+
+import numpy as np
+
+from . import MassFunc
+
+
+class MassFuncBocquet20(MassFunc):
+    """Implements the mass function emulator of `Bocquet et al. 2020
+    <https://arxiv.org/abs/2003.12116>`_. This parametrization is
+    only valid for '200c' masses.
+
+    Args:
+        mass_def (:class:`~pyccl.halos.massdef.MassDef` or :obj:`str`):
+            a mass definition object, or a name string.
+        mass_def_strict (:obj:`bool`): if ``False``, consistency of the mass
+            definition will be ignored.
+    """
+    name = 'Bocquet20'
+
+    def __init__(self, *,
+                 mass_def="200c",
+                 mass_def_strict=True):
+        super().__init__(mass_def=mass_def, mass_def_strict=mass_def_strict)
+        import MiraTitanHMFemulator
+        self.emu = MiraTitanHMFemulator.Emulator()
+
+    def _check_mass_def_strict(self, mass_def):
+        return mass_def.name != '200c'
+
+    def __call__(self, cosmo, M, a):
+        # Set up cosmology
+        h = cosmo['h']
+        om = cosmo['Omega_m']*h**2
+        ob = cosmo['Omega_b']*h**2
+        sigma8 = cosmo['sigma8']
+        if np.isnan(sigma8):
+            raise ValueError("sigma8 must be provided to use the MiraTitan "
+                             "mass function emulator.")
+        onu = cosmo.omega_x(1.0, 'neutrinos_massive')*h**2
+        hmcosmo = {'Ommh2': om,
+                   'Ombh2': ob,
+                   'Omnuh2': onu,
+                   'n_s': cosmo['n_s'],
+                   'h': h,
+                   'sigma_8': sigma8,
+                   'w_0': cosmo['w0'],
+                   'w_a': cosmo['wa']}
+
+        # Filter out masses beyond emulator range
+        M_use = np.atleast_1d(M)
+        # Add h-inverse
+        Mh = M_use * h
+        m_hi = Mh > 1E16
+        m_lo = Mh < 1E13
+        m_good = ~(m_hi | m_lo)
+
+        mfh = np.zeros_like(Mh)
+        # Populate low-halo masses through extrapolation if needed
+        if np.any(m_lo):
+            # Evaluate slope at low masses
+            m0 = 10**np.array([13.0, 13.1])
+            mfp = self.emu.predict(hmcosmo, 1/a-1, m0,
+                                   get_errors=False)[0].flatten()
+            slope = np.log(mfp[1]/mfp[0])/np.log(m0[1]/m0[0])
+            mfh[m_lo] = mfp[0]*(Mh[m_lo]/m0[0])**slope
+        # Predict in good range of masses
+        if np.any(m_good):
+            mfp = self.emu.predict(hmcosmo, 1/a-1, Mh[m_good],
+                                   get_errors=False)[0].flatten()
+            mfh[m_good] = mfp
+        # For masses above emulator range, n(M) will be set to zero
+        # Remove h-inverse and correct for dn/dln(M) -> dn/dlog10(M)
+        # log10 = 2.30258509299
+        mf = mfh*h**3*2.30258509299
+
+        if np.ndim(M) == 0:
+            return mf[0]
+        return mf

pyccl/tests/test_hmfunc.py

@@ -4,7 +4,7 @@
 
 
 COSMO = ccl.Cosmology(
-    Omega_c=0.27, Omega_b=0.045, h=0.67, sigma8=0.8, n_s=0.96,
+    Omega_c=0.27, Omega_b=0.05, h=0.67, sigma8=0.8, n_s=0.96,
     transfer_function='bbks', matter_power_spectrum='linear')
 HMFS = [ccl.halos.MassFuncPress74,
         ccl.halos.MassFuncSheth99,
@@ -14,7 +14,8 @@
         ccl.halos.MassFuncTinker10,
         ccl.halos.MassFuncWatson13,
         ccl.halos.MassFuncDespali16,
-        ccl.halos.MassFuncBocquet16]
+        ccl.halos.MassFuncBocquet16,
+        ccl.halos.MassFuncBocquet20]
 MS = [1E13, [1E12, 1E15], np.array([1E12, 1E15])]
 MFOF = ccl.halos.MassDef('fof', 'matter')
 MVIR = ccl.halos.MassDef('vir', 'critical')
@@ -24,7 +25,9 @@
 M500c = ccl.halos.MassDef(500, 'critical')
 M500m = ccl.halos.MassDef(500, 'matter')
 MDFS = [MVIR, MVIR, MVIR, MVIR,
-        MFOF, MFOF, MVIR, MFOF, MFOF]
+        MFOF, MFOF, MVIR, MFOF, MFOF, MFOF]
+# This is kind of slow to initialize, so let's do it only once
+MF_emu = ccl.halos.MassFuncBocquet20(mass_def='200c')
 
 
 @pytest.mark.parametrize('nM_class', HMFS)
@@ -153,3 +156,35 @@ def test_mass_function_mass_def_strict_always_raises():
     mdef = ccl.halos.MassDef(400, "critical")
     with pytest.raises(ValueError):
         ccl.halos.MassFuncBocquet16(mass_def=mdef, mass_def_strict=False)
+
+
+def test_nM_bocquet20_compare():
+    # Check that the values are sensible (they don't depart from other
+    # parametrisations by more than ~10%
+    Ms = np.geomspace(1E14, 1E15, 128)
+    mf1 = MF_emu
+    mf2 = ccl.halos.MassFuncTinker08(mass_def='200c')
+
+    nM1 = mf1(COSMO, Ms, 1.0)
+    nM2 = mf2(COSMO, Ms, 1.0)
+    assert np.allclose(nM1, nM2, atol=0, rtol=0.1)
+
+
+def test_nM_bocquet20_raises():
+    Ms = np.geomspace(1E12, 1E17, 128)
+
+    # Need A_s
+    cosmo = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.67,
+                          A_s=2E-9, n_s=0.96)
+    with pytest.raises(ValueError):
+        MF_emu(cosmo, Ms, 1.0)
+
+    # Cosmo parameters out of bounds
+    cosmo = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.67,
+                          sigma8=0.8, n_s=2.0)
+    with pytest.raises(ValueError):
+        MF_emu(cosmo, Ms, 1.0)
+
+    # Redshift out of bounds
+    with pytest.raises(ValueError):
+        MF_emu(cosmo, Ms, 0.3)

readthedocs/api/pyccl.halos.hmfunc.bocquet20.rst

@@ -0,0 +1,7 @@
+pyccl.halos.hmfunc.bocquet20 module
+===================================
+
+.. automodule:: pyccl.halos.hmfunc.bocquet20
+   :members:
+   :undoc-members:
+   :show-inheritance:

readthedocs/api/pyccl.halos.hmfunc.rst

@@ -9,6 +9,7 @@ Submodules
 
    pyccl.halos.hmfunc.angulo12
    pyccl.halos.hmfunc.bocquet16
+   pyccl.halos.hmfunc.bocquet20
    pyccl.halos.hmfunc.despali16
    pyccl.halos.hmfunc.jenkins01
    pyccl.halos.hmfunc.press74