Repository: skypyproject/skypy

Issue:
Subhalos in parent halos
## Description

We need a function to add dark matter subhalo to a parent halo (section three of Vale and Ostriker 2005). Possibly related to #177.

## Inputs
 - Parent halo properties (mostly mass, but possibly also redshift).

## Outputs
 - List of subhalo properties.

## References
- [Vale and Ostriker, 2005, astro-ph/0511816.](https://arxiv.org/pdf/astro-ph/0511816.pdf)

Discussion:
Lucia-Fonseca: This is ready for review

aamara: This looks really good. On the tests, it looks like the only thing that fails is codeclimate, which seems to only want four inputs. I don't know how seriously we're taking this. 

The other thought I had is that at the moment the tests only assert the lengths of the arrays. Is there a way to add a couple more asserts to make sure values are consistent. This is probably harder for the sampler, but maybe can be done for the mass function code?
skypybot: Hello @Lucia-Fonseca! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:

* In the file [`skypy/halo/mass.py`](https://github.com/skypyproject/skypy/blob/edec28c8a13e544f45dea4ca88a65111ff9f13f8/skypy/halo/mass.py):

> [Line 452:1](https://github.com/skypyproject/skypy/blob/edec28c8a13e544f45dea4ca88a65111ff9f13f8/skypy/halo/mass.py#L452): [W293](https://duckduckgo.com/?q=pep8%20W293) blank line contains whitespace




Lucia-Fonseca: Some tests need to be changed
Lucia-Fonseca: I am using `skypy.utils.random.schechter`, and I get this error when I run the tests locally on my computer: `TypeError: schechter() got an unexpected keyword argument 'scale'`.  Any idea?
rrjbca: The PR branch is not up to date with master. Click 'Resolve conflicts', resolve the merge conflict and then pull the updated branch. Or fetch master and merge the changes locally.
ntessore: Just as a heads up, it is almost certainly better to replace the calls to `gamma` and `gammaincc` by a small numerical integration, as e.g. done in #244. As it stands you need to work around integer values of `alpha`.
Lucia-Fonseca: @rrjbca Thank you for your comments.
> * gamma_M is defined as an integral from zero to m_cut, but we are sampling from m_min to m_max(=m_cut?) do we need to redefine the limits of the integral defining gamma_M? Or does e.g. the line `n_subhalos = _subhalo_amplitude(M, params) *  gammaincc(1.0 - alpha, x_low) * gamma(1.0 - alpha)` correctly account for this?

- You are right with respect to `m_cut`, that would correspond to the most massive object that you would call "subhalo". That is, the user would be sampling from m_min to m_max(=m_cut). 
- n_subhalos is the integral of equation 3 from m_min to infinity. Perhaps you are correct and the upper limit should be m_cut.

> * Reading sections 3.1 and 3.2 it appears that m_cut is not an independent variable, rather it is determined by the mass loss of the system. Can we understand and implement this?
Indeed you are correct here too. I quote the paper "[m_cut(M)] This value will be set by assuming that the maximum original mass is half the total mass of the parent, M/2".
If I understand this correctly, `m_max` or `m_cut` should not be given by the user but fixed to be `halo_mass / 2`.


Lucia-Fonseca: > Just as a heads up, it is almost certainly better to replace the calls to `gamma` and `gammaincc` by a small numerical integration, as e.g. done in #244. As it stands you need to work around integer values of `alpha`.

Agreed
rrjbca: >   Indeed you are correct here too. I quote the paper "[m_cut(M)] This value will be set by assuming that the maximum original mass is half the total mass of the parent, M/2".
>   If I understand this correctly, `m_max` or `m_cut` should not be given by the user but fixed to be `halo_mass / 2`.

Not quite, it says the maximum "original" mass of the subhalo is M/2, which I think is different from the "present day" or "evolved" mass. The original mass and present day mass are related by the "mass loss factor". I need to read section 3.1 more carefully to understand how we incorporate mass loss.
Lucia-Fonseca: Tests still fail due to `test_quenching.py` failing.
Lucia-Fonseca: All checks passed!

Patch:
diff --git a/skypy/halo/mass.py b/skypy/halo/mass.py
index d83d010..1a2b4cf 100644
--- a/skypy/halo/mass.py
+++ b/skypy/halo/mass.py
@@ -7,30 +7,39 @@ Models
    :nosignatures:
    :toctree: ../api/
 
-   press_schechter
+   ellipsoidal_collapse_function
    halo_mass_function
    halo_mass_sampler
-   ellipsoidal_collapse_function
-   press_schechter_collapse_function
-   sheth_tormen_collapse_function
+   number_subhalos
    press_schechter
+   press_schechter_collapse_function
+   press_schechter_mass_function
    sheth_tormen
+   sheth_tormen_collapse_function
+   sheth_tormen_mass_function
+   subhalo_mass_sampler
 '''
 
 import numpy as np
 from scipy import integrate
+from scipy.special import gamma
+from skypy.utils.special import gammaincc
 from functools import partial
 from astropy import units
+from skypy.utils.random import schechter
 
 __all__ = [
+    'ellipsoidal_collapse_function',
     'halo_mass_function',
     'halo_mass_sampler',
-    'sheth_tormen_collapse_function',
+    'number_subhalos',
+    'press_schechter',
     'press_schechter_collapse_function',
-    'sheth_tormen_mass_function',
     'press_schechter_mass_function',
     'sheth_tormen',
-    'press_schechter',
+    'sheth_tormen_collapse_function',
+    'sheth_tormen_mass_function',
+    'subhalo_mass_sampler',
  ]
 
 
@@ -152,6 +161,7 @@ def halo_mass_sampler(m_min, m_max, resolution, wavenumber, power_spectrum,
     size: int, optional
         Output shape of samples. Default is None.
 
+
     Returns
     --------
     sample: (size,) array_like
@@ -312,3 +322,124 @@ def _sigma_squared(M, k, Pk, growth_function, cosmology):
 def _dlns_dlnM(sigma, M):
     ds = np.gradient(sigma, M)
     return np.absolute((M / sigma) * ds)
+
+
+def number_subhalos(halo_mass, alpha, beta, gamma_M, x, m_min, noise=True):
+    r'''Number of subhalos.
+    This function calculates the number of subhalos above a given initial mass
+    for a parent halo of given mass according to the model of Vale & Ostriker
+    2004 [1]_ equation (7). The number of subhalos returned can optionally be
+    sampled from a Poisson distribution with that mean.
+
+    Parameters
+    -----------
+    halo_mass : (nm, ) array_like
+        The mass of the halo parent, in units of solar mass.
+    alpha, beta : float
+        Parameters that determines the subhalo Schechter function. Its the amplitude
+        is defined by equation (2) in [1].
+    gamma_M : float
+        Present day mass fraction in subhalos.
+    x : float
+        Parameter that accounts for the added mass of the original, unstripped
+        subhalos.
+    m_min : array_like
+        Original mass of the least massive subhalo, in units of solar mass.
+        Current stripped mass is given by :math:`x m_{min}`.
+    noise : bool, optional
+        Poisson-sample the number of subhalos. Default is `True`.
+
+    Returns
+    --------
+    nsubhalos: array_like
+        Array of the number of subhalos assigned to parent halos with mass halo_mass.
+
+    Examples
+    ---------
+    >>> import numpy as np
+    >>> from skypy.halo import mass
+
+    This gives the number of subhalos in a parent halo of mass :math:`10^{12} M_\odot`
+
+    >>> halo, min_sh = 1.0e12, 1.0e6
+    >>> alpha, beta, gamma_M = 1.9, 1.0, 0.3
+    >>> x = 3
+    >>> nsh = mass.number_subhalos(halo, alpha, beta, gamma_M, x, min_sh)
+
+    References
+    ----------
+    .. [1] Vale, A. and Ostriker, J.P. (2005), arXiv: astro-ph/0402500.
+    '''
+    # m_min is the minimum original subhalo mass
+    x_low = m_min / (x * beta * halo_mass)
+    # Subhalo amplitude from equation [2]
+    A = gamma_M / (beta * gamma(2.0 - alpha))
+    # The mean number of subhalos above a mass threshold
+    # can be obtained by integrating equation (1) in [1]
+    n_subhalos = A * gammaincc(1.0 - alpha, x_low) * gamma(1.0 - alpha)
+
+    # Random number of subhalos following a Poisson distribution
+    # with mean n_subhalos
+    return np.random.poisson(n_subhalos) if noise else n_subhalos
+
+
+def subhalo_mass_sampler(halo_mass, nsubhalos, alpha, beta,
+                         x, m_min, resolution=100):
+    r'''Subhalo mass sampler.
+    This function samples the original, unstriped masses of subhaloes from the
+    subhalo mass function of their parent halo with a constant mass stripping
+    factor given by equation (1) in [1]_ and the upper subhalo mass limit from
+    [2]_.
+
+    Parameters
+    -----------
+    halo_mass : (nm, ) array_like
+        The mass of the halo parent, in units of solar mass.
+    nsubhalos: (nm, ) array_like
+        Array of the number of subhalos assigned to parent halos with mass `halo_mass`.
+    alpha, beta : float
+        Parameters that determines the subhalo Schechter function. Its the amplitude
+        is defined by equation 2 in [1].
+    x : float
+        Parameter that accounts for the added mass of the original, unstripped
+        subhalos.
+    m_min : float
+        Original mass of the least massive subhalo, in units of solar mass.
+        Current stripped mass is given by :math:`m_{min} / x`.
+    resolution: int, optional
+        Resolution of the inverse transform sampling spline. Default is 100.
+
+    Returns
+    --------
+    sample: (nh, ) array_like
+        List of original masses drawn from the subhalo mass function for each
+        parent halo, in units of solar mass. The length corresponds to the total
+        number of subhalos for all parent halos, i.e. `np.sum(nsubhalos)`.
+
+    Examples
+    ---------
+    >>> import numpy as np
+    >>> from skypy.halo import mass
+
+    This example samples 100 subhalos for a parent halo of mass 1.0E12 Msun:
+
+    >>> halo, min_sh = 1.0e12, 1.0e6
+    >>> alpha, beta, gamma_M = 1.9, 1.0, 0.3
+    >>> x = 3
+    >>> nsh = 100
+    >>> sh = mass.subhalo_mass_sampler(halo, nsh, alpha, beta, x, min_sh)
+
+    References
+    ----------
+    .. [1] Vale, A. and Ostriker, J.P. (2004), arXiv: astro-ph/0402500.
+    .. [2] Vale, A. and Ostriker, J.P. (2004), arXiv: astro-ph/0511816.
+    '''
+    halo_mass = np.atleast_1d(halo_mass)
+    nsubhalos = np.atleast_1d(nsubhalos)
+    subhalo_list = []
+    for M, n in zip(halo_mass, nsubhalos):
+        x_min = m_min / (x * beta * M)
+        x_max = 0.5 / (x * beta)
+        subhalo_mass = schechter(alpha, x_min, x_max, resolution, size=n, scale=x * beta * M)
+        subhalo_list.append(subhalo_mass)
+    return np.concatenate(subhalo_list)


Test patch:
diff --git a/skypy/halo/tests/test_mass.py b/skypy/halo/tests/test_mass.py
index 7fda6a8..0eed83d 100644
--- a/skypy/halo/tests/test_mass.py
+++ b/skypy/halo/tests/test_mass.py
@@ -103,3 +103,47 @@ def test_ellipsoidal_collapse_function():
     sigma = np.sqrt(mass._sigma_squared(m_array, k, Pk, 1.0, cosmo))
     fps = mass.press_schechter_collapse_function(sigma)
     assert allclose(fps, PS_fsigma)
+
+
+def test_number_subhalos():
+    # Test analytic solution for the mean number of subhalos
+    halo_mass = 1.0e12
+    shm_min = halo_mass / 100
+    alpha, beta, gamma_M, x = 0.0, 0.39, 0.18, 3.0
+    nsh_output = mass.number_subhalos(halo_mass, alpha, beta, gamma_M, x, shm_min, noise=False)
+    nsh_mean = (gamma_M / beta) * np.exp(- shm_min / (x * beta * halo_mass))
+
+    assert round(nsh_output) == round(nsh_mean)
+
+    # Test for array input of halo parents
+    halo_parents = np.array([1.0e12, 1.0e14])
+    shm_min = 1.0e6
+    alpha, beta, gamma_M, x = 1.9, 1.0, 0.3, 1.0
+    array_nsh = mass.number_subhalos(halo_parents, alpha, beta, gamma_M, x, shm_min)
+
+    assert len(array_nsh) == len(halo_parents)
+
+
+def test_subhalo_mass_sampler():
+    # Test the output shape is correct given the sample size
+    halo_mass, shm_min = 1.0e12, 1.0e6
+    alpha, beta, x = 1.9, 1.0, 1.0
+    nsh = 20
+    array_output = mass.subhalo_mass_sampler(halo_mass, nsh, alpha, beta, x, shm_min, 100)
+
+    assert len(array_output) == np.sum(nsh)
+
+    # For each halo test that each subhalo satisfy shm_min < m < shm_max
+    shm_max = 0.5 * halo_mass
+
+    assert np.all(array_output > shm_min) and np.all(array_output < shm_max)
+
+    # Repeat the tests for arrays of halos
+    halo_mass = np.array([1.0e12, 1.0e14])
+    nsh = np.array([10, 100])
+    shm_max = 0.5 * halo_mass
+    array_output = mass.subhalo_mass_sampler(halo_mass, nsh, alpha, beta, x, shm_min, 100)
+
+    assert len(array_output) == np.sum(nsh)
+    assert np.all(array_output[:10] > shm_min) and np.all(array_output[:10] < 0.5 * halo_mass[0])
+    assert np.all(array_output[10:] > shm_min) and np.all(array_output[10:] < 0.5 * halo_mass[1])