make utility functions private and add unit tests for them

pyccl/pkresponse.py

@@ -75,7 +75,7 @@ def Pmm_resp(
     k_use = np.exp(lk_arr)
 
     # set h-modified cosmology to take finite differencing
-    cosmo_hp, cosmo_hm = set_hmodified_cosmology(cosmo, deltah, extra_parameters)
+    cosmo_hp, cosmo_hm = _set_hmodified_cosmology(cosmo, deltah, extra_parameters)
 
     # Growth factor
     Dp = cosmo_hp.growth_factor_unnorm(a_arr)
@@ -178,7 +178,7 @@ def darkemu_Pgm_resp(
     # set h-modified cosmology to take finite differencing
     hp = h + deltah
     hm = h - deltah
-    cosmo_hp, cosmo_hm = set_hmodified_cosmology(cosmo, deltah)
+    cosmo_hp, cosmo_hm = _set_hmodified_cosmology(cosmo, deltah)
 
     # Growth factor
     Dp = cosmo_hp.growth_factor_unnorm(a_arr)
@@ -218,7 +218,7 @@ def darkemu_Pgm_resp(
             logMfor_hmf, hmf(cosmo, 10**logMfor_hmf, aa)
         )  # Mpc^-3
 
-        darkemu_set_cosmology(emu, cosmo)
+        _darkemu_set_cosmology(emu, cosmo)
         for m in range(nM):
             Pth[m] = emu.get_phm_massthreshold(k_emu, Mh[m], z) * (1 / h) ** 3
             Pbin[m] = emu.get_phm_mass(k_emu, Mh[m], z) * (1 / h) ** 3
@@ -231,7 +231,7 @@ def darkemu_Pgm_resp(
                 dndlog10m_emu(logM1) * hbf(cosmo, (10**logM1), aa), dx=dlogM1
             ) / integrate.romb(dndlog10m_emu(logM1), dx=dlogM1)
 
-        darkemu_set_cosmology(emu, cosmo_hp)
+        _darkemu_set_cosmology(emu, cosmo_hp)
         for m in range(nM):
             Pnth_hp[m] = (
                 emu.get_phm(
@@ -240,7 +240,7 @@ def darkemu_Pgm_resp(
                 * (1 / hp) ** 3
             )
 
-        darkemu_set_cosmology(emu, cosmo_hm)
+        _darkemu_set_cosmology(emu, cosmo_hm)
         for m in range(nM):
             Pnth_hm[m] = (
                 emu.get_phm(
@@ -279,7 +279,7 @@ def darkemu_Pgm_resp(
         # Eq. 19
         bgE2 = (
             integrate.romb(
-                dndlog10m_emu(logM) * Ng * b2H17(hbf(cosmo, M, aa)),
+                dndlog10m_emu(logM) * Ng * _b2H17(hbf(cosmo, M, aa)),
                 dx=dlogM,
                 axis=0,
             )
@@ -448,7 +448,7 @@ def darkemu_Pgg_resp(
         )  # Mpc^-3
 
         for m in range(nM):
-            nths[m] = mass_to_dens(dndlog10m_emu, cosmo, M[m])
+            nths[m] = _mass_to_dens(dndlog10m_emu, cosmo, M[m])
             logM1 = np.linspace(logM[m], logM[-1], 2**5 + 1)
             dlogM1 = logM[1] - logM[0]
 
@@ -457,7 +457,7 @@ def darkemu_Pgg_resp(
             ) / integrate.romb(dndlog10m_emu(logM1), dx=dlogM1)
 
         # set cosmology for dark emulator
-        darkemu_set_cosmology(emu, cosmo)
+        _darkemu_set_cosmology(emu, cosmo)
         for m in range(nM):
             for n in range(nM):
                 Pth[m, n] = (
@@ -471,13 +471,13 @@ def darkemu_Pgg_resp(
                 )
 
                 Pth_bin[m, n] = (
-                    get_phh_massthreshold_mass(
+                    _get_phh_massthreshold_mass(
                         emu, k_emu, nths[m] / (h**3), Mh[n], z
                     )
                     * (1 / h) ** 3
                 )
 
-        darkemu_set_cosmology_forAsresp(emu, cosmo, deltalnAs)
+        _darkemu_set_cosmology_forAsresp(emu, cosmo, deltalnAs)
         for m in range(nM):
             for n in range(nM):
                 Pth_Ap[m, n] = (
@@ -490,7 +490,7 @@ def darkemu_Pgg_resp(
                     * (1 / h) ** 3
                 )
 
-        darkemu_set_cosmology_forAsresp(emu, cosmo, -deltalnAs)
+        _darkemu_set_cosmology_forAsresp(emu, cosmo, -deltalnAs)
         for m in range(nM):
             for n in range(nM):
                 Pth_Am[m, n] = (
@@ -540,7 +540,7 @@ def darkemu_Pgg_resp(
         #Eq. 19
         bgE2 = (
             integrate.romb(
-                dndlog10m_emu(logM) * Ng * b2H17(b1), dx=dlogM, axis=0
+                dndlog10m_emu(logM) * Ng * _b2H17(b1), dx=dlogM, axis=0
             )
             / ng
         )
@@ -691,25 +691,6 @@ def _mass_to_dens(dndlog10m, cosmo, mass_thre):
 
     return dens
 
-
-def _dens_to_mass(dndlog10m_emu, cosmo, dens, nint=60):
-    """Convert the cumulative number density to the halo mass threshold 
-    for the current cosmological model at redshift z.
-    """
-    mlist = np.linspace(8, np.log10(10**15.8 / cosmo["h"]), nint)
-    dlist = np.log(
-        np.array(
-            [
-                mass_to_dens(dndlog10m_emu, cosmo, 10 ** mlist[i])
-                for i in range(nint)
-            ]
-        )
-    )
-    d_to_m_interp = ius(-dlist, mlist)
-
-    return 10 ** d_to_m_interp(-np.log(dens))
-
-
 def _get_phh_massthreshold_mass(emu, k_emu, dens1, Mbin, redshift):
     """Compute the halo-halo power spectrum between mass bin halo sample 
     and mass threshold halo sample specified by the corresponding cumulative number density.

pyccl/tests/test_pkresponse.py

@@ -1,6 +1,9 @@
 import numpy as np
 import pyccl as ccl
-from pyccl.pkresponse import Pmm_resp, darkemu_Pgm_resp, darkemu_Pgg_resp
+import pytest
+from pyccl.pkresponse import *
+from pyccl.pkresponse import _mass_to_dens, _get_phh_massthreshold_mass , _b2H17, _darkemu_set_cosmology, _darkemu_set_cosmology_forAsresp, _set_hmodified_cosmology
+from .test_cclobject import check_eq_repr_hash
 
 cosmo = ccl.Cosmology(
     Omega_c=0.27,
@@ -14,11 +17,18 @@
 deltah = 0.02
 deltalnAs = 0.03
 lk_arr = np.log(np.geomspace(1e-3, 1e1, 100))
+k_use = np.exp(lk_arr)
 a_arr = np.array([1.0])
 mass_def = ccl.halos.MassDef(200, "matter")
 cm = ccl.halos.ConcentrationDuffy08(mass_def=mass_def)
 prof_hod = ccl.halos.HaloProfileHOD(mass_def=mass_def, concentration=cm)
 
+# initialize dark emulator class
+emu = darkemu.de_interface.base_class()
+
+cosmo.compute_linear_power()
+pk2dlin = cosmo.get_linear_power("delta_matter:delta_matter")
+
 
 def test_Pmm_resp():
     response = Pmm_resp(cosmo, deltah=deltah, lk_arr=lk_arr, a_arr=a_arr)
@@ -39,4 +49,65 @@ def test_Pgg_resp():
         cosmo, prof_hod, deltalnAs=deltalnAs, lk_arr=lk_arr, a_arr=a_arr
     )
 
-    assert np.all(np.isfinite(response))
+    assert np.all(np.isfinite(response))
+
+# Tests for the utility functions
+def test_mass_to_dens():
+    def dndlog10m(logM):
+        return np.ones_like(logM)
+
+    mass_thre = 1e13
+    dens = _mass_to_dens(dndlog10m, cosmo, mass_thre)
+
+    assert dens > 0
+
+def test_get_phh_massthreshold_mass():
+    # set cosmology for dark emulator
+    _darkemu_set_cosmology(emu, cosmo)
+    h = cosmo["h"]
+    k_emu = k_use / h  # [h/Mpc]
+    dens1 = 1e-3
+    Mbin = 1e13
+    redshift = 0.0    
+    phh = _get_phh_massthreshold_mass(emu, k_emu, dens1, Mbin, redshift)
+    pklin = pk2dlin(k_use, 1.0, cosmo)*h**3 
+
+    valid = (k_emu > 1e-3) & (k_emu < 1e-2)
+
+    # phh is power spectrum of biased tracer
+    assert np.all(phh[valid] > pklin[valid])
+
+def test_b2H17():
+    b2 = _b2H17(0.0)
+
+    assert b2 == 0.77
+
+def test_darkemu_set_cosmology():
+    # 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(ValueError):
+        _darkemu_set_cosmology(emu, cosmo_wr)
+
+def test_darkemu_set_cosmology_forAsresp():
+    # 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(ValueError):
+        _darkemu_set_cosmology_forAsresp(emu, cosmo_wr, deltalnAs=100.0)
+
+
+def test_set_hmodified_cosmology():
+    cosmo_hp, cosmo_hm = _set_hmodified_cosmology(cosmo, deltah, extra_parameters=None)
+
+    for cosmo_test in [cosmo_hp, cosmo_hm]:
+        cosmo_dict = cosmo_test.to_dict()
+        cosmo_dict["h"] = cosmo["h"]
+        cosmo_dict["Omega_c"] = cosmo["Omega_c"]
+        cosmo_dict["Omega_b"] = cosmo["Omega_b"]
+        cosmo_dict["extra_parameters"] = cosmo["extra_parameters"]
+        cosmo_new = cosmology.Cosmology(**cosmo_dict)
+
+        # make sure the output cosmologies are exactly the same as the input one, except for the modified parameters.
+        assert check_eq_repr_hash(cosmo, cosmo_new)
+