{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c2dac99b",
   "metadata": {},
   "source": [
    "# Fornax 2019 Models\n",
    "\n",
    "Some I/O and plotting using the Fornax 3D models from [Vartanyan, Burrows, et al., MNRAS 482(1):351, 2019](https://arxiv.org/abs/1809.05106), which express the angular dependence of the neutrino luminosity from CCSNe in terms of a spherical harmonic expansion up to degree $\\ell=2$. Note: Caching and plotting the angular dependence requires the `healpy` package, which can be installed via `pip install healpy`.\n",
    "\n",
    "The data are available on [the Burrows group website](https://www.astro.princeton.edu/~burrows/nu-emissions.3d/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4df4d343",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snewpy.neutrino import Flavor\n",
    "from snewpy.models.ccsn import Fornax_2019\n",
    "\n",
    "from astropy import units as u\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "try:\n",
    "    import healpy as hp\n",
    "    healpy_available = True\n",
    "except ModuleNotFoundError:\n",
    "    healpy_available = False\n",
    "\n",
    "mpl.rc('font', size=16)\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61a9c046",
   "metadata": {},
   "source": [
    "## Initialization and Time Evolution of the Energy Spectrum\n",
    "\n",
    "To start, let’s see what progenitors are available for the `Fornax_2019` model. We can use the `param` property to view all physics parameters and their possible values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83a8ee2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "Fornax_2019.param"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbe1497e",
   "metadata": {},
   "source": [
    "We’ll initialise one of these progenitors and plot the evolution of the predicted neutrino spectrum across four time steps. If this is the first time you’re using a progenitor, snewpy will automatically download the required data files for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c50291f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Fornax_2019(progenitor_mass=10*u.solMass)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7ea4be2",
   "metadata": {},
   "outputs": [],
   "source": [
    "E = np.arange(0,111) * u.MeV\n",
    "theta = 23.55*u.degree\n",
    "phi = 22.5*u.degree\n",
    "\n",
    "fig, axes = plt.subplots(2,2, figsize=(11,8), sharex=True, sharey=True,\n",
    "                         gridspec_kw = {'hspace':0.04, 'wspace':0.025})\n",
    "axes = axes.flatten()\n",
    "\n",
    "for ax, t in zip(axes, [10., 100., 300., 600.]):\n",
    "    t *= u.ms\n",
    "    spectra = model.get_initial_spectra(t, E, theta, phi)\n",
    "    \n",
    "    lines = ['--', '-.', '-', ':']\n",
    "    for line, flavor in zip(lines, Flavor):\n",
    "        ax.plot(E, spectra[flavor].to('1e52 erg/(s MeV)'), ls=line, label=flavor.to_tex())\n",
    "    ax.legend(fontsize=14, ncol=2, loc='upper right', title='{:g}'.format(t))\n",
    "    ax.set(xlim=(E[0].to_value('MeV'), E[-1].to_value('MeV')))\n",
    "    ax.grid(ls=':')\n",
    "\n",
    "axes[0].set(ylabel=r'$dL/dE$ [$10^{52}$ erg s$^{-1}$ MeV$^{-1}$]')\n",
    "axes[2].set(xlabel=r'energy [MeV]');\n",
    "\n",
    "fig.suptitle('Time Evolution of Spectra'.format(t))\n",
    "fig.subplots_adjust(top=0.925, bottom=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89906e58",
   "metadata": {},
   "source": [
    "## Test Caching of Full Model Angular Dependence\n",
    "\n",
    "Here, if the user specifies `cache_flux=True` in the model constructor, the entire angular dependence of the model will be computed using a [HEALPix grid](https://healpy.readthedocs.io/en/latest/) with 192 pixels, using a HEALPix grid with parameter `nside=4` and `ordering=RING`. The model constructor will calculate the CCSN fluxes on the unit sphere at the centers of the 192 healpixels for all energy and time bins in the model.\n",
    "\n",
    "The resulting tables are then saved to a FITS file in the same location as the input HDF5 file. If the user instantiates this model with `cache_flux=True` in the future, the model will be initialized using the cached FITS file.\n",
    "\n",
    "Comments:\n",
    "1. The HDF5 files store the angular dependence of the models in a highly compressed format using the spherical harmonic expansion up to degree $\\ell=2$.\n",
    "1. The time bins are equal across all flavors but the energy binning is a function of both time and flavor.\n",
    "1. Computing the fluxes at 192 locations on the sky takes about 30 seconds per flavor due to the overhead of reading out the HDF5 files. So construction of the cached model is **slow** when called for the first time.\n",
    "1. With `nside=4` and 192 bins, the angular resolution per bin is about $15^\\circ$, which is reasonable given the highest-resolution features of the models are the quadrupole moments ($\\ell=2$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03fcdbb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmodel = Fornax_2019(progenitor_mass=10*u.solMass, cache_flux=True)\n",
    "cmodel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfc21f0c",
   "metadata": {},
   "source": [
    "### Test Equivalence of the Cached and Uncached Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec5bc0c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 100 * u.ms\n",
    "\n",
    "Es, dEs, spectras = model._get_binnedspectra(t, theta, phi)\n",
    "Ec, dEc, spectrac = cmodel._get_binnedspectra(t, theta, phi)\n",
    "\n",
    "for flavor in Flavor:\n",
    "    print('\\n{}'.format(str(flavor)))\n",
    "    print(Es[flavor])\n",
    "    print(Ec[flavor])\n",
    "    print(spectras[flavor])\n",
    "    print(spectrac[flavor])\n",
    "    \n",
    "    j = np.argmax(spectrac[flavor])\n",
    "    print('Max value:')\n",
    "    print(Ec[flavor][j], spectrac[flavor][j])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de0e525b",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 600 * u.ms\n",
    "E = np.arange(0,111) * u.MeV\n",
    "theta = 23.55*u.degree\n",
    "phi = 22.5*u.degree\n",
    "\n",
    "spectra = model.get_initial_spectra(t, E, theta, phi)\n",
    "cspectra = cmodel.get_initial_spectra(t, E, theta, phi)\n",
    "    \n",
    "fig, axes = plt.subplots(2,2, figsize=(11,8), sharex=True, sharey=True,\n",
    "                         gridspec_kw = {'hspace':0.04, 'wspace':0.025})\n",
    "axes = axes.flatten()\n",
    "\n",
    "spectra = model.get_initial_spectra(t, E, theta, phi)\n",
    "\n",
    "for ax, flavor in zip(axes, Flavor):\n",
    "    ax.plot(E, spectra[flavor].to('1e52 erg/(s MeV)'), ls='-', lw=1, label='{}: model'.format(flavor.to_tex()))\n",
    "    ax.plot(E, cspectra[flavor].to('1e52 erg/(s MeV)'), ls=':', lw=3, label='{}: cached model'.format(flavor.to_tex()))\n",
    "    ax.legend(fontsize=14, ncol=1, loc='upper right')\n",
    "    \n",
    "    ax.set(xlim=(E[0].to_value('MeV'), E[-1].to_value('MeV')))\n",
    "    ax.grid(ls=':')\n",
    "\n",
    "axes[0].set(ylabel=r'$dL/dE$ [$10^{52}$ erg s$^{-1}$ MeV$^{-1}$]')\n",
    "axes[2].set(xlabel=r'energy [MeV]')\n",
    "\n",
    "fig.suptitle('Spectra at $t={{{:g}}}$'.format(t))\n",
    "fig.subplots_adjust(top=0.925, bottom=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f9fed98",
   "metadata": {},
   "source": [
    "### Plot $L_\\nu(t,\\theta,\\varphi)$ at Several Times\n",
    "\n",
    "Plot the angular dependence $L(t)$ for several values of $t$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d880473",
   "metadata": {},
   "outputs": [],
   "source": [
    "if healpy_available:\n",
    "    times = np.asarray([10., 100., 300., 600.]) * u.ms\n",
    "    nt = len(times)\n",
    "\n",
    "    fig, axes = plt.subplots(nt,3, figsize=(8*3, nt*4.5))\n",
    "\n",
    "    L_nue = cmodel.luminosity[Flavor.NU_E].to('1e52*erg/s')\n",
    "    L_nue_bar = cmodel.luminosity[Flavor.NU_E_BAR].to('1e52*erg/s')\n",
    "    L_nux = cmodel.luminosity[Flavor.NU_X].to('1e52*erg/s')\n",
    "\n",
    "    for i, t in enumerate(times):\n",
    "        j = (np.abs(t - cmodel.time)).argmin()\n",
    "        \n",
    "        ax = axes[i,0]\n",
    "        plt.axes(ax)\n",
    "        hp.mollview(L_nue[j], title=r'$L_{{\\nu_e}}(t={:.3f})$'.format(cmodel.time[j]),\n",
    "                    unit=r'$10^{52}$ erg s$^{-1}$', cmap='viridis',\n",
    "                    hold=True)\n",
    "        \n",
    "        ax = axes[i,1]\n",
    "        plt.axes(ax)\n",
    "        hp.mollview(L_nue_bar[j], title=r'$L_{{\\bar{{\\nu}}_e}}(t={:.3f})$'.format(cmodel.time[j]),\n",
    "                    unit=r'$10^{52}$ erg s$^{-1}$', cmap='magma',\n",
    "                    hold=True)\n",
    "        \n",
    "        ax = axes[i,2]\n",
    "        plt.axes(ax)\n",
    "        hp.mollview(L_nux[j], title=r'$L_{{\\nu_X}}(t={:.3f})$'.format(cmodel.time[j]),\n",
    "                    unit=r'$10^{52}$ erg s$^{-1}$', cmap='cividis',\n",
    "                    hold=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "824a3550",
   "metadata": {},
   "source": [
    "### Plot $\\Delta L_\\nu(t,\\theta,\\varphi) / \\langle L_\\nu(t,\\theta,\\varphi)\\rangle$ at Several Times\n",
    "\n",
    "Same plot as above, but here plot the deviation from the average over the sky.\n",
    "\n",
    "This demonstrates how the deviation from isotropy is very small initially and approaches 20% at later times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "128e8cac",
   "metadata": {},
   "outputs": [],
   "source": [
    "if healpy_available:\n",
    "    times = np.asarray([10., 100., 300., 600.]) * u.ms\n",
    "    nt = len(times)\n",
    "\n",
    "    fig, axes = plt.subplots(nt,3, figsize=(8*3, nt*4.5))\n",
    "\n",
    "    L_nue = cmodel.luminosity[Flavor.NU_E].to('1e52*erg/s')\n",
    "    L_nue_bar = cmodel.luminosity[Flavor.NU_E_BAR].to('1e52*erg/s')\n",
    "    L_nux = cmodel.luminosity[Flavor.NU_X].to('1e52*erg/s')\n",
    "\n",
    "    vmin, vmax = -0.2, 0.2\n",
    "\n",
    "    for i, t in enumerate(times):\n",
    "        j = (np.abs(t - cmodel.time)).argmin()\n",
    "        \n",
    "        ax = axes[i,0]\n",
    "        plt.axes(ax)\n",
    "        Lavg = np.average(L_nue[j])\n",
    "        dL_on_L = (L_nue[j] - Lavg) / Lavg\n",
    "        hp.mollview(dL_on_L, title=r'$\\nu_e$: $t=${:.3f}'.format(cmodel.time[j]),\n",
    "                    unit=r'$(L_\\nu - \\langle L_\\nu\\rangle) / \\langle L_\\nu\\rangle$', cmap='viridis',\n",
    "                    min=vmin, max=vmax,\n",
    "                    hold=True)\n",
    "        \n",
    "        ax = axes[i,1]\n",
    "        plt.axes(ax)\n",
    "        Lavg = np.average(L_nue_bar[j])\n",
    "        dL_on_L = (L_nue_bar[j] - Lavg) / Lavg\n",
    "        hp.mollview(dL_on_L, title=r'$\\bar{{\\nu}}_e$: $t=${:.3f}'.format(cmodel.time[j]),\n",
    "                    unit=r'$\\Delta L_{\\bar{\\nu}_e}/\\langle L_{\\bar{\\nu}_e}\\rangle$', cmap='magma',\n",
    "                    min=vmin, max=vmax,\n",
    "                    hold=True)\n",
    "        \n",
    "        ax = axes[i,2]\n",
    "        plt.axes(ax)\n",
    "        Lavg = np.average(L_nux[j])\n",
    "        dL_on_L = (L_nux[j] - Lavg) / Lavg\n",
    "        hp.mollview(dL_on_L, title=r'$\\nu_X$: $t=${:.3f}'.format(cmodel.time[j]),\n",
    "                    unit=r'$(L_\\nu - \\langle L_\\nu\\rangle) / \\langle L_\\nu\\rangle$', cmap='cividis',\n",
    "                    min=vmin, max=vmax,\n",
    "                    hold=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6beba86",
   "metadata": {},
   "source": [
    "### Superimpose $L_\\nu(t,\\theta,\\varphi)$ at All Locations on the Sphere\n",
    "\n",
    "Plot $L(t)$ at all locations as well as the average, and then the deviation from the average."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c342a3c8",
   "metadata": {},
   "source": [
    "#### Electron Neutrinos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea2d9627",
   "metadata": {},
   "outputs": [],
   "source": [
    "if healpy_available:\n",
    "    Lavg = np.average(L_nue, axis=1)\n",
    "    dL_over_L = (L_nue - Lavg[:, np.newaxis]) / Lavg[:, np.newaxis]\n",
    "\n",
    "    fig, axes = plt.subplots(2, 1, figsize=(10, 8),\n",
    "                             gridspec_kw={'height_ratios': [3, 1], 'hspace': 0},\n",
    "                             sharex=True, tight_layout=True)\n",
    "\n",
    "    ax = axes[0]\n",
    "    ax.plot(cmodel.time, L_nue, color='gray', alpha=0.05)\n",
    "    ax.plot(cmodel.time, Lavg, color='r', ls='--')\n",
    "    ax.set(xlim=cmodel.time[0::len(cmodel.time)-1].value,\n",
    "           ylabel=r'$L_{\\nu_e}(t,\\theta,\\varphi)$ [$10^{52}$ erg s$^{-1}$]',\n",
    "           ylim=(0, 1.1*np.max(L_nue).value),\n",
    "           title=r'$L_{\\nu_e}(t,\\theta,\\varphi)$')\n",
    "    ax.grid(ls=':')\n",
    "\n",
    "    ax = axes[1]\n",
    "    ax.plot(cmodel.time, dL_over_L, color='gray', alpha=0.1)\n",
    "    ax.plot(cmodel.time, np.zeros(cmodel.time.shape), ls='--', color='r')\n",
    "    ax.set(xlabel='time [s]',\n",
    "           ylabel=r'$(L_\\nu - \\bar{L}_\\nu)/\\bar{L}_\\nu$',\n",
    "           ylim=(-0.3, 0.3))\n",
    "    ax.grid(ls=':')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94bfc6a9",
   "metadata": {},
   "source": [
    "#### Electron Antineutrinos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "643757c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if healpy_available:\n",
    "    Lavg = np.average(L_nue_bar, axis=1)\n",
    "    dL_over_L = (L_nue_bar - Lavg[:, np.newaxis]) / Lavg[:, np.newaxis]\n",
    "\n",
    "    fig, axes = plt.subplots(2, 1, figsize=(10, 8),\n",
    "                             gridspec_kw={'height_ratios': [3, 1], 'hspace': 0},\n",
    "                             sharex=True, tight_layout=True)\n",
    "\n",
    "    ax = axes[0]\n",
    "    ax.plot(cmodel.time, L_nue_bar, color='gray', alpha=0.05)\n",
    "    ax.plot(cmodel.time, Lavg, color='r', ls='--')\n",
    "    ax.set(xlim=cmodel.time[0::len(cmodel.time)-1].value,\n",
    "           ylabel=r'$L_{\\bar{\\nu}_e}(t,\\theta,\\varphi)$ [$10^{52}$ erg s$^{-1}$]',\n",
    "           ylim=(0, 1.1*np.max(L_nue_bar).value),\n",
    "           title=r'$L_{\\bar{\\nu}_e}(t,\\theta,\\varphi)$')\n",
    "    ax.grid(ls=':')\n",
    "\n",
    "    ax = axes[1]\n",
    "    ax.plot(cmodel.time, dL_over_L, color='gray', alpha=0.1)\n",
    "    ax.plot(cmodel.time, np.zeros(cmodel.time.shape), ls='--', color='r')\n",
    "    ax.set(xlabel='time [s]',\n",
    "           ylabel=r'$(L_\\bar{\\nu} - \\bar{L}_\\bar{\\nu})/\\bar{L}_\\bar{\\nu}$',\n",
    "           ylim=(-0.3, 0.3))\n",
    "    ax.grid(ls=':')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de60253b",
   "metadata": {},
   "source": [
    "#### Flavor X (Stand-in for Mu and Tau Neutrinos)\n",
    "\n",
    "Note that in this model, the mu and tau antineutrino flux is identical to the neutrino flux so we won't bother plotting it here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee70cfa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "if healpy_available:\n",
    "    Lavg = np.average(L_nux, axis=1)\n",
    "    dL_over_L = (L_nux - Lavg[:,np.newaxis]) / Lavg[:,np.newaxis]\n",
    "\n",
    "    fig, axes = plt.subplots(2,1, figsize=(10,8),\n",
    "                        gridspec_kw = {'height_ratios':[3,1], 'hspace':0},\n",
    "                        sharex=True, tight_layout=True)\n",
    "\n",
    "    ax = axes[0]\n",
    "    ax.plot(cmodel.time, L_nux, color='gray', alpha=0.05)\n",
    "    ax.plot(cmodel.time, Lavg, color='r', ls='--')\n",
    "    ax.set(xlim=cmodel.time[0::len(cmodel.time)-1].value,\n",
    "            ylabel=r'$L_{\\nu_X}(t,\\theta,\\varphi)$ [$10^{52}$ erg s$^{-1}$]',\n",
    "            ylim=(0, 1.1*np.max(L_nux).value),\n",
    "            title=r'$L_{\\nu_X}(t,\\theta,\\varphi)$')\n",
    "    ax.grid(ls=':')\n",
    "\n",
    "    ax = axes[1]\n",
    "    ax.plot(cmodel.time, dL_over_L, color='gray', alpha=0.1)\n",
    "    ax.plot(cmodel.time, np.zeros(cmodel.time.shape), ls='--', color='r')\n",
    "    ax.set(xlabel='time [s]',\n",
    "            ylabel=r'$(L_\\nu - \\bar{L}_\\nu)/\\bar{L}_\\nu$',\n",
    "            ylim=(-0.3, 0.3))\n",
    "    ax.grid(ls=':');"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.5 ('snews')",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "e2528887d751495e023d57d695389d9a04f4c4d2e5866aaf6dc03a1ed45c573e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}