{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "41c096f8-fe6c-4ccf-b471-df63f90c337e",
   "metadata": {},
   "source": [
    "# Example 12 Using Neural Network transfer functions (NNTFs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c48fc8f-f4a2-4e66-af03-f69880f91e5a",
   "metadata": {},
   "source": [
    "In this notebook, we demonstrate how to generate a transfer function from a Neural Network (more specifically, a Multi-layer Perceptron network), compare it against table transfer functions (if available), and use it in `main.evolve` with the `tf_mode` switch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7acfb272-4d5f-453d-890f-500976062b8e",
   "metadata": {},
   "source": [
    "## Notebook Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7806d8d-069e-4682-bf93-ba073f649e49",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "sys.path.append('..') # append DarkHistory's path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47dae39e-df71-4898-b50b-4d97734a000d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "matplotlib.rc_file('matplotlibrc_nntf')\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e077738a-08f2-48bb-9e39-6246c90c8053",
   "metadata": {},
   "outputs": [],
   "source": [
    "from darkhistory.config import load_data\n",
    "from darkhistory.nntf.load import load_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4cfa4d3-f882-4ced-8d5c-cdb596f0ca24",
   "metadata": {},
   "source": [
    "## 1. Generate a NNTF\n",
    "We use the coarsened high energy photon transfer function `hep_p12` as an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11a4261c-7f3e-42c3-a2f6-ade6469b78d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "hep_p12_nntf = load_model('dep_nntf')['hep_p12']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "555fb646-dbcf-464b-9757-39e144d0ba56",
   "metadata": {},
   "source": [
    "For some value of redshift $rs$=`rs`, Hydrogen ionization fraction $x_\\text{HII}$=`xH`, Helium first ionization fraction $x_\\text{HeII}$=`xHe`, let the Neural Network (NN) model predict the transfer function value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fb46d48-ee82-4345-bc12-91844718d80a",
   "metadata": {},
   "outputs": [],
   "source": [
    "hep_p12_nntf.predict_TF(rs=300, xH=0.6, xHe=0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d68a96ce-aa7d-44d1-a35e-4ddf168f77f5",
   "metadata": {},
   "source": [
    "Now plot the transfer functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bc5dbad-3214-4cb9-96f7-a4feb2e42842",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10,8))\n",
    "\n",
    "extent = np.log10((hep_p12_nntf.abscs[0][0], hep_p12_nntf.abscs[0][-1],\n",
    "                   hep_p12_nntf.abscs[1][0], hep_p12_nntf.abscs[1][-1]))\n",
    "\n",
    "posim = ax.imshow(np.flipud(np.log10(np.clip(hep_p12_nntf.TF, 1e-50, 1e50))), vmax=7, vmin=-50, cmap='plasma', extent=extent)\n",
    "negim = ax.imshow(np.flipud(np.log10(-hep_p12_nntf.TF)), vmax=7, vmin=-10, cmap='viridis', extent=extent)\n",
    "fig.colorbar(negim, ax=ax, shrink=0.8, label=r'$\\log_{10}$ of negative value')\n",
    "fig.colorbar(posim, ax=ax, shrink=0.8, label=r'$\\log_{10}$ of positive value')\n",
    "\n",
    "ax.set(title='High energy photon transfer function NN prediction',\n",
    "       xlabel=r'$\\log_{10}E_\\mathrm{out}/\\mathrm{eV}$', ylabel=r'$\\log_{10}E_\\mathrm{in}/\\mathrm{eV}$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fdb222e-2b0e-4c4d-a7ee-58dc6ca7949d",
   "metadata": {},
   "source": [
    "## 2. Comparison with table transfer functions (if tables available)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64beb39a-7cf1-40f2-bb64-87a7419fb08e",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    hep_tf = load_data('dep_tf')['highengphot']\n",
    "    tf_E = load_data('tf_helper')['tf_E']\n",
    "except (FileNotFoundError, ValueError) as err:\n",
    "    print(type(err).__name__, ':', err)\n",
    "    print('\\nNeed table transfer functions for comparison.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f1395ce-e19d-43b7-a0ca-55ca3684c469",
   "metadata": {},
   "source": [
    "Again using high energy photon transfer function as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc7929ec-a6f9-473f-82e4-ed2ac4cb0086",
   "metadata": {},
   "outputs": [],
   "source": [
    "rs = 300\n",
    "xH = 0.6\n",
    "xHe = 0.00\n",
    "\n",
    "hep_tf_slice = hep_tf.get_tf(xH, xHe, rs)\n",
    "z_tb = np.linalg.matrix_power(hep_tf_slice.grid_vals, 12)\n",
    "\n",
    "hep_E = tf_E.get_val(xH, xHe, rs)[0]\n",
    "hep_p12_nntf.predict_TF(rs=rs, xH=xH, xHe=xHe, E_arr=hep_E)\n",
    "z_nn = hep_p12_nntf.TF\n",
    "\n",
    "z_ediff = z_nn-z_tb\n",
    "absc_in  = load_data('binning')['phot']\n",
    "absc_out = load_data('binning')['phot']\n",
    "for ii, Ein in enumerate(absc_in):\n",
    "    for io, Eout in enumerate(absc_out):\n",
    "        z_ediff[ii][io] *= (Eout/Ein)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb8f9f25-a2d4-4378-83f8-99e874f3df13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot definitions\n",
    "i_absc, o_absc = hep_p12_nntf.abscs\n",
    "extent = np.log10((o_absc[0], o_absc[-1], i_absc[0], i_absc[-1]))\n",
    "xlabel_str = r'$\\log_{10}E_\\mathrm{out}/\\mathrm{eV}$'\n",
    "ylabel_str = r'$\\log_{10}E_\\mathrm{in} /\\mathrm{eV}$'\n",
    "EPSILON = 1e-50\n",
    "plot_val = lambda x: np.flipud(np.log10(np.clip(np.abs(x), EPSILON, None)))\n",
    "contour_color = 'k'\n",
    "\n",
    "zc_tb = z_tb.copy()\n",
    "zc_tb[:,100:] = 1e-50 # set consistent epsilons/zeros\n",
    "\n",
    "# plot\n",
    "fig, axs = plt.subplots(1, 4, figsize=(18, 6))\n",
    "\n",
    "imtf = axs[0].imshow(plot_val(z_nn), vmin=-40, vmax=5, cmap='plasma', extent=extent)\n",
    "axs[0].contour(z_nn+1e-50, levels=[0], extent=extent, colors=contour_color, linestyles='--', linewidths=2)\n",
    "axs[0].plot([extent[0]+0.05, extent[0]+0.05, -0.7], [3.5, extent[3]-0.05, extent[3]-0.05], color=contour_color, linestyle='--', linewidth=3)\n",
    "\n",
    "axs[1].imshow(plot_val(z_tb), vmin=-40, vmax=5, cmap='plasma', extent=extent)\n",
    "axs[1].contour(zc_tb+1e-50, levels=[0], extent=extent, colors=contour_color, linestyles='--', linewidths=2)\n",
    "axs[1].plot([extent[0]+0.05, extent[0]+0.05, -0.7], [3.5, extent[3]-0.05, extent[3]-0.05], color=contour_color, linestyle='--', linewidth=3)\n",
    "\n",
    "imd1 = axs[2].imshow(plot_val(z_nn)-plot_val(z_tb), vmin=-1, vmax=1, cmap='coolwarm', extent=extent)\n",
    "imd2 = axs[3].imshow(np.flipud(z_ediff), vmin=-0.002, vmax=0.002, cmap='coolwarm', extent=extent)\n",
    "\n",
    "axs[0].set(title=r'NN $\\log_{10}P$', xlabel=xlabel_str, ylabel=ylabel_str)\n",
    "axs[1].set(title=r'Table $\\log_{10}P$', xlabel=xlabel_str)\n",
    "axs[2].set(title=r'Error of $\\log_{10}P$', xlabel=xlabel_str)\n",
    "axs[3].set(title=r'Error in energy transfer', xlabel=xlabel_str)\n",
    "\n",
    "fig.colorbar(imtf, ax=axs[:2], orientation='horizontal', aspect=40)\n",
    "fig.colorbar(imd1, ax=axs[2], orientation='horizontal', aspect=20)\n",
    "#fig.colorbar(imd2, ax=axs[3], orientation='horizontal', aspect=20)\n",
    "cbar = fig.colorbar(imd2, ax=axs[3], orientation='horizontal', aspect=20,\n",
    "                 ticks=[-0.002, -0.001, 0, 0.001, 0.002])\n",
    "cbar.ax.set_xticklabels([r'$-0.2$\\%', r'$-0.1$\\%', r'$0$\\%', r'$0.1$\\%', r'$0.2$\\%'])\n",
    "\n",
    "axs[0].text(5, -2, '$1+z=%d$\\n$x_\\\\mathrm{HII}=%.1f$\\n$x_\\\\mathrm{HeII}=%.1f$' % (rs, xH, xHe), color='white', fontsize=18)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86d57ffa-20a8-4656-acaa-2e275581b1a7",
   "metadata": {},
   "source": [
    "Note that in the area encircled by the black dashed line the transfer function value is negative."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac8c14b0-c4da-48f9-831c-fba75ca31723",
   "metadata": {},
   "source": [
    "## 3. Using NN option in `main.evolve`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d499126-492f-43fe-ad5f-b4e1a61ec9b2",
   "metadata": {},
   "source": [
    "We will be using the same example scenario as in `Example_8_Ionization_and_Thermal_History.ipynb`: a 100 MeV dark matter particle decaying to electron positron pairs with a lifetime of $\\tau=3\\times10^{25}$ s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91124543-69b3-4114-852e-b2cfae828ba9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import darkhistory.main as main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8e0e8bf-5642-47ef-93e8-8e5964f50506",
   "metadata": {},
   "outputs": [],
   "source": [
    "soln = main.evolve(\n",
    "    DM_process     = 'decay',\n",
    "    mDM            = 1e8,\n",
    "    lifetime       = 3e25,\n",
    "    primary        = 'elec_delta',\n",
    "    backreaction   = True,\n",
    "    helium_TLA     = True,\n",
    "    reion_switch   = True,\n",
    "    \n",
    "    start_rs       = 3000,\n",
    "    end_rs         = 4,\n",
    "    coarsen_factor = 12, # Will be set to 12 automatically if using NN transfer functions.\n",
    "    tf_mode        = 'nn',\n",
    "    \n",
    "    use_tqdm       = True,\n",
    "    verbose        = 1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9fcd069-8d5d-415d-a241-9edb3d4e8ba2",
   "metadata": {},
   "source": [
    "## 3.1 Plot matter temperature and ionization history"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "264ad641-fdef-4758-aa3b-c723628db837",
   "metadata": {},
   "source": [
    "First get the standard cosmology temperature and ionization history solution with helium three-level atom and reionization turned on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e3a9e8-1dd2-4e5e-a59e-36d67baadd07",
   "metadata": {},
   "outputs": [],
   "source": [
    "import darkhistory.history.tla as tla\n",
    "\n",
    "soln_baseline = tla.get_history(soln['rs'], helium_TLA=True, reion_switch=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8af0f56c-3ee0-4aa3-9295-e2fbfeb5b972",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 3, figsize=(16,4.5))\n",
    "\n",
    "x = soln['rs']\n",
    "y = soln['Tm']\n",
    "ydflt = soln_baseline[:,0]\n",
    "axs[0].plot(x, y    , 'r'  , lw=2, alpha=1.0, label='DM (using NN)')\n",
    "axs[0].plot(x, ydflt, 'k--', lw=2, alpha=0.5, label='no DM')\n",
    "axs[0].set(xscale='log', xlim=(4,3000), yscale='log', ylim=(1e-4,10))\n",
    "axs[0].set(title='Matter temperature', xlabel=r'Redshift $1+z$', ylabel=r'$T_m$ [eV]')\n",
    "textstr =  '$\\chi \\\\to e^+e^-$, $m_\\chi = 0.1$~GeV\\n'\n",
    "textstr += '$\\\\tau = 3 \\\\times 10^{25}$ s'\n",
    "axs[0].text(30, 0.3, textstr)\n",
    "axs[0].legend()\n",
    "\n",
    "y = soln['x'][:,0]\n",
    "ydflt = soln_baseline[:,1]\n",
    "axs[1].plot(x, y    , 'r'  , lw=2, alpha=1.0, label='DM')\n",
    "axs[1].plot(x, ydflt, 'k--', lw=2, alpha=0.5, label='no DM')\n",
    "axs[1].set(xscale='log', xlim=(None, None), yscale='log', ylim=(1e-4, 2))\n",
    "axs[1].set(title='Ionized hydrogen fraction', xlabel=r'Redshift $1+z$', ylabel=r'$x_\\mathrm{HII}=n_\\mathrm{HII}/n_\\mathrm{H}$')\n",
    "\n",
    "y = soln['x'][:,1]\n",
    "ydflt = soln_baseline[:,2]\n",
    "axs[2].plot(x, y    , 'r'  , lw=2, alpha=1.0, label='DM')\n",
    "axs[2].plot(x, ydflt, 'k--', lw=2, alpha=0.5, label='no DM')\n",
    "axs[2].set(xscale='log', yscale='log', ylim=(1e-10, 1))\n",
    "axs[2].set(title='Singly ionized helium fraction', xlabel=r'Redshift $1+z$', ylabel=r'$x_\\mathrm{HeII}=n_\\mathrm{HeII}/n_\\mathrm{H}$')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7b91c86-d398-48bd-9561-94d40e65a16a",
   "metadata": {},
   "source": [
    "## 3.2 Plot low energy photon spectral distortion\n",
    "(not contribution from low energy electrons and atomic transitions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de7b70a2-e782-4684-a89d-675e6f44ce8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "specs = soln['lowengphot']\n",
    "summed_spec = specs[0]*0\n",
    "for i in range(len(soln['rs'])):\n",
    "    spec = specs[i]*1\n",
    "    spec.redshift(1)\n",
    "    summed_spec += spec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "221ba469-422d-444e-bc78-93660ad96f5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "ax.plot(summed_spec.eng, summed_spec.N, color='r', label='DM (using NN)')\n",
    "ax.axhline(y=0, color='k', alpha=0.5)\n",
    "ax.set(xscale='log', xlim=(8e-5,1e2))\n",
    "ax.ticklabel_format(axis='y', style='sci', scilimits=(0,0))\n",
    "ax.set_title(r'Photon spectral distortion', fontsize=18)\n",
    "ax.set_xlabel(r'$E$ [eV]', fontsize=16)\n",
    "ax.set_ylabel(r'$E\\,\\mathrm{d}N/\\mathrm{d}E$', fontsize=16)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d91bd5ea",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
