DarkEmulator HaloMassFunction (#1138)

* DarkEmulator HaloMassFunction * add benchmark for DarkEmulator halo mass function * add benchmark for DarkEmulator halo mass function * add assert when input A_s is nan * add unit test of Dark Emulator halo mass function * extrapolate option is written in the input of mass definition * extrapolate option is written in the input of mass definition * added unit test of Dark Emulator halo mass function * added the requirement to install dark emulator to .github/environment.yml and updated readthedocs * flaked * try bumping python * george on conda * try 3.10 * tighter test * Slightly degrade nishimichi MF benchmark * Update environment.yml * self-addressed comments * check precision * loosen --------- Co-authored-by: Hironao Miyatake <[email protected]> Co-authored-by: damonge <[email protected]> Co-authored-by: Chisari, N.E. (Nora) <[email protected]>
LSSTDESC · Jul 12, 2024 · aae09e9 · aae09e9
1 parent 132cfc1
commit aae09e9
Show file tree

Hide file tree

Showing 7 changed files with 204 additions and 3 deletions.
diff --git a/.github/environment.yml b/.github/environment.yml
@@ -1,7 +1,7 @@
 name: test  # default testing environment name from conda-incubator
 dependencies:
   - gfortran=12.2.0  # for isitgr
-  - python=3.8
+  - python=3.10
   - pip
   - setuptools_scm
   - cmake
@@ -10,16 +10,22 @@ dependencies:
   - swig
   - pyyaml
   - numpy
-  - scipy
+  # The below is only because the currnt version of fast-pt uses deprecated scipy functions.
+  # We should remove the <1.14 flag as soon as this is fixed.
+  - scipy<1.14
   - camb
   - fast-pt
   - pytest
   - pytest-cov
   - pytables  # for baccoemu
   - pyfftw  # for velocileptors
+  - george  # for dark_emu
   - pip:
       #- classy<3
       - isitgr<=1.0.2
       - velocileptors @ git+https://github.com/sfschen/velocileptors
       - baccoemu @ git+https://bitbucket.org/rangulo/baccoemu.git@master
       - MiraTitanHMFemulator
+      - dark_emulator==1.1.2
+      - colossus
+
diff --git a/benchmarks/data/hmf_nishimichi19.txt b/benchmarks/data/hmf_nishimichi19.txt
@@ -0,0 +1,10 @@
+# M nM(z=0) nM(z=0.5) nM(z=1)
+1.000000000000000000e+11 3.207791791342633203e-02 3.482133312565023620e-02 3.822309639517154634e-02
+3.162277660168378906e+11 1.167383294369832118e-02 1.216923001455169108e-02 1.247601394854842144e-02
+1.000000000000000000e+12 4.248354770567464196e-03 4.252857253129659239e-03 4.072169413884454969e-03
+3.162277660168379395e+12 1.547301974962943987e-03 1.481065705413910820e-03 1.316111455631085829e-03
+1.000000000000000000e+13 5.557626903970683003e-04 4.916126927948273606e-04 3.856351878777588183e-04
+3.162277660168379297e+13 1.898410647090543048e-04 1.454945823101581545e-04 9.167663603194084086e-05
+1.000000000000000000e+14 5.747023496462301149e-05 3.389934548831424228e-05 1.442650889090236593e-05
+3.162277660168379375e+14 1.326166758048492651e-05 4.795489435014995849e-06 9.867927017794248036e-07
+1.000000000000000000e+15 1.672588114601259732e-06 2.336662646897412591e-07 1.177117474730898818e-08
diff --git a/benchmarks/test_nishimichi19_hmf.py b/benchmarks/test_nishimichi19_hmf.py
@@ -0,0 +1,25 @@
+import os
+import numpy as np
+import pyccl as ccl
+
+# Set cosmology
+# based on fiducial cosmology (Planck 2015)
+cosmo = ccl.Cosmology(Omega_c=0.2650, Omega_b=0.0492, h=0.6724,
+                      A_s=2.2065e-9, n_s=0.9645, w0=-1)
+# Read data
+dirdat = os.path.join(os.path.dirname(__file__), 'data')
+
+# Redshifts
+zs = np.array([0., 0.5, 1.])
+
+
+def test_hmf_nishimichi19():
+    hmd = ccl.halos.MassDef200m
+    mf = ccl.halos.MassFuncNishimichi19(mass_def=hmd, extrapolate=True)
+    d_hmf = np.loadtxt(os.path.join(dirdat, 'hmf_nishimichi19.txt'),
+                       unpack=True)
+    m = d_hmf[0]
+    for iz, z in enumerate(zs):
+        nm_d = d_hmf[iz+1]
+        nm_h = mf(cosmo, m, 1. / (1 + z))
+        assert np.all(np.fabs(nm_h / nm_d - 1) < 1E-4)
diff --git a/pyccl/halos/hmfunc/__init__.py b/pyccl/halos/hmfunc/__init__.py
@@ -9,3 +9,4 @@
 from .tinker10 import *
 from .watson13 import *
 from .bocquet20 import *
+from .nishimichi19 import *
diff --git a/pyccl/halos/hmfunc/nishimichi19.py b/pyccl/halos/hmfunc/nishimichi19.py
@@ -0,0 +1,96 @@
+__all__ = ("MassFuncNishimichi19",)
+
+import numpy as np
+
+from . import MassFunc
+
+
+class MassFuncNishimichi19(MassFunc):
+    """Implements the mass function emulator of `Nishimichi et al. 2019
+    <https://arxiv.org/abs/1811.09504>`_.
+    documentation is
+    `here <https://dark-emulator.readthedocs.io/en/latest/index.html#>`_.
+    This parametrization is only valid for '200m' 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.
+        extrapolate (:obj:`bool`): if ``True``, extrapolation will be used
+            to calculate the the mass function for small halo masses, below
+            the range supported by the Dark Emulator. Otherwise, an error
+            is thrown when evaluating the mass function on these masses.
+    """
+    name = 'Nishimichi19'
+
+    def __init__(self, *,
+                 mass_def="200m",
+                 mass_def_strict=True,
+                 extrapolate=False):
+        super().__init__(mass_def=mass_def, mass_def_strict=mass_def_strict)
+        from dark_emulator import model_hod
+        self.hod = model_hod.darkemu_x_hod({"fft_num": 1})
+        self.extrapolate = extrapolate
+
+    def _check_mass_def_strict(self, mass_def):
+        return mass_def.name != '200m'
+
+    def __call__(self, cosmo, M, a):
+        # Set up cosmology
+        h = cosmo['h']
+        ob = cosmo['Omega_b']*h**2
+        oc = cosmo['Omega_c']*h**2
+        As = cosmo['A_s']
+        if np.isnan(As):
+            raise ValueError("A_s must be provided to use the Dark Emulator "
+                             "halo mass function.")
+        ns = cosmo['n_s']
+        w0 = cosmo['w0']
+        onu = 0.00064  # we fix this value (Nishimichi et al. 2019)
+        Ode = 1.-(ob+oc+onu)/h**2.
+        # (omega_b,omega_c,Omega_de,ln(10^10As),ns,w)
+        cparam = np.array([ob, oc, Ode, np.log(As*1E10), ns, w0])
+        self.hod.set_cosmology(cparam)
+        # calculating interpolated dndM array is (ln(M), ln(dndM))
+        self.hod._compute_dndM_spl(redshift=1./a-1.)
+
+        # 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 < 1E12
+        m_good = ~(m_hi | m_lo)
+
+        mfh = np.zeros_like(Mh)
+        # Populate low-halo masses through extrapolation if needed
+        if np.any(m_lo):
+            if self.extrapolate:
+                # Evaluate slope at low masses
+                m0 = 10**np.array([12.0, 12.1])
+                mfp = self.hod.dndM_spl(np.log(m0))
+                mfp = np.exp(mfp)
+                slope = np.log(mfp[1]/mfp[0])/np.log(m0[1]/m0[0])
+                mfh[m_lo] = mfp[0]*(Mh[m_lo]/m0[0])**slope
+            else:
+                raise RuntimeError(
+                    "Input mass range is not supported. "
+                    "The supported range is from 10^12 to 10^16 Msun/h. "
+                    "If you want to obtain mass function at "
+                    "M200m<10^12 Msun/h, use 'extrapolate=True'. "
+                    "This function will then extrapolate outside the "
+                    "supported range.")
+        # Predict in good range of masses
+        if np.any(m_good):
+            mfp = self.hod.dndM_spl(np.log(Mh[m_good]))
+            mfp = np.exp(mfp)
+            mfh[m_good] = mfp
+        # For masses above emulator range, n(M) will be set to zero
+        # Remove h-inverse and correct for dn/dM -> dn/dlog10(M)
+        # ln10 = 2.30258509299
+        mf = mfh*Mh*2.30258509299*h**3.
+
+        if np.ndim(M) == 0:
+            return mf[0]
+        return mf
diff --git a/pyccl/tests/test_hmfunc.py b/pyccl/tests/test_hmfunc.py
@@ -6,6 +6,9 @@
 COSMO = ccl.Cosmology(
     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')
+# Dark Emulator needs A_s not sigma8, so cosmological params are redifined.
+COSMO_DE = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.67,
+                         A_s=2.2e-9, n_s=0.96, w0=-1)
 HMFS = [ccl.halos.MassFuncPress74,
         ccl.halos.MassFuncSheth99,
         ccl.halos.MassFuncJenkins01,
@@ -26,8 +29,10 @@
 M500m = ccl.halos.MassDef(500, 'matter')
 MDFS = [MVIR, MVIR, MVIR, MVIR,
         MFOF, MFOF, MVIR, MFOF, MFOF, MFOF]
-# This is kind of slow to initialize, so let's do it only once
+# These are kinds of slow to initialize, so let's do it only once
 MF_emu = ccl.halos.MassFuncBocquet20(mass_def='200c')
+# Dark Emulator needs A_s not sigma8, so cosmological params are defined later.
+MF_demu = ccl.halos.MassFuncNishimichi19(mass_def='200m', extrapolate=True)
 
 
 @pytest.mark.parametrize('nM_class', HMFS)
@@ -188,3 +193,52 @@ def test_nM_bocquet20_raises():
     # Redshift out of bounds
     with pytest.raises(ValueError):
         MF_emu(cosmo, Ms, 0.3)
+
+
+def test_nM_nishimichi_smoke():
+    for m in MS:
+        n = MF_demu(COSMO_DE, m, 0.9)
+        assert np.all(np.isfinite(n))
+        assert np.shape(n) == np.shape(m)
+
+
+def test_nM_nishimichi19_compare():
+    # Check that the values are sensible (they don't depart from other
+    # parametrisations by more than ~4%
+    # Msun, under supported range(10^12-16 Msun/h)
+    Ms = np.geomspace(1.5E12, 1E15, 128)
+    mf1 = MF_demu
+    mf2 = ccl.halos.MassFuncTinker10(mass_def='200m')
+
+    nM1 = mf1(COSMO_DE, Ms, 1.0)
+    nM2 = mf2(COSMO_DE, Ms, 1.0)
+    assert np.allclose(nM1, nM2, atol=0, rtol=0.04)
+
+
+def test_nM_nishimichi19_raises():
+    Ms = np.geomspace(1.5E12, 1E15, 128)
+    # mdef raise
+    with pytest.raises(ValueError):
+        ccl.halos.MassFuncNishimichi19(mass_def=MFOF)
+
+    # contains sigma8 not A_s
+    cosmo_s = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.67,
+                            sigma8=0.8, n_s=0.96)
+    with pytest.raises(ValueError):
+        MF_demu(cosmo_s, Ms, 1.0)
+
+    # Cosmo parameters out of bounds
+    cosmo_wr = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.67,
+                             A_s=2.2e-9, n_s=2.0)
+    with pytest.raises(RuntimeError):
+        MF_demu(cosmo_wr, Ms, 1.0)
+
+    # contain unsupported range
+    # you can pass it when you set "extrapolate=True" in input of mass
+    # function definition even you use unsupported range.
+    # default is "extrapolate=False"
+    Ms = np.geomspace(1E10, 1E15, 128)
+    MF_demu_exFal = ccl.halos.MassFuncNishimichi19(mass_def=M200m,
+                                                   extrapolate=False)
+    with pytest.raises(RuntimeError):
+        MF_demu_exFal(COSMO_DE, Ms, 1.0)
diff --git a/readthedocs/source/installation.rst b/readthedocs/source/installation.rst
@@ -163,6 +163,15 @@ MiraTitan mass function emulator
 
    $ python3 -m pip install MiraTitanHMFemulator
 
+Dark Emulator
+--------------------------------
+
+`Source code <https://github.com/DarkQuestCosmology/dark_emulator_public>`__. Installation:
+
+.. code-block:: bash
+
+   $ python3 -m pip install dark_emulator==1.1.2
+
 
 .. _getting-cmake: