examples/example_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Some basic test examples to run with exo21cmFAST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import py21cmfast as p21c\n",
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some useful functions for the plots\n",
    "\n",
    "def myfigure(xlims:list = None, ylims:list = None, logx:bool = True, logy:bool = True):\n",
    "    \n",
    "    fig = plt.figure(figsize=(5,4))\n",
    "    ax = fig.gca()\n",
    "    ax.grid(True, alpha = 0.5, linewidth=0.5)\n",
    "    ax.grid(True, alpha = 0.5, linewidth=0.5)\n",
    "\n",
    "    if logx is True:\n",
    "        ax.set_xscale('log')\n",
    "    if logy is True:\n",
    "        ax.set_yscale('log')\n",
    "        \n",
    "    if xlims is not None:\n",
    "        ax.set_xlim(xlims)\n",
    "    if ylims is not None:\n",
    "        ax.set_ylim(ylims)\n",
    "    \n",
    "    return fig, ax \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 0: The case without exotic energy injection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lightcone_no_exo = p21c.run_lightcone(\n",
    "            redshift = 5,        # Minimal value of redshift -> Here we will go slightly lower (cannot go below z = 4 for DarkHistory)\n",
    "            user_params = {\n",
    "                \"BOX_LEN\":                  50,    # Default value: 300  (Box length Mpc) 1000\n",
    "                \"DIM\":                      None,  # Default value: None / gives DIM=3*HII_DIM (High resolution) None\n",
    "                \"HII_DIM\":                  20,    # Default value: 200  (HII cell resolution) 350\n",
    "            },\n",
    "            flag_options = {\n",
    "                \"USE_TS_FLUCT\"               : True,  # Turn on IGM spin temperature fluctuations    \n",
    "                \"FORCE_DEFAULT_INIT_COND\"    : True,  # Forces the code to use the default initial conditions from RECFAST \n",
    "            },\n",
    "            lightcone_quantities = (),\n",
    "            global_quantities    = ('brightness_temp', 'density', 'xH_box', 'x_e_box', 'Ts_box', 'Tk_box'),\n",
    "            verbose_ntbk = True,\n",
    "            direc='./cache', \n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Plotting the results\n",
    "\n",
    "redshift_no_exo_arr   = lightcone_no_exo.node_redshifts\n",
    "Tb_global_no_exo_arr  = lightcone_no_exo.global_quantities['brightness_temp']\n",
    "xH_global_no_exo_arr  = lightcone_no_exo.global_quantities['xH_box']\n",
    "xe_global_no_exo_arr  = lightcone_no_exo.global_quantities['x_e_box']\n",
    "TS_global_no_exo_arr  = lightcone_no_exo.global_quantities['Ts_box']\n",
    "TK_global_no_exo_arr  = lightcone_no_exo.global_quantities['Tk_box']\n",
    "\n",
    "fig, ax = myfigure(xlims = [5, 35], ylims = [-180, 50], logx=False, logy=False)\n",
    "ax.set_xlabel(r\"$z$\")\n",
    "ax.set_ylabel(r\"$\\overline{T_{\\rm b}} ~ \\rm [mK]$\")\n",
    "ax.plot(redshift_no_exo_arr, Tb_global_no_exo_arr, 'k-', linewidth=1)\n",
    "fig.savefig('figures/global_Tb_with_no_exo.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Using DarkHistory for $\\chi \\to e^+e^-$\n",
    "\n",
    "In this example we evaluate the impact of dark matter decaying into electrons and positions. We choose a mass of $130$ MeV, a lifetime of $10^{26}$ s and we include the effect of backreaction. To save storage space we only output global quantities of the lightcone. To run this section you need to have installed DarkHistory and added its path to the $PYTHONPATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lightcone_DH, output_exotic_energy_injection_DH = p21c.run_lightcone(\n",
    "            redshift = 5,        # Minimal value of redshift -> Here we will go slightly lower (cannot go below z = 4 for DarkHistory)\n",
    "            user_params = {\n",
    "                \"BOX_LEN\":                  50,    # Default value: 300  (Box length Mpc) 1000\n",
    "                \"DIM\":                      None,  # Default value: None / gives DIM=3*HII_DIM (High resolution) None\n",
    "                \"HII_DIM\":                  20,    # Default value: 200  (HII cell resolution) 350\n",
    "            },\n",
    "            astro_params = {\n",
    "                \"DM_LOG10_MASS\"     : np.log10(1.3e+8),  # DM mass in eV\n",
    "                \"DM_LOG10_LIFETIME\" : np.log10(1e+26),   # Lifetime | relevant only if DM_PROCESS = 'decay'\n",
    "            },\n",
    "            flag_options = {\n",
    "                \"USE_TS_FLUCT\"               : True,          # Turn on IGM spin temperature fluctuations\n",
    "                \"USE_DM_ENERGY_INJECTION\"    : True,          # Turn on DM energy injection\n",
    "                \"DM_PROCESS\"                 : 'decay',       # Energy injection process 'swave', 'decay', ... \n",
    "                \"DM_PRIMARY\"                 : 'elec_delta',  # Primary particles (see list in user_params description)\n",
    "                \"DM_BACKREACTION\"            : True,          # Turns on backreaction       \n",
    "            },\n",
    "            lightcone_quantities = (),\n",
    "            global_quantities    = ('brightness_temp', 'density', 'xH_box', 'x_e_box', 'Ts_box', 'Tk_box'),\n",
    "            verbose_ntbk = True,\n",
    "            output_exotic_data=True,\n",
    "            direc='./cache', \n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Plotting the results\n",
    "\n",
    "redshift_DH_arr   = lightcone_DH.node_redshifts\n",
    "Tb_global_DH_arr  = lightcone_DH.global_quantities['brightness_temp']\n",
    "xH_global_DH_arr  = lightcone_DH.global_quantities['xH_box']\n",
    "xe_global_DH_arr  = lightcone_DH.global_quantities['x_e_box']\n",
    "TS_global_DH_arr  = lightcone_DH.global_quantities['Ts_box']\n",
    "TK_global_DH_arr  = lightcone_DH.global_quantities['Tk_box']\n",
    "\n",
    "fig, ax = myfigure(xlims = [5, 35], ylims = [-180, 50], logx=False, logy=False)\n",
    "ax.set_xlabel(r\"$z$\")\n",
    "ax.set_ylabel(r\"$\\overline{T_{\\rm b}} ~ \\rm [mK]$\")\n",
    "ax.plot(redshift_DH_arr, Tb_global_DH_arr, 'b-', linewidth=1)\n",
    "ax.plot(redshift_no_exo_arr, Tb_global_no_exo_arr, 'k-', linewidth=1)\n",
    "fig.savefig('figures/global_Tb_with_DH.pdf', bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Using $f_{\\rm heat}$ templates\n",
    "\n",
    "In this example we evaluate the impact of dark matter decay with $f_{\\rm heat}$ represented by a Schechter function, \n",
    "\\begin{equation}\n",
    "f_{\\rm heat}(z) = f_0 e^{a (z-z_0)} \\left( \\frac{z}{z_0}\\right)^b\n",
    "\\end{equation}\n",
    "where $z_0=15$ and the parameters $(f_0, a, b)$ are set by hand (here matching values obtained from DarkHistory). We choose a mass of $10$ GeV (although irrelevant in that case for decay) and a lifetime of $10^{26}$ s. Here we also impose by hand initial conditions for $x_e$ and $T_{\\rm K}$ (here matching values obtained from DarkHistory). The values of the deposition fractions in other canals are evaluated from $f_{\\rm heat}$ according to\n",
    "\\begin{equation}\n",
    "\\begin{cases}\n",
    "& f_{\\rm ion, H} =  2.2 f_{\\rm heat}\\\\\n",
    "& f_{\\rm ion, He} =  0.07 f_{\\rm heat}\\\\\n",
    "& f_{\\rm exc} =  1.7 f_{\\rm heat} \\, .\n",
    "\\end{cases}\n",
    "\\end{equation}\n",
    "Custom values for the coefficients of proportionality can be set in astro_params with the flag_option \"USE_DM_CUSTOM_F_RATIOS\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lightcone_templates, output_exotic_energy_injection_templates = p21c.run_lightcone(\n",
    "            redshift = 5,        # Minimal value of redshift -> Here we will go slightly lower (cannot go below z = 4 for DarkHistory)\n",
    "            user_params = {\n",
    "                \"BOX_LEN\":                  50,    # Default value: 300  (Box length Mpc) 1000\n",
    "                \"DIM\":                      None,  # Default value: None / gives DIM=3*HII_DIM (High resolution) None\n",
    "                \"HII_DIM\":                  20,    # Default value: 200  (HII cell resolution) 350\n",
    "            },\n",
    "            astro_params = {\n",
    "                \"DM_LOG10_MASS\"     : np.log10(1.3e+8),  # DM mass in eV\n",
    "                \"DM_LOG10_LIFETIME\" : np.log10(1e+26),   # Lifetime | relevant only if DM_PROCESS = 'decay'\n",
    "\n",
    "                \"DM_FHEAT_APPROX_PARAM_LOG10_F0\" : np.log10(0.1531),     # Parameter to feed to the template of fheat\n",
    "                \"DM_FHEAT_APPROX_PARAM_A\"        : 0.003209,            # Parameter to feed to the template of fheat\n",
    "                \"DM_FHEAT_APPROX_PARAM_B\"        : 0.1295,               # Parameter to feed to the template of fheat\n",
    "\n",
    "                \"LOG10_XION_at_Z_HEAT_MAX\"      : np.log10(0.0031),      # Ionization fraction at redshift z = Z_HEAT_MAX = 35 (or max_redshift)\n",
    "                \"LOG10_TK_at_Z_HEAT_MAX\"        : np.log10(118)          # Temperature in K at redshift z = Z_HEAT_MAX = 35 (or max_redshift)\n",
    "            },\n",
    "            flag_options = {\n",
    "                \"USE_TS_FLUCT\"               : True,          # Turn on IGM spin temperature fluctuations\n",
    "                \"USE_DM_ENERGY_INJECTION\"    : True,          # Turn on DM energy injection\n",
    "                \"USE_DM_EFFECTIVE_DEP_FUNCS\" : True ,         # Treat the energy injection with approximate templates (instead of DarkHistory)\n",
    "                \"DM_PROCESS\"                 : 'decay',       # Energy injection process 'swave', 'decay', ... \n",
    "                \"DM_PRIMARY\"                 : 'elec_delta',  # Primary particles (see list in user_params description)\n",
    "                \"DM_BACKREACTION\"            : True,          # Turns on backreaction \n",
    "                \"DM_FHEAT_APPROX_SHAPE\"      : 'schechter',    # Shape of the template for f_heat\n",
    "                \"USE_CUSTOM_INIT_COND\"       : True           # Forces the code to use the init conditions defined in astro_params\n",
    "            },\n",
    "            lightcone_quantities = (),\n",
    "            global_quantities    = ('brightness_temp', 'density', 'xH_box', 'x_e_box', 'Ts_box', 'Tk_box'),\n",
    "            verbose_ntbk = True,\n",
    "            output_exotic_data=True,\n",
    "            direc='./cache', \n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Plotting the results\n",
    "\n",
    "redshift_templates_arr   = lightcone_templates.node_redshifts\n",
    "Tb_global_templates_arr  = lightcone_templates.global_quantities['brightness_temp']\n",
    "xH_global_templates_arr  = lightcone_templates.global_quantities['xH_box']\n",
    "xe_global_templates_arr  = lightcone_templates.global_quantities['x_e_box']\n",
    "TS_global_templates_arr  = lightcone_templates.global_quantities['Ts_box']\n",
    "TK_global_templates_arr  = lightcone_templates.global_quantities['Tk_box']\n",
    "\n",
    "## Plotting the brightness temperature\n",
    "## Comment the corresponding lines if you have not run the first two examples before\n",
    "fig, ax = myfigure(xlims = [5, 35], ylims = [-180, 50], logx=False, logy=False)\n",
    "ax.set_xlabel(r\"$z$\")\n",
    "ax.set_ylabel(r\"$\\overline{T_{\\rm b}} ~ \\rm [mK]$\")\n",
    "ax.plot(redshift_templates_arr, Tb_global_templates_arr, 'r-', linewidth=1)\n",
    "ax.plot(redshift_DH_arr, Tb_global_DH_arr, 'b-', linewidth=1)\n",
    "ax.plot(redshift_no_exo_arr, Tb_global_no_exo_arr, 'k-', linewidth=1)\n",
    "fig.savefig('figures/global_Tb_with_templates.pdf', bbox_inches='tight')\n",
    "\n",
    "## Plotting the spin and the IGM temperature\n",
    "## Comment the corresponding lines if you have not run the first two examples before\n",
    "fig, ax = myfigure(xlims = [5, 35], logx=False)\n",
    "ax.set_xlabel(r\"$z$\")\n",
    "ax.set_ylabel(r\"$\\overline{T} ~ \\rm [K]$\")\n",
    "line1, = ax.plot(redshift_templates_arr, TS_global_templates_arr, 'r-', linewidth=1)\n",
    "line2, = ax.plot(redshift_templates_arr, TK_global_templates_arr, 'r--', linewidth=1)\n",
    "line1, = ax.plot(redshift_DH_arr, TS_global_DH_arr, 'b-', linewidth=1)\n",
    "line2, = ax.plot(redshift_DH_arr, TK_global_DH_arr, 'b--', linewidth=1)\n",
    "line1, = ax.plot(redshift_no_exo_arr, TS_global_no_exo_arr, 'k-', linewidth=1)\n",
    "line2, = ax.plot(redshift_no_exo_arr, TK_global_no_exo_arr, 'k--', linewidth=1)\n",
    "ax.legend([line1, line2], [r\"$\\rm T_{\\rm S}$\", r\"$\\rm T_{\\rm K}$\"])\n",
    "fig.savefig('figures/global_TS_TK_with_templates.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.6 ('exo21cmFAST_release': venv)",
   "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.9.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "b978f4ecbf43ad133bda540fcbea492cd113875f61480fb8ce8c71a82633a1e6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}