Repository: skypyproject/skypy
Pull request #394 (2021-01-20 11:45:54)

refactor schechter stellar mass sampling
## Description

Refactors `skypy.galaxy.stellar_mass` to be in `skypy.galaxy._schechter` and hence the sampling of `skypy.galaxy._schechter.schechter_smf` to work in the same manner as `skypy.galaxy._schechter.schechter_lf`.

## Checklist
- [x] Follow the [Contributor Guidelines](https://github.com/skypyproject/skypy/blob/master/CONTRIBUTING.md)
- [x] Write unit tests
- [x] Write documentation strings
- [x] Assign someone from your working team to review this pull request
- [x] Assign someone from the infrastructure team to review this pull request


itrharrison: Currently fails due to the dependent_argument of redshift for stellar_mass.schechter_smf_mass.

The function itself has no reason to depend on redshift, but the parameter alpha can be a function of redshift.

Is the way to fix this add the redshift argument as
`def schechter_smf_mass(redshift, m_star, alpha, m_min, m_max, size=None,
                       resolution=1000):`
which will then be unused? Or to remove the `dependent_argument` decorators above the function?

Patch:
diff --git a/docs/galaxy.rst b/docs/galaxy.rst
index c097cc6..f6d076e 100644
--- a/docs/galaxy.rst
+++ b/docs/galaxy.rst
@@ -50,6 +50,7 @@ The following models are found in the `skypy.galaxy.redshift` package.
 
    redshifts_from_comoving_density
    schechter_lf_redshift
+   schechter_smf_redshift
    smail
 
 
@@ -93,7 +94,7 @@ The following models are found in the `skypy.galaxy.stellar_mass` package.
 .. autosummary::
   :nosignatures:
 
-  schechter_smf
+  schechter_smf_mass
 
 
 Reference/API
diff --git a/skypy/galaxy/__init__.py b/skypy/galaxy/__init__.py
index 51487af..c8ea4d2 100644
--- a/skypy/galaxy/__init__.py
+++ b/skypy/galaxy/__init__.py
@@ -15,3 +15,4 @@ from . import spectrum  # noqa F401,F403
 from . import stellar_mass  # noqa F401,F403
 
 from ._schechter import schechter_lf  # noqa
+from ._schechter import schechter_smf  # noqa
diff --git a/skypy/galaxy/_schechter.py b/skypy/galaxy/_schechter.py
index 59efed9..477c543 100644
--- a/skypy/galaxy/_schechter.py
+++ b/skypy/galaxy/_schechter.py
@@ -2,12 +2,14 @@
 
 import numpy as np
 
-from .redshift import schechter_lf_redshift
+from .redshift import schechter_lf_redshift, schechter_smf_redshift
+from .stellar_mass import schechter_smf_mass
 from .luminosity import schechter_lf_magnitude
 from astropy import units
 
 __all__ = [
     'schechter_lf',
+    'schechter_smf',
 ]
 
 
@@ -73,3 +75,68 @@ def schechter_lf(redshift, M_star, phi_star, alpha, m_lim, sky_area, cosmology,
     M = schechter_lf_magnitude(z, M_star, alpha, m_lim, cosmology)
 
     return z, M
+
+
+@units.quantity_input(sky_area=units.sr)
+def schechter_smf(redshift, m_star, phi_star, alpha, m_min, m_max, sky_area, cosmology, noise=True):
+    r'''Sample redshifts and stellar masses from a Schechter mass function.
+
+    Sample the redshifts and stellar masses of galaxies following a Schechter
+    mass function with potentially redshift-dependent parameters, limited
+    by maximum and minimum masses `m_min`, `m_max`, for a sky area `sky_area`.
+
+    Parameters
+    ----------
+    redshift : array_like
+        Input redshift grid on which the Schechter function parameters are
+        evaluated. Galaxies are sampled over this redshift range.
+    m_star : array_like or function
+        Characteristic mass of the Schechter function. Can be a
+        single value, an array of values for each `redshift`, or a function of
+        redshift.
+    phi_star : array_like or function
+        Normalisation of the Schechter function. Can be a single value, an
+        array of values for each `redshift`, or a function of redshift.
+    alpha : array_like or function
+        Schechter function power law index. Can be a single value, an array of
+        values for each `redshift`, or a function of redshift.
+    m_min, m_max : float
+        Lower and upper bounds for the stellar mass.
+    sky_area : `~astropy.units.Quantity`
+        Sky area over which galaxies are sampled. Must be in units of solid angle.
+    cosmology : Cosmology
+        Cosmology object to calculate comoving densities.
+    noise : bool, optional
+        Poisson-sample the number of galaxies. Default is `True`.
+
+    Notes
+    -----
+
+    Effectively calls `~skypy.galaxy.redshift.schechter_smf_redshift` and
+    `~skypy.galaxy.stellar_mass.schechter_smf_mass` internally and returns
+    the tuple of results.
+
+    Returns
+    -------
+    redshifts, stellar masses : tuple of array_like
+        Redshifts and stellar masses of the galaxy sample described by the Schechter
+        stellar mass function.
+
+    '''
+
+    # sample halo redshifts
+    z = schechter_smf_redshift(redshift, m_star, phi_star, alpha, m_min, m_max,
+                               sky_area, cosmology, noise)
+
+    # if a function is NOT given for M_star, phi_star, alpha, interpolate to z
+    if not callable(m_star) and np.ndim(m_star) > 0:
+        m_star = np.interp(z, redshift, m_star)
+    if not callable(phi_star) and np.ndim(phi_star) > 0:
+        phi_star = np.interp(z, redshift, phi_star)
+    if not callable(alpha) and np.ndim(alpha) > 0:
+        alpha = np.interp(z, redshift, alpha)
+
+    # sample galaxy mass for redshifts
+    m = schechter_smf_mass(z, alpha, m_star, m_min, m_max)
+
+    return z, m
diff --git a/skypy/galaxy/redshift.py b/skypy/galaxy/redshift.py
index 659e756..6350f6f 100644
--- a/skypy/galaxy/redshift.py
+++ b/skypy/galaxy/redshift.py
@@ -15,6 +15,7 @@ from ..utils import broadcast_arguments, dependent_argument
 __all__ = [
     'redshifts_from_comoving_density',
     'schechter_lf_redshift',
+    'schechter_smf_redshift',
     'smail',
 ]
 
@@ -157,6 +158,77 @@ def schechter_lf_redshift(redshift, M_star, phi_star, alpha, m_lim, sky_area,
                                            sky_area=sky_area, cosmology=cosmology, noise=noise)
 
 
+@dependent_argument('m_star', 'redshift')
+@dependent_argument('phi_star', 'redshift')
+@dependent_argument('alpha', 'redshift')
+@broadcast_arguments('redshift', 'm_star', 'phi_star', 'alpha')
+@units.quantity_input(sky_area=units.sr)
+def schechter_smf_redshift(redshift, m_star, phi_star, alpha, m_min, m_max, sky_area,
+                           cosmology, noise=True):
+    r'''Sample redshifts from Schechter function.
+
+    Sample the redshifts of galaxies following a Schechter function
+    with potentially redshift-dependent parameters, limited by stellar masses
+    `m_max` and `m_min`, for a sky area `sky_area`.
+
+    Parameters
+    ----------
+    redshift : array_like
+        Input redshift grid on which the Schechter function parameters are
+        evaluated. Galaxies are sampled over this redshift range.
+    m_star : array_like or function
+        Characteristic stellar mass of the Schechter function. Can be a
+        single value, an array of values for each `redshift`, or a function of
+        redshift.
+    phi_star : array_like or function
+        Normalisation of the Schechter function. Can be a single value, an
+        array of values for each `redshift`, or a function of redshift.
+    alpha : array_like or function
+        Schechter function power law index. Can be a single value, an array of
+        values for each `redshift`, or a function of redshift.
+    m_min : float
+        Minimum stellar mass.
+    m_max : float
+        Maximum stellar mass.
+    sky_area : `~astropy.units.Quantity`
+        Sky area over which galaxies are sampled. Must be in units of solid angle.
+    cosmology : Cosmology
+        Cosmology object to convert comoving density.
+    noise : bool, optional
+        Poisson-sample the number of galaxies. Default is `True`.
+
+    Returns
+    -------
+    redshifts : array_like
+        Redshifts of the galaxy sample described by the Schechter
+        function.
+
+    '''
+
+    lnxmin = np.log(m_min)
+    lnxmin -= np.log(m_star)
+
+    lnxmax = np.log(m_max)
+    lnxmax -= np.log(m_star)
+
+    # gamma function integrand
+    def f(lnx, a):
+        return np.exp((a + 1)*lnx - np.exp(lnx)) if lnx < lnxmax.max() else 0.
+
+    # integrate gamma function for each redshift
+    gam = np.empty_like(alpha)
+
+    for i, _ in np.ndenumerate(gam):
+        gam[i], _ = scipy.integrate.quad(f, lnxmin[i], lnxmax[i], args=(alpha[i],))
+
+    # comoving number density is normalisation times upper incomplete gamma
+    density = phi_star*gam
+
+    # sample redshifts from the comoving density
+    return redshifts_from_comoving_density(redshift=redshift, density=density,
+                                           sky_area=sky_area, cosmology=cosmology, noise=noise)
+
+
 @units.quantity_input(sky_area=units.sr)
 def redshifts_from_comoving_density(redshift, density, sky_area, cosmology, noise=True):
     r'''Sample redshifts from a comoving density function.
diff --git a/skypy/galaxy/stellar_mass.py b/skypy/galaxy/stellar_mass.py
index fd30ac2..044c696 100644
--- a/skypy/galaxy/stellar_mass.py
+++ b/skypy/galaxy/stellar_mass.py
@@ -3,30 +3,34 @@
 
 import numpy as np
 
-from skypy.utils.random import schechter
+from ..utils.random import schechter
+from ..utils import broadcast_arguments, dependent_argument
 
 
 __all__ = [
-    'schechter_smf',
+    'schechter_smf_mass',
 ]
 
 
-def schechter_smf(alpha, m_star, x_min, x_max, resolution=100, size=None):
+def schechter_smf_mass(redshift, alpha, m_star, m_min, m_max, size=None,
+                       resolution=1000):
     r""" Stellar masses following the Schechter mass function [1]_.
 
     Parameters
     ----------
+    redshift : array_like
+        Galaxy redshifts for which to sample magnitudes.
     alpha : float
         The alpha parameter in the Schechter stellar mass function.
     m_star : (nm,) array-like
-        Characteristic stellar mass M_*.
+        Characteristic stellar mass m_*.
     size: int, optional
          Output shape of stellar mass samples. If size is None and m_star
          is a scalar, a single sample is returned. If size is None and
          m_star is an array, an array of samples is returned with the same
          shape as m_star.
-    x_min, x_max : float
-        Lower and upper bounds for the random variable x in units of M_*.
+    m_min, m_max : float
+        Lower and upper bounds for the stellar mass.
     resolution : int, optional
         Resolution of the inverse transform sampling spline. Default is 100.
 
@@ -48,16 +52,6 @@ def schechter_smf(alpha, m_star, x_min, x_max, resolution=100, size=None):
 
     From this pdf one can sample the stellar masses.
 
-    Examples
-    --------
-    >>> from skypy.galaxy import stellar_mass
-
-    Sample 100 stellar masses values at redshift z = 1.0 with alpha = -1.4,
-    m_star = 10**10.67, x_min = 0.0002138 and x_max = 213.8
-
-    >>> masses = stellar_mass.schechter_smf(-1.4, 10**10.67, 0.0002138,
-    ...                               213.8, size=100)
-
     References
     ----------
     .. [1] Mo, H., Van den Bosch, F., & White, S. (2010). Galaxy Formation and
@@ -66,9 +60,18 @@ def schechter_smf(alpha, m_star, x_min, x_max, resolution=100, size=None):
 
     """
 
-    if size is None and np.shape(m_star):
-        size = np.shape(m_star)
+    # only alpha scalars supported at the moment
+    if np.ndim(alpha) > 0:
+        raise NotImplementedError('only scalar alpha is supported')
+
+    if size is None and np.shape(redshift):
+        size = np.shape(redshift)
+
+    # convert m_min, m_max to units of m_star
+    x_min = m_min / m_star
+    x_max = m_max / m_star
 
-    x_sample = schechter(alpha, x_min, x_max, resolution=resolution, size=size)
+    # sample masses
+    m = schechter(alpha, x_min, x_max, resolution, size=size, scale=m_star)
 
-    return m_star * x_sample
+    return m


Test patch:
diff --git a/skypy/galaxy/tests/test_redshift.py b/skypy/galaxy/tests/test_redshift.py
index ddf5c53..563ca99 100644
--- a/skypy/galaxy/tests/test_redshift.py
+++ b/skypy/galaxy/tests/test_redshift.py
@@ -54,6 +54,52 @@ def test_schechter_lf_redshift():
     assert p > 0.01, 'D = {}, p = {}'.format(D, p)
 
 
+@pytest.mark.flaky
+def test_schechter_smf_redshift():
+
+    from skypy.galaxy.redshift import schechter_smf_redshift
+    from astropy.cosmology import FlatLambdaCDM
+    from astropy import units
+    from scipy.special import gamma, gammaincc
+
+    # fix this cosmology
+    cosmo = FlatLambdaCDM(H0=70, Om0=0.3)
+
+    # parameters for the sampling
+    z = np.linspace(1e-10, 2., 1000)
+    m_star = 10 ** 10.67
+    phi_star = 1e-3
+    alpha = -0.5
+    m_min = 10 ** 7
+    m_max = 10 ** 13
+    sky_area = 1.0 * units.deg**2
+
+    # sample redshifts
+    z_gal = schechter_smf_redshift(z, m_star, phi_star, alpha, m_min, m_max, sky_area, cosmo, noise=False)
+
+    # density with factor from upper incomplete gamma function
+    density = phi_star*gamma(alpha+1)*gammaincc(alpha+1, m_min / m_star)
+
+    # turn into galaxies/surface area
+    density *= (sky_area * cosmo.differential_comoving_volume(z)).to_value('Mpc3')
+
+    # integrate total number
+    n_gal = np.trapz(density, z, axis=-1)
+
+    # make sure noise-free sample has right size
+    assert np.isclose(len(z_gal), n_gal, atol=1.0)
+
+    # turn density into CDF
+    cdf = density  # same memory
+    np.cumsum((density[1:]+density[:-1])/2*np.diff(z), out=cdf[1:])
+    cdf[0] = 0
+    cdf /= cdf[-1]
+
+    # check distribution of sample
+    D, p = kstest(z_gal, lambda z_: np.interp(z_, z, cdf))
+    assert p > 0.01, 'D = {}, p = {}'.format(D, p)
+
+
 @pytest.mark.flaky
 def test_redshifts_from_comoving_density():
 
diff --git a/skypy/galaxy/tests/test_schechter.py b/skypy/galaxy/tests/test_schechter.py
index 09a7cd3..75c72da 100644
--- a/skypy/galaxy/tests/test_schechter.py
+++ b/skypy/galaxy/tests/test_schechter.py
@@ -35,3 +35,37 @@ def test_schechter_lf():
     # alpha array not yet supported
     with raises(NotImplementedError):
         z_gal, M_gal = schechter_lf(z, M_star, phi_star, alpha, m_lim, sky_area, cosmo)
+
+
+def test_schechter_smf():
+
+    from pytest import raises
+    from skypy.galaxy import schechter_smf
+    from astropy import units
+
+    # redshift and magnitude distributions are tested separately
+    # only test that output is consistent here
+
+    # parameters for the sampling
+    z = np.linspace(0., 1., 100)
+    m_star = 10 ** 10.67
+    phi_star = 1e-3
+    alpha = -1.5
+    m_min = 1.e7
+    m_max = 1.e13
+    sky_area = 1.0 * units.deg**2
+    cosmo = default_cosmology.get()
+
+    # sample redshifts and magnitudes
+    z_gal, m_gal = schechter_smf(z, m_star, phi_star, alpha, m_min, m_max, sky_area, cosmo)
+
+    # check length
+    assert len(z_gal) == len(m_gal)
+
+    # turn m_star, phi_star, alpha into arrays
+    z, m_star, phi_star, alpha = np.broadcast_arrays(z, m_star, phi_star, alpha)
+
+    # sample s.t. arrays need to be interpolated
+    # alpha array not yet supported
+    with raises(NotImplementedError):
+        z_gal, m_gal = schechter_smf(z, m_star, phi_star, alpha, m_min, m_max, sky_area, cosmo)
diff --git a/skypy/galaxy/tests/test_stellar_mass.py b/skypy/galaxy/tests/test_stellar_mass.py
index 8b20e0f..7322748 100644
--- a/skypy/galaxy/tests/test_stellar_mass.py
+++ b/skypy/galaxy/tests/test_stellar_mass.py
@@ -12,8 +12,8 @@ from skypy.utils import special
 def test_exponential_distribution():
     # When alpha=0, M*=1 and x_min~0 we get a truncated exponential
     q_max = 1e2
-    sample = stellar_mass.schechter_smf(0, 1, size=1000,
-                                        x_min=1e-10, x_max=q_max,
+    sample = stellar_mass.schechter_smf_mass(0, 0, 1, size=1000,
+                                        m_min=1e-10, m_max=q_max,
                                         resolution=1000)
     d, p_value = scipy.stats.kstest(sample, 'truncexpon', args=(q_max,))
     assert p_value >= 0.01
@@ -24,12 +24,12 @@ def test_stellar_masses():
     # Test that error is returned if m_star input is an array but size !=
     # None and size != m_star,size
     with pytest.raises(ValueError):
-        stellar_mass.schechter_smf(-1.4, np.array([1e10, 2e10]), 2, 1, size=3)
+        stellar_mass.schechter_smf_mass(0., -1.4, np.array([1e10, 2e10]), 1.e7, 1.e13, size=3)
 
     # Test that an array with the sme shape as m_star is returned if m_star is
     # an array and size = None
     m_star = np.array([1e10, 2e10])
-    sample = stellar_mass.schechter_smf(-1.4, m_star, 2, 1, size=None)
+    sample = stellar_mass.schechter_smf_mass(0., -1.4, m_star, 1.e7, 1.e13, size=None)
     assert m_star.shape == sample.shape
 
     # Test that sampling corresponds to sampling from the right pdf.
@@ -57,8 +57,8 @@ def test_stellar_masses():
     m_star = 10 ** 10.67
     m_min = 10 ** 7
     m_max = 10 ** 13
-    sample = stellar_mass.schechter_smf(-1.4, m_star, m_min / m_star,
-                                        m_max / m_star, size=1000,
+    sample = stellar_mass.schechter_smf_mass(0., -1.4, m_star, m_min,
+                                        m_max, size=1000,
                                         resolution=100)
     p_value = scipy.stats.kstest(sample, calc_cdf)[1]
     assert p_value >= 0.01