{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 11 21cm Sensitivity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "import sys\n",
    "sys.path.append(\"..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoreload\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.rc_file('matplotlibrc')\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from astropy.io import fits\n",
    "\n",
    "import numpy as np\n",
    "import copy\n",
    "import pickle\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "\n",
    "import darkhistory.physics as phys\n",
    "import darkhistory.spec.spectools as spectools\n",
    "from darkhistory.history.tla import get_history\n",
    "\n",
    "import darkhistory.main as main"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 21-cm Brightness Temperature $T_{21}$ and Matter Temperature $T_m$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "21-cm cosmology is one of the most important experimental methods of understanding the thermal and ionization histories of the universe. The global 21-cm brightness temperature is related to the hydrogen spin temperature through the following relation:\n",
    "\n",
    "$$ T_\\text{21} (z) \\approx x_\\text{HI}(z) \\left( \\frac{0.15}{\\Omega_m} \\right)^{1/2} \\left( \\frac{\\Omega_b h}{0.02} \\right) \\left( \\frac{1+z}{10} \\right)^{1/2} \\left[ 1 - \\frac{T_R(z)}{T_S(z)} \\right] 23 \\text{ mK}, $$\n",
    "\n",
    "where $x_\\text{HI}$ is the neutral fraction of hydrogen, $\\Omega_m$ and $\\omega_b$ are the matter and baryon energy density as a fraction of the critical density, $h$ is the Hubble parameter today in units of 100 km s$^{-1}$ Mpc$^{-1}$, $T_R$ is the background 21-cm radiation temperature and $T_S$ is the spin temperature. Under standard assumptions, before the onset of star formation heats up the gas, $T_m < T_S < T_\\text{CMB}$, where $T_m$ is the matter temperature; this implies that measuring the global 21-cm brightness temperature at a particular redshift sets an upper bound on the matter temperature at that point in time. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The expected $T_{21}$ assuming no exotic sources of injection is approximately -200 mK at $z \\sim 17$. In this notebook, we will use DarkHistory to estimate the limits that can be set on $T_m$, assuming we measure a $T_{21}$ value of -50 mK. We will also compare the difference with and without backreaction taken into account, to give users a sense of the possible differences in constraints.\n",
    "\n",
    "Before we proceed, we should note that recent measurements performed by the EDGES collaboration [[1]](#cite_EDGES) seems to suggest a $T_{21}(z = 17) \\lesssim $ -300 mK at the 99% confidence level, which would suggest that $T_m$ is actually *colder* than expected, assuming that $T_R = T_\\text{CMB}$. The introduction of new baryon-DM interactions may be able to produce such an effect, and DarkHistory can be easily modified to accommodate this if desired."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 21-cm Constraints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's compute the limit set on $T_m$ if we take $T_{21}$ = -50 mK at $z = 17$, assuming that $T_R = T_\\text{CMB}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_21 = -50\n",
    "rs = 18 # We always define redshifts as 1 + z by convention in the code. \n",
    "one_minus_TR_over_TS = (T_21/23.) / (0.15 / phys.omega_m)**0.5 / (phys.omega_baryon*phys.h/0.02) / (rs/10)**0.5\n",
    "Tm_max = 1/(1 - one_minus_TR_over_TS) * phys.TCMB(rs) / phys.kB # in K\n",
    "print('The upper bound on Tm is: {:03.1f} K'.format(Tm_max))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the user would like to save the results that this notebook produces, we include an option to make this easy.  Simply set `save_data=True` below, and set `direc=/directory/where/the/data/gets/saved/`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_data = True\n",
    "direc = #'/foo/bar'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decaying Dark Matter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first get the constraints for decays. We'll consider two channels: $\\chi \\to e^+e^-$ and $\\chi \\to \\gamma \\gamma$, with no electroweak corrections. To compute the constraints with no backreaction, we can use the tabulated $f_c(z)$ tables, and use these as inputs for [*tla.get_history()*](https://darkhistory.readthedocs.io/en/master/_autosummary/darkhistory/history/tla/darkhistory.history.tla.get_history.html) and obtain the temperature history. \n",
    "\n",
    "With backreaction, however, the spectra of particles that get produced at each redshift step must be computed, and the temperature history integrated as a function of these spectra. This is done by the function [*main.evolve()*](https://darkhistory.readthedocs.io/en/master/_autosummary/main/main.evolve.html). \n",
    "\n",
    "For decays, we can specify the primary that we are interested in, the dark matter process, the mass of the dark matter, and the decay lifetime. We also use some coarsening to speed up the calculation. \n",
    "\n",
    "We'll define a function called ``get_T_decay()`` to run either of these functions, and return the temperature at $z = 17$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_T_decay(mDM, tau, br, pri):\n",
    "    \"\"\"Gets the temperature at z = 17, decays.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    mDM : float\n",
    "        The mass of the dark matter in eV. \n",
    "    tau : float\n",
    "        The decay lifetime in s. \n",
    "    br : bool\n",
    "        Whether to use backreaction or not. \n",
    "    pri : {'elec', 'phot'}\n",
    "        The primary particles that DM decays to. \n",
    "        \n",
    "    Returns\n",
    "    --------\n",
    "    float\n",
    "        The temperature in K at z = 17. \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    if pri == 'elec':\n",
    "        pri_str = 'elec_delta'\n",
    "        cf = 12\n",
    "    if pri == 'phot':\n",
    "        pri_str = 'phot_delta'\n",
    "        cf = 4\n",
    "        \n",
    "    if br:\n",
    "    \n",
    "        result = main.evolve(\n",
    "            primary=pri_str,\n",
    "            DM_process='decay', mDM=mDM, lifetime=tau,\n",
    "            start_rs = 3000, end_rs=18,\n",
    "            coarsen_factor=cf, backreaction=True\n",
    "        )\n",
    "        \n",
    "        return result['Tm'][-1]/phys.kB\n",
    "    \n",
    "    else:\n",
    "        # get_history takes a redshift vector \n",
    "        rs_vec = np.flipud(np.arange(18, 3000., 0.1))\n",
    "        \n",
    "        result = get_history(\n",
    "            rs_vec, baseline_f = True,\n",
    "            inj_particle = pri, DM_process = 'decay',\n",
    "            mDM=mDM, lifetime = tau\n",
    "        )\n",
    "        \n",
    "        return result[-1,0]/phys.kB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now all set up to do the scan. Let's scan over $\\chi \\to \\gamma \\gamma$, for $\\log_{10} (m_\\chi/\\text{eV}) \\in [4, 12.5] $. We'll do this in a clever way to save the user time here. We'll scan over lifetimes $\\log_{10}(\\tau/\\text{s}) \\in [23, 28.5]$, and start with a guess of $10^{25}$ s. Starting from this guess, we compute the matter temperature at $z = 17$. If the temperature lies below `Tm_max`, we will try a shorter lifetime to increase the energy deposition rate, and vice versa with a higher temperature. We'll terminate the scan once we have computed $T_m(z = 17)$ over enough points to linearly interpolate for the lifetime required to produce `Tm_max` at $z = 17$. This lifetime is the *minimum* lifetime allowed by $T_{21} = -$50 mK.\n",
    "\n",
    "Let's set up a 2D array, indexed by mass and lifetime, one without backreaction, and one with. First specify the decay products, then set up the array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pri = 'phot'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if pri == 'phot':\n",
    "    log10_m_chi_arr = np.arange(4.01, 12.76, 0.25)\n",
    "else:\n",
    "    log10_m_chi_arr = np.arange(6.01, 12.76, 0.25)\n",
    "\n",
    "log10_tau_arr   = np.arange(23., 28.25, 0.25)\n",
    "\n",
    "Tm_dec_no_br = np.zeros((log10_m_chi_arr.size, log10_tau_arr.size))\n",
    "Tm_dec_br = np.zeros_like(Tm_dec_no_br)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start the scan! To do the scan quickly, we'll first run the case with no backreaction first, and then use the results as an initial guess for [*main.evolve()*](https://darkhistory.readthedocs.io/en/master/_autosummary/main/main.evolve.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoreload\n",
    "print('###### Target Temperature: {:03.1f} K'.format(Tm_max), ' ######')\n",
    "\n",
    "for i, log10mDM in enumerate(tqdm(log10_m_chi_arr)):\n",
    "    \n",
    "    print('****** log10(mDM): ', log10mDM, ' ******')\n",
    "    mDM = 10**log10mDM\n",
    "    \n",
    "    below_target = False\n",
    "    above_target = False\n",
    "    if i==0:\n",
    "        j = 8\n",
    "    \n",
    "    while (\n",
    "        not below_target or not above_target\n",
    "    ):\n",
    "\n",
    "#         print('****************************')\n",
    "        tau = 10**log10_tau_arr[j]\n",
    "        \n",
    "        Tm_dec_no_br[i,j] = get_T_decay(mDM, tau, br=False, pri=pri)\n",
    "        print('lifetime: {:03.1e}'.format(tau), ' T(z = 17) w/o backreaction: {:03.1f} K'.format(Tm_dec_no_br[i,j]))\n",
    "             \n",
    "        if Tm_dec_no_br[i,j] <= Tm_max:\n",
    "            \n",
    "            below_target = True\n",
    "            # Below target T_m means we need a shorter lifetime.\n",
    "            j -= 1\n",
    "            \n",
    "        elif Tm_dec_no_br[i,j] >= Tm_max:\n",
    "            \n",
    "            above_target = True\n",
    "            # Above target T_m means we need a longer lifetime. \n",
    "            j += 1\n",
    "\n",
    "# Save the data that was just produced\n",
    "if save_data:\n",
    "    pickle.dump(Tm_dec_no_br, open(direc+pri+'_decay_no_br.dat','wb'))\n",
    "\n",
    "print('###### Calculation Complete! ######')           "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having obtained an (incomplete) list of temperature values, we write a simple function to perform a linear interpolation at each mass point to obtain the limit on the minimum lifetime, and then compute it in `min_lifetime_no_br`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_log10_min_lifetime(log10_lifetime_arr, Tm_arr, max_Tm):\n",
    "    \"\"\"Returns the log10 minimum lifetime from temperature calculations.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    log10_lifetime_arr : ndarray\n",
    "        Abscissa of log10 lifetimes used, in increasing order.\n",
    "    Tm_arr : ndarray\n",
    "        Computed matter temperature for each entry in lifetime_arr. \n",
    "    max_Tm : float\n",
    "        The maximum Tm allowed at this point. \n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    ndarray\n",
    "        The interpolated minimum lifetime. \n",
    "    \n",
    "    \"\"\"\n",
    "    # Find the last index at which we have Tm > max_Tm in Tm_arr. \n",
    "    above = np.where((Tm_arr > 0) & (Tm_arr > max_Tm))[0][-1]\n",
    "    # Since increasing lifetime decreases Tm, the next index must\n",
    "    # lie below max_Tm.\n",
    "    below = above+1\n",
    "    \n",
    "    # Obtain the two temperature values. \n",
    "    Ts = [Tm_arr[below], Tm_arr[above]]\n",
    "    # Obtain the two log10 lifetime values. \n",
    "    log10_taus = [log10_lifetime_arr[below], log10_lifetime_arr[above]]\n",
    "    \n",
    "    # Perform a linear interpolation.\n",
    "    return (np.diff(log10_taus)/np.diff(Ts)*(max_Tm-Ts[0]) + log10_taus[0])[0]\n",
    "\n",
    "min_lifetime_no_br = np.array(\n",
    "    [get_log10_min_lifetime(log10_tau_arr, row, Tm_max) for row in Tm_dec_no_br]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's just make a simple plot to make sure we understand what we have."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "plt.loglog()\n",
    "\n",
    "plt.plot(10**log10_m_chi_arr, 10**min_lifetime_no_br)\n",
    "\n",
    "plt.title(r'$\\chi \\to \\gamma \\gamma$ Limits, $T_{21} = -50$ mK')\n",
    "plt.ylabel(r'Minimum Decay Lifetime $\\tau$ [s]')\n",
    "plt.xlabel(r'Dark Matter Mass $m_\\chi$ [eV]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's do this now with backreaction. The backreaction calculation can be initialized using the $\\tau$ value obtained without backreaction as an initial guess. That significantly narrows down the amount of scanning that needs to be done. Otherwise, it works in the same way as before, but it takes longer because the code needs to compute all of the spectra at each step. \n",
    "\n",
    "When using *main.evolve()* for the first time, the code will take some time to load all of the transfer functions required to perform this calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('###### Target Temperature: {:03.1f} K'.format(Tm_max), ' ######')\n",
    "\n",
    "for i, log10mDM in enumerate(tqdm(log10_m_chi_arr)):\n",
    "    \n",
    "    print('****** log10(mDM): ', log10mDM, ' ******')\n",
    "    mDM = 10**log10mDM\n",
    "    \n",
    "    below_target = False\n",
    "    above_target = False\n",
    "    \n",
    "    # Get the estimate from min_lifetime_no_br\n",
    "    j = np.where(log10_tau_arr > min_lifetime_no_br[i])[0][0]\n",
    "    \n",
    "    while (\n",
    "        not below_target or not above_target\n",
    "    ):\n",
    "\n",
    "#         print('****************************')\n",
    "        tau = 10**log10_tau_arr[j]\n",
    "        \n",
    "        Tm_dec_br[i,j] = get_T_decay(mDM, tau, br=True, pri=pri)\n",
    "        print('lifetime: {:03.1e}'.format(tau), ' T(z = 17) w/o backreaction: {:03.1f} K'.format(Tm_dec_br[i,j]))\n",
    "             \n",
    "        if Tm_dec_br[i,j] <= Tm_max:\n",
    "            \n",
    "            below_target = True\n",
    "            # Below target T_m means we need a shorter lifetime.\n",
    "            j -= 1\n",
    "            \n",
    "        elif Tm_dec_br[i,j] >= Tm_max:\n",
    "            \n",
    "            above_target = True\n",
    "            # Above target T_m means we need a longer lifetime. \n",
    "            j += 1\n",
    "            \n",
    "# Save the data that was just produced\n",
    "if save_data:\n",
    "    pickle.dump(Tm_dec_br, open(direc+pri+'_decay_br.dat','wb'))\n",
    "\n",
    "print('###### Calculation Complete! ######')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Annihilating Dark Matter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Annihilating dark matter sensitivity can also be evaluated in a very similar manner. The only difference here is that structure formation now boosts the annihilation rate at late times. We can load one of DarkHistory's default structure formation boost factors that can be obtained using the function [*physics.struct_boost_func()*](https://darkhistory.readthedocs.io/en/master/_autosummary/darkhistory/physics/darkhistory.physics.struct_boost_func.html). The first thing we need to do is obtain the function that returns the boost: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct_boost_einasto_subs = phys.struct_boost_func(model='einasto_subs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now define the analogous *get_T_swave* function with this structure formation prescription included with the keyword *struct_boost* passed to [*main.evolve()*](https://darkhistory.readthedocs.io/en/master/_autosummary/main/main.evolve.html):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoreload 2\n",
    "def get_T_swave(mDM, sigmav, br, pri):\n",
    "    \"\"\"Gets the temperature at z = 17, s-wave annihilation.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    mDM : float\n",
    "        The mass of the dark matter in eV. \n",
    "    sigmav : float\n",
    "        The annihilation cross-section in cm^-3 s^-1.  \n",
    "    br : bool\n",
    "        Whether to use backreaction or not. \n",
    "    pri : {'elec', 'phot'}\n",
    "        The primary particles that DM decays to. \n",
    "        \n",
    "    Returns\n",
    "    --------\n",
    "    float\n",
    "        The temperature in K at z = 17. \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    if pri == 'elec':\n",
    "        pri_str = 'elec_delta'\n",
    "        cf = 12\n",
    "    if pri == 'phot':\n",
    "        pri_str = 'phot_delta'\n",
    "        cf = 4\n",
    "        \n",
    "    if br:\n",
    "    \n",
    "        result = main.evolve(\n",
    "            primary=pri_str,\n",
    "            DM_process='swave', mDM=mDM, sigmav=sigmav,\n",
    "            start_rs = 3000, end_rs=18,\n",
    "            coarsen_factor=cf, backreaction=True,\n",
    "            struct_boost=struct_boost_einasto_subs\n",
    "        )\n",
    "        \n",
    "        return result['Tm'][-1]/phys.kB\n",
    "    \n",
    "    else:\n",
    "        # get_history takes a redshift vector \n",
    "        rs_vec = np.flipud(np.arange(18, 3000., 0.1))\n",
    "        \n",
    "        result = get_history(\n",
    "            rs_vec, baseline_f = True,\n",
    "            inj_particle = pri, DM_process = 'swave',\n",
    "            mDM=mDM, sigmav = sigmav,\n",
    "            struct_boost=None\n",
    "        )\n",
    "        \n",
    "        return result[-1,0]/phys.kB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now initialize the data! We use $\\langle \\sigma v \\rangle / m_\\chi$ as the range to scan over, since this quantity determines the energy injection rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pri = 'phot'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if pri == 'phot':\n",
    "    log10_m_chi_arr = np.arange(4, 12.75, 0.25)\n",
    "else:\n",
    "    log10_m_chi_arr = np.arange(6.01, 12.76, 0.25)\n",
    "log10_sigmav_over_m_arr   = np.arange(-29, -22, 0.25)\n",
    "\n",
    "Tm_ann_no_br = np.zeros((log10_m_chi_arr.size, log10_sigmav_over_m_arr.size))\n",
    "Tm_ann_br = np.zeros_like(Tm_ann_no_br)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And start the scan!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%autoreload\n",
    "print('###### Target Temperature: {:03.1f} K'.format(Tm_max), ' ######')\n",
    "\n",
    "for i, log10mDM in enumerate(tqdm(log10_m_chi_arr)):\n",
    "    \n",
    "    print('****** log10(mDM): ', log10mDM, ' ******')\n",
    "    mDM = 10**log10mDM\n",
    "    \n",
    "    below_target = False\n",
    "    above_target = False\n",
    "    if i == 0:\n",
    "        j = 0\n",
    "    \n",
    "    while (\n",
    "        not below_target or not above_target\n",
    "    ):\n",
    "\n",
    "        sigmav = 10**log10_sigmav_over_m_arr[j]*mDM/1e9\n",
    "        \n",
    "        Tm_ann_no_br[i,j] = get_T_swave(mDM, sigmav, br=False, pri=pri)\n",
    "        print('sigmav: {:03.1e}'.format(sigmav), ' T(z = 17) w/o backreaction: {:03.1f} K'.format(Tm_ann_no_br[i,j]))\n",
    "             \n",
    "        if Tm_ann_no_br[i,j] <= Tm_max:\n",
    "            \n",
    "            below_target = True\n",
    "            # Below target T_m means we need a larger sigmav.\n",
    "            j += 1\n",
    "            \n",
    "        elif Tm_ann_no_br[i,j] >= Tm_max:\n",
    "            \n",
    "            above_target = True\n",
    "            # Above target T_m means we need a smaller sigmav. \n",
    "            j -= 1\n",
    "\n",
    "# Save the data that was just produced\n",
    "if save_data:\n",
    "    pickle.dump(Tm_ann_no_br, open(direc+pri+'_swave_no_br.dat','wb'))\n",
    "\n",
    "print('###### Calculation Complete! ######')        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As above, we use a simple linear interpolation to determine the maximum $\\left(\\frac{\\langle \\sigma v \\rangle}{m}\\right)_\\text{max}$ over our list of dark matter masses and store the result in `max_xsec_no_br`,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_log10_max_xsec(log10_sigmav_over_m_arr, Tm_arr, max_Tm):\n",
    "    \"\"\"Returns the log10 maximum sigma v over m from temperature calculations.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    log10_lifetime_arr : ndarray\n",
    "        Abscissa of log10 sigma v over m's used, in increasing order.\n",
    "    Tm_arr : ndarray\n",
    "        Computed matter temperature for each entry in lifetime_arr. \n",
    "    max_Tm : float\n",
    "        The maximum Tm allowed at this point. \n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    ndarray\n",
    "        The interpolated minimum lifetime. \n",
    "    \n",
    "    \"\"\"\n",
    "    # Find the first index at which we have Tm < max_Tm in Tm_arr. \n",
    "    below = np.where((Tm_arr > 0) & (Tm_arr < max_Tm))[0][-1]\n",
    "    # Since increasing lifetime decreases Tm, the next index must\n",
    "    # lie below max_Tm.\n",
    "    above = below+1\n",
    "    \n",
    "    # Obtain the two temperature values. \n",
    "    Ts = [Tm_arr[below], Tm_arr[above]]\n",
    "    # Obtain the two log10 lifetime values. \n",
    "    log10_xsecs = [log10_sigmav_over_m_arr[below], log10_sigmav_over_m_arr[above]]\n",
    "    \n",
    "    # Perform a linear interpolation.\n",
    "    return (np.diff(log10_xsecs)/np.diff(Ts)*(max_Tm-Ts[0]) + log10_xsecs[0])[0]\n",
    "\n",
    "max_xsec_no_br = np.array(\n",
    "    [get_log10_max_xsec(log10_sigmav_over_m_arr, row, Tm_max) for row in Tm_ann_no_br]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and then we plot $\\left(\\frac{\\langle \\sigma v \\rangle}{m}\\right)_\\text{max}$ as a function of $m_\\chi$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "plt.loglog()\n",
    "\n",
    "plt.plot(10**log10_m_chi_arr/10**9, 10**(max_xsec_no_br))\n",
    "\n",
    "if pri == 'phot':\n",
    "    plt.axis([1e-5,2e3,3e-27,1.5e-25])\n",
    "else:\n",
    "    plt.axis([1e-3,2e3,5e-29,1.5e-25])\n",
    "\n",
    "plt.title(r'$\\chi\\chi \\to \\gamma \\gamma$ Limits, $T_{21} = -50$ mK')\n",
    "plt.ylabel(r'Maximum $\\langle \\sigma v \\rangle$ / $m_\\chi$ [cm$^3$ s$^{-1}$ GeV$^{-1}$]')\n",
    "plt.xlabel(r'Dark Matter Mass $m_\\chi$ [eV]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and finally compute $\\left(\\frac{\\langle \\sigma v \\rangle}{m}\\right)_\\text{max}$ including back-reaction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoreload\n",
    "print('###### Target Temperature: {:03.1f} K'.format(Tm_max), ' ######')\n",
    "\n",
    "for i, log10mDM in enumerate(tqdm(log10_m_chi_arr)):\n",
    "    \n",
    "    print('****** log10(mDM): ', log10mDM, ' ******')\n",
    "    mDM = 10**log10mDM\n",
    "    \n",
    "    below_target = False\n",
    "    above_target = False\n",
    "    j = np.where(log10_sigmav_over_m_arr > max_xsec_no_br[i])[0][0]\n",
    "    \n",
    "    while (\n",
    "        not below_target or not above_target\n",
    "    ):\n",
    "\n",
    "        sigmav = 10**log10_sigmav_over_m_arr[j] * mDM/1e9\n",
    "        \n",
    "        Tm_ann_br[i,j] = get_T_swave(mDM, sigmav, br=True, pri=pri)\n",
    "        print('sigmav/m: {:03.1e}'.format(sigmav/mDM), ' T(z = 17) w/ backreaction: {:03.1f} K'.format(Tm_ann_no_br[i,j]))\n",
    "             \n",
    "        if Tm_ann_br[i,j] <= Tm_max:\n",
    "            \n",
    "            below_target = True\n",
    "            # Below target T_m means we need a larger sigmav.\n",
    "            j += 1\n",
    "            \n",
    "        elif Tm_ann_br[i,j] >= Tm_max:\n",
    "            \n",
    "            above_target = True\n",
    "            # Above target T_m means we need a smaller sigmav. \n",
    "            j -= 1\n",
    "            \n",
    "# Save the data that was just produced\n",
    "if save_data:\n",
    "    pickle.dump(Tm_ann_br, open(direc+pri+'_swave_br.dat','wb'))\n",
    "\n",
    "print('###### Calculation Complete! ######')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this final section, we show the user how to plot the 21cm Sensitivity as is done in Fig. 6 of paper I.  We assume that the user has at some point run this notebook with `save_data=True` so that the data has been saved in the directory `direc=/directory/where/the/data/gets/saved/` as described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inj_type = 'decay'\n",
    "pri = 'phot'\n",
    "direc = #'/foo/bar/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "then load the raw temperature data as a function of $m_\\chi$ and $\\tau$ or $\\langle \\sigma v \\rangle$,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Tm_no_br = pickle.load(open(direc+pri+'_'+inj_type+'_no_br.dat','rb'))\n",
    "Tm_br    = pickle.load(open(direc+pri+'_'+inj_type+'_br.dat','rb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "then process the data to the $\\tau$ or $\\langle \\sigma v \\rangle$ bounds as a function of $m_\\chi$,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if inj_type == 'decay':\n",
    "    log10_tau_arr   = np.arange(23., 28.25, 0.25)\n",
    "    thresh_params_std = np.array(\n",
    "        [get_log10_min_lifetime(log10_tau_arr, row, Tm_max) for row in Tm_no_br]\n",
    "    )\n",
    "\n",
    "    thresh_params_BR = np.array(\n",
    "        [get_log10_min_lifetime(log10_tau_arr, row, Tm_max) for row in Tm_br]\n",
    "    )\n",
    "else:\n",
    "    log10_sigmav_over_m_arr = np.arange(-29, -22, 0.25)\n",
    "    thresh_params_std = np.array(\n",
    "        [get_log10_max_xsec(log10_sigmav_over_m_arr, row, Tm_max) for row in Tm_no_br]\n",
    "    )\n",
    "\n",
    "    thresh_params_BR = np.array(\n",
    "        [get_log10_max_xsec(log10_sigmav_over_m_arr, row, Tm_max) for row in Tm_br]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make an interpolation function for the points $(m_\\text{DM}, \\tau_\\text{min})$ or $(m_\\text{DM}, \\sigma_\\text{max})$ with and without backreaction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if pri == 'phot':\n",
    "    log10_m_chi_arr = np.arange(4.01, 12.76, 0.25)\n",
    "else:\n",
    "    log10_m_chi_arr = np.arange(6.01, 12.76, 0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "std_interp = interp1d(log10_m_chi_arr, thresh_params_std, kind='quadratic')\n",
    "BR_interp = interp1d(log10_m_chi_arr, thresh_params_BR, kind='quadratic')\n",
    "\n",
    "log10mDM_arr_fine = np.arange(log10_m_chi_arr[0], log10_m_chi_arr[-1], .01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "ax = plt.gca()\n",
    "ax.loglog()\n",
    "\n",
    "xmin = 10**log10mDM_arr_fine[0]/1e9\n",
    "xmax = 10**log10mDM_arr_fine[-1]/1e9\n",
    "\n",
    "if pri=='elec':\n",
    "    if inj_type == 'swave':\n",
    "        plt.title(r'\\bf{21-cm Sensitivity,} $\\chi \\chi \\to e^+ e^-$')\n",
    "        ymin = 3e-29\n",
    "        ymax = 1e-25\n",
    "    else:\n",
    "        plt.title(r'\\bf{21-cm Sensitivity,} $\\chi \\to e^+ e^-$')\n",
    "        ymin = 1e24\n",
    "        ymax = 5e28\n",
    "else:\n",
    "    if inj_type == 'swave':\n",
    "        plt.title(r'\\bf{21-cm Sensitivity,} $\\chi\\chi \\to \\gamma\\gamma$')\n",
    "        ymin = 3e-27\n",
    "        ymax = 1.5e-25\n",
    "    else:\n",
    "        plt.title(r'\\bf{21-cm Sensitivity,} $\\chi \\to \\gamma\\gamma$')\n",
    "        ymin = 2e23\n",
    "        ymax = 4e26\n",
    "    \n",
    "plt.xlabel(r'Dark Matter Mass $m_\\chi$ [GeV]')\n",
    "if inj_type == 'swave':\n",
    "    plt.ylabel(r'Maximum $\\langle \\sigma v \\rangle / m_\\chi$ [cm$^3$ s$^{-1}$ GeV$^{-1}$]')\n",
    "else:\n",
    "    plt.ylabel(r'Minimum Lifetime $\\tau \\, [\\mathrm{s}]$')\n",
    "    \n",
    "plt.axis([xmin, xmax, ymin, ymax])\n",
    "\n",
    "plt_std, = plt.plot(10**log10mDM_arr_fine/1e9, 10**std_interp(log10mDM_arr_fine), \n",
    "        linewidth=1, label='No Backreaction', color='#6a176e', linestyle='--')\n",
    "plt_BR, = plt.plot(10**log10mDM_arr_fine/1e9, 10**BR_interp(log10mDM_arr_fine), \n",
    "         linewidth=3, label='With Backreaction', color='#f37819')\n",
    "\n",
    "ax.fill_between(\n",
    "    10**log10mDM_arr_fine/1e9, \n",
    "    1e-15, \n",
    "    (10**std_interp(log10mDM_arr_fine)), \n",
    "    interpolate=True, \n",
    "    color='#6a176e',\n",
    "    alpha=0.2\n",
    ")\n",
    "\n",
    "ax.fill_between(\n",
    "    10**log10mDM_arr_fine/1e9, \n",
    "    1e-15, \n",
    "    (10**BR_interp(log10mDM_arr_fine)), \n",
    "    interpolate=True, \n",
    "    color='#f37819',\n",
    "    alpha=.1\n",
    ")\n",
    "\n",
    "if inj_type == 'swave':\n",
    "    plt.text(0.52, 0.22, r'$T_{21}$ = $-50$ mK', fontsize=20, transform=ax.transAxes) \n",
    "else:\n",
    "    plt.text(0.49, 0.89, r'$T_{21}$ = $-50$ mK', fontsize=20, transform=ax.transAxes)\n",
    "\n",
    "if inj_type == 'swave':\n",
    "    if inj_str == 'phot':\n",
    "        leg = plt.legend(handles=[plt_std, plt_BR])\n",
    "    else:\n",
    "        leg = plt.legend(handles=[plt_std, plt_BR], loc='lower right')\n",
    "else:\n",
    "    leg = plt.legend(handles=[plt_std, plt_BR], loc=(0.4, 0.72))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bibliography"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[1]<a id='cite_EDGES'></a> Judd D. Bowman, Alan E. E. Rogers, Raul A. Monsalve, Thomas J. Mozdzen, and Nivedita Mahesh, “An absorption profile centred at 78 megahertz in the sky-averaged spectrum,” Nature 555, 67–70 (2018)."
   ]
  }
 ],
 "metadata": {
  "cite2c": {
   "citations": {
    "6232408/9ZW7JWLJ": {
     "URL": "https://www.nature.com/articles/nature25792",
     "accessed": {
      "day": 1,
      "month": 4,
      "year": 2019
     },
     "id": "6232408/9ZW7JWLJ",
     "title": "An absorption profile centred at 78 megahertz in the sky-averaged spectrum | Nature",
     "type": "webpage"
    },
    "6232408/FH99AISX": {
     "DOI": "10.1038/nature25792",
     "URL": "https://www.nature.com/articles/nature25792",
     "abstract": "After stars formed in the early Universe, their ultraviolet light is expected, eventually, to have penetrated the primordial hydrogen gas and altered the excitation state of its 21-centimetre hyperfine line. This alteration would cause the gas to absorb photons from the cosmic microwave background, producing a spectral distortion that should be observable today at radio frequencies of less than 200 megahertz1. Here we report the detection of a flattened absorption profile in the sky-averaged radio spectrum, which is centred at a frequency of 78 megahertz and has a best-fitting full-width at half-maximum of 19 megahertz and an amplitude of 0.5 kelvin. The profile is largely consistent with expectations for the 21-centimetre signal induced by early stars; however, the best-fitting amplitude of the profile is more than a factor of two greater than the largest predictions2. This discrepancy suggests that either the primordial gas was much colder than expected or the background radiation temperature was hotter than expected. Astrophysical phenomena (such as radiation from stars and stellar remnants) are unlikely to account for this discrepancy; of the proposed extensions to the standard model of cosmology and particle physics, only cooling of the gas as a result of interactions between dark matter and baryons seems to explain the observed amplitude3. The low-frequency edge of the observed profile indicates that stars existed and had produced a background of Lyman-α photons by 180 million years after the Big Bang. The high-frequency edge indicates that the gas was heated to above the radiation temperature less than 100 million years later.",
     "accessed": {
      "day": 1,
      "month": 4,
      "year": 2019
     },
     "author": [
      {
       "family": "Bowman",
       "given": "Judd D."
      },
      {
       "family": "Rogers",
       "given": "Alan E. E."
      },
      {
       "family": "Monsalve",
       "given": "Raul A."
      },
      {
       "family": "Mozdzen",
       "given": "Thomas J."
      },
      {
       "family": "Mahesh",
       "given": "Nivedita"
      }
     ],
     "container-title": "Nature",
     "id": "6232408/FH99AISX",
     "issue": "7694",
     "issued": {
      "month": 3,
      "year": 2018
     },
     "language": "en",
     "page": "67-70",
     "page-first": "67",
     "title": "An absorption profile centred at 78 megahertz in the sky-averaged spectrum",
     "type": "article-journal",
     "volume": "555"
    }
   }
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.1.-1"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
