ENH: Add a mass sampler for the colossus halo mass function (#423)

Co-authored-by: Richard R <58728519+rrjbca@users.noreply.github.com>

Author: sutieng

setup.cfg

@@ -34,9 +34,11 @@ test =
     speclite>=0.11
 all =
     speclite>=0.11
+    colossus
 docs =
     sphinx-astropy
     matplotlib
+    colossus
     speclite>=0.11
 
 [options.package_data]

skypy/halos/_colossus.py

@@ -0,0 +1,74 @@
+"""Colossus dark matter halo properties.
+
+This module contains functions that interfaces with the external code
+`Colossus <http://www.benediktdiemer.com/code/colossus/>`_.
+
+"""
+
+import numpy as np
+from scipy import integrate
+
+__all__ = [
+    'colossus_mass_sampler',
+]
+
+try:
+    import colossus  # noqa F401
+except ImportError:
+    HAS_COLOSSUS = False
+else:
+    HAS_COLOSSUS = True
+
+
+def colossus_mass_sampler(redshift, model, mdef, m_min, m_max,
+                          cosmology, sigma8, ns, size=None, resolution=1000):
+    """Colossus halo mass sampler.
+
+    This function generate a sample of halos from a mass function which
+    is available in COLOSSUS [1]_.
+
+    Parameters
+    ----------
+    redshift : float
+        The redshift values at which to sample the halo mass.
+    model : string
+        Mass function model which is available in colossus.
+    mdef : str
+        Halo mass definition for spherical overdensities used by colossus.
+    m_min, m_max : float
+        Lower and upper bounds for halo mass in units of Solar mass, Msun.
+    cosmology : astropy.cosmology.Cosmology
+        Astropy cosmology object
+    sigma8 : float
+        Cosmology parameter, amplitude of the (linear) power spectrum on the
+        scale of 8 Mpc/h.
+    ns : float
+        Cosmology parameter, spectral index of scalar perturbation power spectrum.
+    size : int, optional
+        Number of halos to sample. If size is None (default), a single value is returned.
+    resolution : int, optional
+        Resolution of the inverse transform sampling spline. Default is 1000.
+
+    Returns
+    -------
+    sample : (size,) array_like
+        Samples drawn from the mass function, in units of solar masses.
+
+    References
+    ----------
+    .. [1] Diemer B., 2018, ApJS, 239, 35
+
+    """
+    from colossus.cosmology.cosmology import fromAstropy
+    from colossus.lss import mass_function
+    fromAstropy(cosmology, sigma8=sigma8, ns=ns)
+    h0 = cosmology.h
+    m_h0 = np.logspace(np.log10(m_min*h0), np.log10(m_max*h0), resolution)  # unit: Msun/h
+    dndm = mass_function.massFunction(m_h0, redshift, mdef=mdef, model=model,
+                                      q_out='dndlnM', q_in='M')/m_h0
+    m = m_h0/h0
+    CDF = integrate.cumtrapz(dndm, (m), initial=0)
+    CDF = CDF / CDF[-1]
+    n_uniform = np.random.uniform(size=size)
+    masssample = np.interp(n_uniform, CDF, m)
+    return masssample

skypy/halos/mass.py

@@ -7,6 +7,7 @@
    :nosignatures:
    :toctree: ../api/
 
+   colossus_mass_sampler
    ellipsoidal_collapse_function
    halo_mass_function
    halo_mass_sampler
@@ -27,8 +28,10 @@
 from functools import partial
 from astropy import units
 from skypy.utils.random import schechter
+from ._colossus import colossus_mass_sampler  # noqa F401,F403
 
 __all__ = [
+    'colossus_mass_sampler',
     'ellipsoidal_collapse_function',
     'halo_mass_function',
     'halo_mass_sampler',

skypy/halos/tests/test_colossus.py

@@ -0,0 +1,34 @@
+import numpy as np
+from scipy import integrate
+from scipy.stats import kstest
+import pytest
+from skypy.halos._colossus import HAS_COLOSSUS
+
+
+@pytest.mark.skipif(not HAS_COLOSSUS, reason='test requires colossus')
+@pytest.mark.flaky
+def test_colossus_mass_sampler():
+    from astropy.cosmology import WMAP9
+    from colossus.cosmology.cosmology import fromAstropy
+    from colossus.lss import mass_function
+    from skypy.halos.mass import colossus_mass_sampler
+    m_min, m_max, size = 1e+12, 1e+16, 100
+    sample = colossus_mass_sampler(redshift=0.1, model='despali16',
+                                   mdef='500c', m_min=m_min, m_max=m_max,
+                                   cosmology=WMAP9, sigma8=0.8200, ns=0.9608,
+                                   size=size, resolution=1000)
+    assert np.all(sample >= m_min)
+    assert np.all(sample <= m_max)
+    assert np.shape(sample) == (size,)
+    fromAstropy(WMAP9, sigma8=0.8200, ns=0.9608)
+    h0 = WMAP9.h
+    m_h0 = np.logspace(np.log10(1e+12*h0), np.log10(1e+16*h0), 1000)
+    dndm = mass_function.massFunction(m_h0, 0.1, mdef='500c', model='despali16',
+                                      q_out='dndlnM', q_in='M')/m_h0
+    m = m_h0/h0
+    CDF = integrate.cumtrapz(dndm, (m), initial=0)
+    CDF = CDF / CDF[-1]
+
+    D, p = kstest(sample, lambda t: np.interp(t, m, CDF))
+
+    assert p > 0.01, 'D = {}, p = {}'.format(D, p)