From 541c4ea54c4511a097e797d8aa7f21f2b6429c4b Mon Sep 17 00:00:00 2001
From: Ian Harrison <itrharrison@gmail.com>
Date: Fri, 9 Oct 2020 16:30:04 +0100
Subject: [PATCH] added P&H2017 neutral hydrogen mass model

---
 skypy/neutral_hydrogen/__init__.py        | 16 +++++
 skypy/neutral_hydrogen/mass.py            | 85 +++++++++++++++++++++++
 skypy/neutral_hydrogen/tests/test_mass.py | 29 ++++++++
 3 files changed, 130 insertions(+)
 create mode 100644 skypy/neutral_hydrogen/__init__.py
 create mode 100644 skypy/neutral_hydrogen/mass.py
 create mode 100644 skypy/neutral_hydrogen/tests/test_mass.py

diff --git a/skypy/neutral_hydrogen/__init__.py b/skypy/neutral_hydrogen/__init__.py
new file mode 100644
index 000000000..5cbf75fee
--- /dev/null
+++ b/skypy/neutral_hydrogen/__init__.py
@@ -0,0 +1,16 @@
+"""
+This module contains methods that model the properties of neutral hydrogen
+emission in the Universe.
+
+Merger Rates
+================
+
+.. autosummary::
+   :nosignatures:
+   :toctree: ../api/
+
+   pr_halo_model
+
+"""
+
+from .mass import *  # noqa F401,F403
diff --git a/skypy/neutral_hydrogen/mass.py b/skypy/neutral_hydrogen/mass.py
new file mode 100644
index 000000000..30df6937b
--- /dev/null
+++ b/skypy/neutral_hydrogen/mass.py
@@ -0,0 +1,85 @@
+import numpy as np
+from astropy import units
+
+
+r"""HI mass module.
+This module provides functions to calculate neutral Hydrogen masses for haloes.
+
+"""
+
+__all__ = [
+           'pr_halo_model',
+]
+
+
+def pr_halo_model(halo_mass, circular_velocity, cosmology,
+                  alpha=0.17,
+                  beta=-0.55,
+                  v_c0=np.exp(1.57) * (units.km / units.s),
+                  v_c1=np.exp(4.39) * (units.km / units.s),
+                  Y_p=0.24):
+    r"""Model of Padmanabhan & Refregier (2017), equation 1.
+
+    Halo neutral Hydrogen (HI 21-cm) masses as a function of halo mass,
+    halo circular velocity, and cosmology.
+
+    Parameters
+    ----------
+    halo_mass : (nhalos,) `~astropy.Quantity`
+        The masses of the underlying halos in units of solar mass.
+    circular_velocity : (nhalos,) `~astropy.Quantity`
+        The circular velocity for the halos in units of [km s-1].
+    cosmology : `~astropy.cosmology.Cosmology`
+        Cosmology object providing values of omega_baryon and omega_matter.
+    alpha : float
+        Linear model constant, fit from data.
+    beta : float
+        Model exponent, fit from data.
+    v_c0 : `~astropy.Quantity`
+        Velocity for the lower exponential cutoff of the model
+        in units of [km s-1].
+    v_c1 : `~astropy.Quantity`
+        Velocity for the upper exponential cutoff of the model
+        in units of [km s-1].
+    Y_p : float
+        Cosmic Helium fraction.
+
+    Returns
+    -------
+    m_hone : (nhalos,) `~astropy.Quantity`
+        Neutral hydrogen mass contained in the halo, in units of solar mass.
+
+    References
+    ----------
+    .. Padmanabhan & Refregier 2017, MNRAS, Volume 464, Issue 4, p. 4008
+        https://arxiv.org/abs/1607.01021
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from astropy import units
+    >>> from astropy.cosmology import Planck15
+    >>> from skypy.neutral_hydrogen import mass
+
+    Sample halo masses on a grid and use a simple model for the circular velocity.
+
+    >>> m_halo = np.logspace(8,15,128) * units.Msun
+    >>> v_halo = (96.6 * (units.km / units.s)) * m_halo / (1.e11 * units.Msun)
+
+    Calculate the neutral Hydrogen mass within each halo.
+
+    >>> m_hone = pr_halo_model(m_halo, v_halo, Planck15)
+
+    """
+
+    f_Hc = 1. - Y_p * (cosmology.Ob0 / cosmology.Om0)
+
+    lower_cutoff = np.exp(-(v_c0 / circular_velocity)**3.)
+    upper_cutoff = np.exp(-(circular_velocity / v_c1)**3.)
+
+    m_pivot = (1.e11 / cosmology.h) * units.Msun
+
+    m_hone = (alpha * f_Hc * halo_mass * np.power(halo_mass / m_pivot, beta) *
+              lower_cutoff * upper_cutoff)
+
+    return m_hone
diff --git a/skypy/neutral_hydrogen/tests/test_mass.py b/skypy/neutral_hydrogen/tests/test_mass.py
new file mode 100644
index 000000000..870a37c35
--- /dev/null
+++ b/skypy/neutral_hydrogen/tests/test_mass.py
@@ -0,0 +1,29 @@
+import numpy as np
+import pytest
+
+from astropy.cosmology import FlatLambdaCDM
+from astropy import units
+
+pr_data = np.asarray([0.00000000e+00, 8.08547768e-45, 6.40736312e+09, 3.66714636e+09,
+                      0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]) * units.Msun
+
+
+def test_pr_halo_model():
+
+    from skypy.neutral_hydrogen.mass import pr_halo_model
+
+    cosmology = FlatLambdaCDM(Om0=1.0, H0=70.0, Ob0=0.045)
+
+    m_halo = np.logspace(8, 15, 8) * units.Msun
+    v_halo = (96.6 * (units.km / units.s)) * m_halo / (1.e11 * units.Msun)
+
+    m_hone = pr_halo_model(m_halo, v_halo, cosmology)
+
+    # test scalar output
+    assert np.isscalar(pr_halo_model(m_halo[0], v_halo[0], cosmology).value)
+
+    # test array output
+    assert pr_halo_model(m_halo, v_halo, cosmology).value.shape == (8,)
+
+    # test pre-computed example
+    assert units.allclose(pr_halo_model(m_halo, v_halo, cosmology), pr_data)