ENH: Halos from COLOSSUS joint mass-redshift distributions

docs/halos/index.rst

@@ -2,8 +2,6 @@
 Dark Matter Halos (`skypy.halos`)
 *********************************
 
-.. automodule:: skypy.halos
-
 You can reproduce figure 2 in Sheth and Tormen 1999
 and plot the collapse functions for different halo models. For this, you can use
 a python script, for example.
@@ -91,6 +89,12 @@ file and run the pipeline, for example.
     plt.show()
 
 
+Halos (`skypy.halos`)
+=====================
+
+.. automodule:: skypy.halos
+
+
 Abundance Matching (`skypy.halos.abundance_matching`)
 =====================================================
 

skypy/halos/__init__.py

@@ -1,8 +1,22 @@
-"""
+"""Halos module.
+
 This module contains methods that model the properties of dark matter halo
 populations.
+
+Models
+======
+.. autosummary::
+   :nosignatures:
+   :toctree: ../api/
+
+   colossus_mf
 """
 
+__all__ = [
+    'colossus_mf',
+]
+
 from . import abundance_matching  # noqa F401,F403
 from . import mass  # noqa F401,F403
 from . import quenching  # noqa F401,F403
+from ._colossus import colossus_mf  # noqa F401,F403

skypy/halos/_colossus.py

@@ -5,8 +5,11 @@
 
 """
 
+from astropy.cosmology import z_at_value
+from astropy import units
 import numpy as np
 from scipy import integrate
+from skypy.galaxies.redshift import redshifts_from_comoving_density
 
 __all__ = [
     'colossus_mass_sampler',
@@ -72,3 +75,131 @@ def colossus_mass_sampler(redshift, model, mdef, m_min, m_max,
     n_uniform = np.random.uniform(size=size)
     masssample = np.interp(n_uniform, CDF, m)
     return masssample
+
+
+@units.quantity_input(sky_area=units.sr)
+def colossus_mf_redshift(redshift, model, mdef, m_min, m_max, sky_area,
+                         cosmology, sigma8, ns, resolution=1000, noise=True):
+    r'''Sample redshifts from a COLOSSUS halo mass function.
+
+    Sample the redshifts of dark matter halos following a mass function
+    implemented in COLOSSUS [1]_ within given mass and redshift ranges and for
+    a given area of the sky.
+
+    Parameters
+    ----------
+    redshift : array_like
+        Input redshift grid on which the mass function is evaluated. Halos are
+        sampled over this redshift range.
+    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 the halo mass in units of Solar mass, Msun.
+    sky_area : `~astropy.units.Quantity`
+        Sky area over which halos are sampled. Must be in units of solid angle.
+    cosmology : Cosmology
+        Cosmology object to convert comoving density.
+    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.
+    noise : bool, optional
+        Poisson-sample the number of halos. Default is `True`.
+
+    Returns
+    -------
+    redshifts : array_like
+        Redshifts of the halo sample.
+
+    References
+    ----------
+    .. [1] Diemer B., 2018, ApJS, 239, 35
+
+    '''
+    from colossus.cosmology.cosmology import fromAstropy
+    from colossus.lss.mass_function import massFunction
+
+    # Set the cosmology in COLOSSUS
+    fromAstropy(cosmology, sigma8, ns)
+
+    # Integrate the mass function to get the number density of halos at each redshift
+    def dndlnM(lnm, z):
+        return massFunction(np.exp(lnm), z, 'M', 'dndlnM', mdef, model)
+    lnmmin = np.log(m_min*cosmology.h)
+    lnmmax = np.log(m_max*cosmology.h)
+    density = [integrate.quad(dndlnM, lnmmin, lnmmax, args=(z))[0] for z in redshift]
+    density = np.array(density) * np.power(cosmology.h, 3)
+
+    # Sample halo redshifts and assign to bins
+    return redshifts_from_comoving_density(redshift, density, sky_area, cosmology, noise)
+
+
+@units.quantity_input(sky_area=units.sr)
+def colossus_mf(redshift, model, mdef, m_min, m_max, sky_area, cosmology,
+                sigma8, ns, resolution=1000, noise=True):
+    r'''Sample halo redshifts and masses from a COLOSSUS mass function.
+
+    Sample the redshifts and masses of dark matter halos following a mass
+    function implemented in COLOSSUS [1]_ within given mass and redshift ranges
+    and for a given area of the sky.
+
+    Parameters
+    ----------
+    redshift : array_like
+        Defines the edges of a set of redshift bins for which the mass function
+        is evaluated. Must be a monotonically-increasing one-dimensional array
+        of values. Halo redshifts are sampled between the minimum and maximum
+        values in this array.
+    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 the halo mass in units of Solar mass, Msun.
+    sky_area : `~astropy.units.Quantity`
+        Sky area over which halos are sampled. Must be in units of solid angle.
+    cosmology : Cosmology
+        Cosmology object to calculate comoving densities.
+    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.
+    noise : bool, optional
+        Poisson-sample the number of halos. Default is `True`.
+
+    Returns
+    -------
+    redshift, mass : tuple of array_like
+        Redshifts and masses of the halo sample.
+
+    References
+    ----------
+    .. [1] Diemer B., 2018, ApJS, 239, 35
+
+    '''
+
+    # Sample halo redshifts and assign to bins
+    z = colossus_mf_redshift(redshift, model, mdef, m_min, m_max, sky_area,
+                             cosmology, sigma8, ns, resolution, noise)
+    redshift_bin = np.digitize(z, redshift)
+
+    # Calculate the redshift at the centre of each bin
+    comoving_distance = cosmology.comoving_distance(redshift)
+    d_mid = 0.5 * (comoving_distance[:-1] + comoving_distance[1:])
+    z_mid = [z_at_value(cosmology.comoving_distance, d) for d in d_mid]
+
+    # Sample halo masses in each redshift bin
+    m = np.empty_like(z)
+    for i, zm in enumerate(z_mid):
+        mask = redshift_bin == i + 1
+        size = np.count_nonzero(mask)
+        m[mask] = colossus_mass_sampler(zm, model, mdef, m_min, m_max,
+                                        cosmology, sigma8, ns, size, resolution)
+
+    return z, m

skypy/halos/tests/test_colossus.py

@@ -32,3 +32,76 @@ def test_colossus_mass_sampler():
     D, p = kstest(sample, lambda t: np.interp(t, m, CDF))
 
     assert p > 0.01, 'D = {}, p = {}'.format(D, p)
+
+
+@pytest.mark.skipif(not HAS_COLOSSUS, reason='test requires colossus')
+@pytest.mark.flaky
+def test_colossus_mf_redshift():
+
+    from skypy.halos._colossus import colossus_mf_redshift
+    from astropy.cosmology import Planck15
+    from astropy import units
+    from scipy import integrate
+    from colossus.lss.mass_function import massFunction
+
+    # Parameters
+    redshift = np.linspace(0., 2., 100)
+    model, mdef = 'despali16', '500c'
+    m_min, m_max = 1e+12, 1e+16
+    sky_area = 1.0 * units.deg**2
+    cosmology = Planck15
+    sigma8, ns = 0.82, 0.96
+
+    # Sample redshifts
+    z_halo = colossus_mf_redshift(redshift, model, mdef, m_min, m_max, sky_area,
+                                  cosmology, sigma8, ns, noise=False)
+
+    # Integrate the mass function to get the number density of halos at each redshift
+    def dndlnM(lnm, z):
+        return massFunction(np.exp(lnm), z, 'M', 'dndlnM', mdef, model)
+    lnmmin = np.log(m_min*cosmology.h)
+    lnmmax = np.log(m_max*cosmology.h)
+    density = [integrate.quad(dndlnM, lnmmin, lnmmax, args=(z))[0] for z in redshift]
+    density = np.array(density) * np.power(cosmology.h, 3)
+
+    # Integrate total number of halos for the given area of sky
+    density *= (sky_area * cosmology.differential_comoving_volume(redshift)).to_value('Mpc3')
+    n_halo = np.trapz(density, redshift, axis=-1)
+
+    # make sure noise-free sample has right size
+    assert np.isclose(len(z_halo), n_halo, atol=1.0)
+
+    # Halo redshift CDF
+    cdf = density  # same memory
+    np.cumsum((density[1:]+density[:-1])/2*np.diff(redshift), out=cdf[1:])
+    cdf[0] = 0
+    cdf /= cdf[-1]
+
+    # Check distribution of sampled halo redshifts
+    D, p = kstest(z_halo, lambda z_: np.interp(z_, redshift, cdf))
+    assert p > 0.01, 'D = {}, p = {}'.format(D, p)
+
+
+@pytest.mark.skipif(not HAS_COLOSSUS, reason='test requires colossus')
+def test_colossus_mf():
+
+    from skypy.halos import colossus_mf
+    from astropy.cosmology import Planck15
+    from astropy import units
+
+    # redshift and mass distributions are tested separately
+    # only test that output is consistent here
+
+    z = np.linspace(0., 1., 100)
+    model, mdef = 'despali16', '500c'
+    m_min, m_max = 1e+12, 1e+16
+    sky_area = 1.0 * units.deg**2
+    cosmo = Planck15
+    sigma8, ns = 0.82, 0.96
+    z_halo, m_halo = colossus_mf(z, model, mdef, m_min, m_max, sky_area, cosmo, sigma8, ns)
+
+    assert len(z_halo) == len(m_halo)
+    assert np.all(z_halo >= np.min(z))
+    assert np.all(z_halo <= np.max(z))
+    assert np.all(m_halo >= m_min)
+    assert np.all(m_halo <= m_max)