ENH number counts integrals (#787)

* WIP first pass at number counts integrals

* make sure it is 2d

* comments

* try again

* WIP working on benchmark

* WIP working on benchmark

* WIP working on benchmark

* ENH added benchmarks

* code goes here

* working benchmark kind of

* TST add a test based on comp to scipy dblquad

* ENH added changelog entry

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Python library
 - Added more developer documentation (#776).
+- Added number counts integrals for clusters (#787)
 
 ## C library
 -
@@ -23,7 +24,7 @@
 - Improved implementation of halo model quantities (n(M), b(M), c(M)) (#668, #655, #656, #657, #636).
 - Add straightforward single massive neutrino option (#730).
 - Updated MG support via mu / Sigma (scale-independent) with higher accuracy implementation (#737)
-- Added function to estimate the scale at which non-linear structure formation becomes 
+- Added function to estimate the scale at which non-linear structure formation becomes
   important for the computation of the matter power spectrum (#760).
 
 ## C library

benchmarks/data/codes/numcosmo_halo_number_counts.py

@@ -0,0 +1,105 @@
+import math
+import pyccl as ccl
+
+try:
+    import gi
+    gi.require_version('NumCosmo', '1.0')
+except Exception:
+    pass
+
+from gi.repository import NumCosmo as Nc
+from gi.repository import NumCosmoMath as Ncm
+Ncm.cfg_init()
+
+Omega_c = 0.25
+Omega_b = 0.05
+Omega_k = 0.0
+h = 0.7
+A_s = 2.1e-9
+n_s = 0.96
+Neff = 0.0
+w0 = -1.0
+wa = 0.0
+
+cosmo = Nc.HICosmo.new_from_name(
+    Nc.HICosmo, "NcHICosmoDECpl{'massnu-length':<0>}")
+cosmo.omega_x2omega_k()
+cosmo.param_set_by_name("H0", h * 100)
+cosmo.param_set_by_name("Omegak", 0.0)
+cosmo.param_set_by_name("w0", w0)
+cosmo.param_set_by_name("w1", wa)
+cosmo.param_set_by_name("Omegab", Omega_b)
+cosmo.param_set_by_name("Omegac", Omega_c)
+cosmo.param_set_by_name("ENnu", Neff)
+
+# Set Omega_K in a consistent way
+Omega_k = cosmo.Omega_k0()
+Omega_v = cosmo.E2Omega_de(0.0)
+T_CMB = cosmo.T_gamma0()
+ccl_cosmo = ccl.Cosmology(
+    Omega_c=Omega_c, Omega_b=Omega_b, Neff=Neff,
+    h=h, n_s=n_s, Omega_k=Omega_k,
+    w0=w0, wa=wa, Omega_g=0, sigma8=0.8,
+    transfer_function='eisenstein_hu',
+    T_CMB=T_CMB,
+    matter_power_spectrum='linear')
+
+hiprim = Nc.HIPrimPowerLaw.new()
+hiprim.param_set_by_name("ln10e10ASA", math.log(1.0e10 * A_s))
+hiprim.param_set_by_name("n_SA", n_s)
+cosmo.add_submodel(hiprim)
+dist = Nc.Distance.new(5.0)
+dist.prepare(cosmo)
+tf_eh = Nc.TransferFuncEH.new()
+tf_eh.props.CCL_comp = True
+psml = Nc.PowspecMLTransfer.new(tf_eh)
+psml.require_kmin(1.0e-3)
+psml.require_kmax(1.0e3)
+psml.prepare(cosmo)
+
+fact = (0.8 / psml.sigma_tophat_R(cosmo, 1.0e-7, 0.0, 8.0 / cosmo.h()))**2
+hiprim.param_set_by_name("ln10e10ASA", math.log(1.0e10 * A_s * fact))
+psml.require_kmin(1.0e-3)
+psml.require_kmax(1.0e3)
+psml.prepare(cosmo)
+
+# now compute the mass function
+psf = Ncm.PowspecFilter.new(psml, Ncm.PowspecFilterType.TOPHAT)
+psf.set_best_lnr0()
+mulf = Nc.MultiplicityFunc.new_from_name("NcMultiplicityFuncTinkerMean")
+mulf.set_Delta(200)
+mf = Nc.HaloMassFunction.new(dist, psf, mulf)
+
+psf.prepare(cosmo)
+mf.set_area_sd(200.0)
+mf.set_prec(1.0e-9)
+mf.set_eval_limits(cosmo, math.log(1e14), math.log(1e16), 0.0, 2.0)
+mf.prepare(cosmo)
+
+gf = psml.peek_gf()
+gf.prepare(cosmo)
+
+cosmo.params_log_all()
+print(cosmo.sigma8(psf))
+print(cosmo.E(1.0/0.9 - 1))
+print(psml.eval(cosmo, 1/0.9-1, 1.0))
+print(gf.eval(cosmo, 1/0.9-1))
+
+for m in [1e14, 5e14, 1e15, 5e15]:
+    print(
+        math.log10(m), mf.dn_dlnM(cosmo, math.log(m), 1/0.9-1) * math.log(10))
+
+with open("../numcosmo_cluster_counts.txt", "w") as fp:
+    fp.write("# counts  mmin  mmax  zmin  zmax\n")
+
+    nc = mf.n(cosmo, math.log(1.0e14), math.log(2.0e14), 0, 2, True)
+    fp.write(
+        "%20.12g %20.12g %20.12g %20.12g %20.12g\n" % (nc, 1e14, 2e14, 0, 2))
+
+    nc = mf.n(cosmo, math.log(2.0e14), math.log(1.0e15), 0, 2, True)
+    fp.write(
+        "%20.12g %20.12g %20.12g %20.12g %20.12g\n" % (nc, 2e14, 1e15, 0, 2))
+
+    nc = mf.n(cosmo, math.log(1.0e15), math.log(1.0e16), 0, 2, True)
+    fp.write(
+        "%20.12g %20.12g %20.12g %20.12g %20.12g\n" % (nc, 1e15, 1e16, 0, 2))

benchmarks/data/numcosmo_cluster_counts.txt

@@ -0,0 +1,4 @@
+# counts  mmin  mmax  zmin  zmax
+        3578.9272152                1e+14                2e+14                    0                    2
+       959.715661575                2e+14                1e+15                    0                    2
+       7.19881988858                1e+15                1e+16                    0                    2

benchmarks/test_halomod_numbercounts.py

@@ -0,0 +1,62 @@
+import numpy as np
+import pyccl as ccl
+
+
+def test_hmcalculator_number_counts_numcosmo():
+    cosmo = ccl.Cosmology(
+        Omega_c=0.25,
+        Omega_b=0.05,
+        h=0.7,
+        n_s=0.96,
+        sigma8=0.8,
+        Omega_k=0.0,
+        Omega_g=0,
+        Neff=0.0,
+        w0=-1.0,
+        wa=0.0,
+        T_CMB=2.7245,
+        mu_0=0.0,
+        transfer_function='eisenstein_hu',
+        matter_power_spectrum='linear'
+    )
+    mdef = ccl.halos.MassDef(200, 'matter')
+    hmf = ccl.halos.MassFuncTinker08(cosmo, mdef,
+                                     mass_def_strict=False)
+    hbf = ccl.halos.HaloBiasTinker10(cosmo, mass_def=mdef,
+                                     mass_def_strict=False)
+
+    benches = np.loadtxt("./benchmarks/data/numcosmo_cluster_counts.txt")
+
+    for i in range(benches.shape[0]):
+        bench = benches[i, :]
+        hmc = ccl.halos.HMCalculator(
+            cosmo, hmf, hbf, mdef,
+            log10M_min=np.log10(bench[1]),
+            log10M_max=np.log10(bench[2]),
+            integration_method_M='spline')
+
+        a_2 = 1.0 / (1.0 + bench[4])
+        a_1 = 1.0 / (1.0 + bench[3])
+
+        def sel(m, a):
+            m = np.atleast_1d(m)
+            a = np.atleast_1d(a)
+            val = np.zeros_like(m.reshape(-1, 1) * a.reshape(1, -1))
+            msk_a = (a > a_2) & (a < a_1)
+            msk_m = (m > bench[1]) & (m < bench[2])
+            val[msk_m, :] += 1
+            val[:, msk_a] += 1
+            msk = val == 2
+            val[~msk] = 0
+            val[msk] = 1.0
+            return val
+
+        area = 200 * (np.pi / 180)**2
+
+        nc = hmc.number_counts(cosmo, sel, amin=a_2, amax=a_1) * area
+        assert np.isfinite(nc)
+        assert not np.allclose(nc, 0)
+
+        tol = max(0.01, np.sqrt(bench[0]) / bench[0] / 10)
+        print(nc, bench[0], nc/bench[0]-1, tol)
+        assert np.allclose(nc, bench[0], atol=0, rtol=tol)

pyccl/halos/halo_model.py

@@ -8,8 +8,11 @@
 from ..power import linear_matter_power, nonlin_matter_power
 from ..background import rho_x
 from ..pyutils import _spline_integrate
+from .. import background
 import numpy as np
 
+physical_constants = lib.cvar.constants
+
 
 class HMCalculator(object):
     """ This class implements a set of methods that can be used to
@@ -137,6 +140,67 @@ def profile_norm(self, cosmo, a, prof):
         norm = 1. / self._integrate_over_mf(uk0)
         return norm
 
+    def number_counts(self, cosmo, sel, na=128, amin=None, amax=1.0):
+        """ Solves the integral:
+
+        .. math::
+            nc(sel) = \\int dM\\int da\\,\\frac{dV}{dad\\Omega}\\,n(M,a)\\,sel(M,a)
+
+        where :math:`n(M,a)` is the halo mass function, and
+        :math:`sel(M,a)` is the selection function as a function of halo mass
+        and scale factor.
+
+        Note that the selection function is normalized to integrate to unity and
+        assumed to represent the selection probaility per unit scale factor and
+        per unit mass.
+
+        Args:
+            cosmo (:class:`~pyccl.core.Cosmology`): a Cosmology object.
+            sel (callable): function of mass and scale factor that returns the
+                selection function. This function should take in floats or arrays
+                with a signature ``sel(m, a)`` and return an array with shape
+                ``(len(m), len(a))`` according to the numpy broadcasting rules.
+            na (int): number of samples in scale factor to be used in
+                the integrals. Default: 128.
+            amin (float): the minimum scale factor at which to start integrals
+                over the selection function.
+                Default: value of ``cosmo.cosmo.spline_params.A_SPLINE_MIN``
+            amax (float): the maximum scale factor at which to end integrals
+                over the selection function.
+                Default: 1.0
+
+        Returns:
+            float: the total number of clusters
+        """  # noqa
+
+        # get a values for integral
+        if amin is None:
+            amin = cosmo.cosmo.spline_params.A_SPLINE_MIN
+        a = np.linspace(amin, amax, na)
+
+        # compute the volume element
+        abs_dzda = 1 / a / a
+        dc = background.comoving_angular_distance(cosmo, a)
+        ez = background.h_over_h0(cosmo, a)
+        dh = physical_constants.CLIGHT_HMPC / cosmo['h']
+        dvdz = dh * dc**2 / ez
+        dvda = dvdz * abs_dzda
+
+        # now do m intergrals in a loop
+        mint = np.zeros_like(a)
+        for i, _a in enumerate(a):
+            self._get_ingredients(_a, cosmo, False)
+            _selm = np.atleast_2d(sel(self._mass, _a)).T
+            mint[i] = self._integrator(
+                dvda[i] * self.mf[..., :] * _selm[..., :],
+                self._lmass
+            )
+
+        # now do scale factor integral
+        mtot = self._integrator(mint, a)
+
+        return mtot
+
     def I_0_1(self, cosmo, k, a, prof):
         """ Solves the integral: