From d0c1115ced37ba88edef05481f41e0855d498cbc Mon Sep 17 00:00:00 2001
From: Richard R <58728519+rrjbca@users.noreply.github.com>
Date: Thu, 11 Mar 2021 15:45:31 +0000
Subject: [PATCH 1/8] Implement colossus_mf_redshift and colossus_mf

---
 skypy/halos/_colossus.py | 131 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 131 insertions(+)

diff --git a/skypy/halos/_colossus.py b/skypy/halos/_colossus.py
index 7579487c7..aca59aeb2 100644
--- a/skypy/halos/_colossus.py
+++ b/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 dimensinoal 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

From 4292eb7b89b74db29dca65f5eee22b25c6562e2f Mon Sep 17 00:00:00 2001
From: Richard R <58728519+rrjbca@users.noreply.github.com>
Date: Thu, 11 Mar 2021 15:46:03 +0000
Subject: [PATCH 2/8] Unit tests for colossus_mf_redshift and colossus_mf

---
 skypy/halos/tests/test_colossus.py | 72 ++++++++++++++++++++++++++++++
 1 file changed, 72 insertions(+)

diff --git a/skypy/halos/tests/test_colossus.py b/skypy/halos/tests/test_colossus.py
index 52296de1b..9e611ffb4 100644
--- a/skypy/halos/tests/test_colossus.py
+++ b/skypy/halos/tests/test_colossus.py
@@ -32,3 +32,75 @@ 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 Planck18
+    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 = Planck18
+    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._colossus import colossus_mf
+    from astropy.cosmology import Planck18
+    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 = Planck18
+    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)

From 0886c6af3d0911e6023dd72c3dc20ba21205191e Mon Sep 17 00:00:00 2001
From: Richard R <58728519+rrjbca@users.noreply.github.com>
Date: Thu, 11 Mar 2021 16:01:52 +0000
Subject: [PATCH 3/8] Refactor halo submodules and documentation to accommodate
 new functions

---
 docs/halos/index.rst               | 14 ++++++++++++--
 skypy/halos/__init__.py            | 16 +++++++++++++++-
 skypy/halos/redshift.py            | 19 +++++++++++++++++++
 skypy/halos/tests/test_colossus.py |  4 ++--
 4 files changed, 48 insertions(+), 5 deletions(-)
 create mode 100644 skypy/halos/redshift.py

diff --git a/docs/halos/index.rst b/docs/halos/index.rst
index 95e8994f8..8d7f9ccf1 100644
--- a/docs/halos/index.rst
+++ b/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`)
 =====================================================
 
@@ -107,3 +111,9 @@ Quenching (`skypy.halos.quenching`)
 ===================================
 
 .. automodule:: skypy.halos.quenching
+
+
+Redshift (`skypy.halos.redshift`)
+=================================
+
+.. automodule:: skypy.halos.redshift
diff --git a/skypy/halos/__init__.py b/skypy/halos/__init__.py
index 3badc0592..e5ded3234 100644
--- a/skypy/halos/__init__.py
+++ b/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
diff --git a/skypy/halos/redshift.py b/skypy/halos/redshift.py
new file mode 100644
index 000000000..d44f048ed
--- /dev/null
+++ b/skypy/halos/redshift.py
@@ -0,0 +1,19 @@
+"""Halo redshifts.
+
+This module implements functions to sample from redshift distributions of dark
+matter halos.
+
+Models
+======
+.. autosummary::
+   :nosignatures:
+   :toctree: ../api/
+
+   colossus_mf_redshift
+"""
+
+from ._colossus import colossus_mf_redshift
+
+__all__ = [
+    'colossus_mf_redshift',
+]
diff --git a/skypy/halos/tests/test_colossus.py b/skypy/halos/tests/test_colossus.py
index 9e611ffb4..b42325c14 100644
--- a/skypy/halos/tests/test_colossus.py
+++ b/skypy/halos/tests/test_colossus.py
@@ -38,7 +38,7 @@ def test_colossus_mass_sampler():
 @pytest.mark.flaky
 def test_colossus_mf_redshift():
 
-    from skypy.halos._colossus import colossus_mf_redshift
+    from skypy.halos.redshift import colossus_mf_redshift
     from astropy.cosmology import Planck18
     from astropy import units
     from scipy import integrate
@@ -84,7 +84,7 @@ def dndlnM(lnm, z):
 @pytest.mark.skipif(not HAS_COLOSSUS, reason='test requires colossus')
 def test_colossus_mf():
 
-    from skypy.halos._colossus import colossus_mf
+    from skypy.halos import colossus_mf
     from astropy.cosmology import Planck18
     from astropy import units
 

From 2b737d7ee724241ac43e30d98f9e51e478d6c813 Mon Sep 17 00:00:00 2001
From: Richard R <58728519+rrjbca@users.noreply.github.com>
Date: Thu, 11 Mar 2021 16:09:43 +0000
Subject: [PATCH 4/8] flake8 whitespace

---
 skypy/halos/tests/test_colossus.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/skypy/halos/tests/test_colossus.py b/skypy/halos/tests/test_colossus.py
index b42325c14..ac12b5ab7 100644
--- a/skypy/halos/tests/test_colossus.py
+++ b/skypy/halos/tests/test_colossus.py
@@ -81,6 +81,7 @@ def dndlnM(lnm, z):
     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 bc59000c084bc2839180da5c0718ddfb97ae2a2f Mon Sep 17 00:00:00 2001
From: Richard R <58728519+rrjbca@users.noreply.github.com>
Date: Fri, 12 Mar 2021 11:33:26 +0000
Subject: [PATCH 5/8] Update skypy/halos/_colossus.py

Co-authored-by: Lucia F. de la Bella <55983939+Lucia-Fonseca@users.noreply.github.com>
---
 skypy/halos/_colossus.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/skypy/halos/_colossus.py b/skypy/halos/_colossus.py
index aca59aeb2..938ef040c 100644
--- a/skypy/halos/_colossus.py
+++ b/skypy/halos/_colossus.py
@@ -151,7 +151,7 @@ def colossus_mf(redshift, model, mdef, m_min, m_max, sky_area, cosmology,
     ----------
     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 dimensinoal array
+        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

From df9fc2707c3a9ea68a0b8d6f4079540fb0c1e020 Mon Sep 17 00:00:00 2001
From: Richard R <58728519+rrjbca@users.noreply.github.com>
Date: Fri, 12 Mar 2021 11:40:50 +0000
Subject: [PATCH 6/8] Use Planck15 for tests

---
 skypy/halos/tests/test_colossus.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/skypy/halos/tests/test_colossus.py b/skypy/halos/tests/test_colossus.py
index ac12b5ab7..305e1fc2e 100644
--- a/skypy/halos/tests/test_colossus.py
+++ b/skypy/halos/tests/test_colossus.py
@@ -39,7 +39,7 @@ def test_colossus_mass_sampler():
 def test_colossus_mf_redshift():
 
     from skypy.halos.redshift import colossus_mf_redshift
-    from astropy.cosmology import Planck18
+    from astropy.cosmology import Planck15
     from astropy import units
     from scipy import integrate
     from colossus.lss.mass_function import massFunction
@@ -49,7 +49,7 @@ def test_colossus_mf_redshift():
     model, mdef = 'despali16', '500c'
     m_min, m_max = 1e+12, 1e+16
     sky_area = 1.0 * units.deg**2
-    cosmology = Planck18
+    cosmology = Planck15
     sigma8, ns = 0.82, 0.96
 
     # Sample redshifts
@@ -86,7 +86,7 @@ def dndlnM(lnm, z):
 def test_colossus_mf():
 
     from skypy.halos import colossus_mf
-    from astropy.cosmology import Planck18
+    from astropy.cosmology import Planck15
     from astropy import units
 
     # redshift and mass distributions are tested separately
@@ -96,7 +96,7 @@ def test_colossus_mf():
     model, mdef = 'despali16', '500c'
     m_min, m_max = 1e+12, 1e+16
     sky_area = 1.0 * units.deg**2
-    cosmo = Planck18
+    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)
 

From 4d3695e3b41c2d25234d356b2ec13e32a13434d1 Mon Sep 17 00:00:00 2001
From: Richard R <58728519+rrjbca@users.noreply.github.com>
Date: Fri, 12 Mar 2021 11:52:36 +0000
Subject: [PATCH 7/8] Correct Msun/h units

---
 skypy/halos/_colossus.py           | 4 ++--
 skypy/halos/tests/test_colossus.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/skypy/halos/_colossus.py b/skypy/halos/_colossus.py
index 938ef040c..2ecd048d7 100644
--- a/skypy/halos/_colossus.py
+++ b/skypy/halos/_colossus.py
@@ -129,8 +129,8 @@ def colossus_mf_redshift(redshift, model, mdef, m_min, m_max, sky_area,
     # 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)
+    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)
 
diff --git a/skypy/halos/tests/test_colossus.py b/skypy/halos/tests/test_colossus.py
index 305e1fc2e..5ae5e60fe 100644
--- a/skypy/halos/tests/test_colossus.py
+++ b/skypy/halos/tests/test_colossus.py
@@ -59,8 +59,8 @@ def test_colossus_mf_redshift():
     # 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)
+    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)
 

From 4d7182f4d875a6dea59aa0cf472d528c0c6566e4 Mon Sep 17 00:00:00 2001
From: Richard R <58728519+rrjbca@users.noreply.github.com>
Date: Fri, 12 Mar 2021 13:52:49 +0000
Subject: [PATCH 8/8] Remove skypy.halos.redshift submodule

---
 docs/halos/index.rst               |  6 ------
 skypy/halos/redshift.py            | 19 -------------------
 skypy/halos/tests/test_colossus.py |  2 +-
 3 files changed, 1 insertion(+), 26 deletions(-)
 delete mode 100644 skypy/halos/redshift.py

diff --git a/docs/halos/index.rst b/docs/halos/index.rst
index 8d7f9ccf1..d0d82efad 100644
--- a/docs/halos/index.rst
+++ b/docs/halos/index.rst
@@ -111,9 +111,3 @@ Quenching (`skypy.halos.quenching`)
 ===================================
 
 .. automodule:: skypy.halos.quenching
-
-
-Redshift (`skypy.halos.redshift`)
-=================================
-
-.. automodule:: skypy.halos.redshift
diff --git a/skypy/halos/redshift.py b/skypy/halos/redshift.py
deleted file mode 100644
index d44f048ed..000000000
--- a/skypy/halos/redshift.py
+++ /dev/null
@@ -1,19 +0,0 @@
-"""Halo redshifts.
-
-This module implements functions to sample from redshift distributions of dark
-matter halos.
-
-Models
-======
-.. autosummary::
-   :nosignatures:
-   :toctree: ../api/
-
-   colossus_mf_redshift
-"""
-
-from ._colossus import colossus_mf_redshift
-
-__all__ = [
-    'colossus_mf_redshift',
-]
diff --git a/skypy/halos/tests/test_colossus.py b/skypy/halos/tests/test_colossus.py
index 5ae5e60fe..f78155251 100644
--- a/skypy/halos/tests/test_colossus.py
+++ b/skypy/halos/tests/test_colossus.py
@@ -38,7 +38,7 @@ def test_colossus_mass_sampler():
 @pytest.mark.flaky
 def test_colossus_mf_redshift():
 
-    from skypy.halos.redshift import colossus_mf_redshift
+    from skypy.halos._colossus import colossus_mf_redshift
     from astropy.cosmology import Planck15
     from astropy import units
     from scipy import integrate