{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5299d14b",
   "metadata": {},
   "source": [
    "# Fornax 2021 Models\n",
    "\n",
    "Neutrino spectra from the long-time run 2D (axisymmetric) models produced by Burrows and Vartanyan, \n",
    "[Nature **589**:29-39, 2021](https://www.nature.com/articles/s41586-020-03059-w).\n",
    "\n",
    "Data taken from the HDF5 files available for download at [the Princeton group website](https://www.astro.princeton.edu/~burrows/nu-emissions.2d/data/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d3d8a99",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snewpy.neutrino import Flavor\n",
    "from snewpy.models.ccsn import Fornax_2021\n",
    "\n",
    "from astropy import units as u\n",
    "from glob import glob\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "mpl.rc('font', size=16)\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4799f82",
   "metadata": {},
   "source": [
    "## Initialize Models\n",
    "\n",
    "To start, let’s see what progenitors are available for the `Fornax_2021` model. We can use the `param` property to view all physics parameters and their possible values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd8215d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "Fornax_2021.param"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96bfc909",
   "metadata": {},
   "source": [
    "We’ll initialise some of these progenitors and plot the luminosity of different neutrino flavors for two of them. (Note that the `Fornax_2021` simulations didn’t distinguish between $\\nu_x$ and $\\bar{\\nu}_x$, so both have the same luminosity.) 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": "c10c3261",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = {}\n",
    "for m in Fornax_2021.param['progenitor_mass'].value[::2]:\n",
    "    # Initialise every second progenitor\n",
    "    models[m] = Fornax_2021(progenitor_mass=m*u.solMass)\n",
    "\n",
    "models[12]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56eb4178",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharex=True, sharey=True, tight_layout=True)\n",
    "\n",
    "for i, model in enumerate([models[12], models[20]]):\n",
    "    ax = axes[i]\n",
    "    for flavor in Flavor:\n",
    "        ax.plot(model.time, model.luminosity[flavor]/1e51,  # Report luminosity in units foe/s\n",
    "                label=flavor.to_tex(),\n",
    "                color='C0' if flavor.is_electron else 'C1',\n",
    "                ls='-' if flavor.is_neutrino else ':',\n",
    "                lw=2)\n",
    "    ax.set(xlim=(-0.05, 1.0),\n",
    "           xlabel=r'$t-t_{\\rm bounce}$ [s]',\n",
    "           title=r'{} $M_\\odot$'.format(model.metadata['Progenitor mass'].value))\n",
    "    ax.grid()\n",
    "    ax.legend(loc='upper right', ncol=2, fontsize=18)\n",
    "\n",
    "axes[0].set(ylabel=r'luminosity [foe s$^{-1}$]');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a8a47a3",
   "metadata": {},
   "source": [
    "## Spectra of All Flavors vs. Time for the $12M_{\\odot}$ Model\n",
    "\n",
    "### Use Default Linear Interpolation in Flux Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2487b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models[12]  # Use the 12 solar mass model\n",
    "\n",
    "times = np.arange(-0.2, 3.8, 0.2) * u.s\n",
    "E = np.arange(0, 101, 1) * u.MeV\n",
    "\n",
    "fig, axes = plt.subplots(5,4, figsize=(15,12), sharex=True, sharey=True, tight_layout=True)\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "\n",
    "spectra = model.get_initial_spectra(times, E)\n",
    "\n",
    "for i, ax in enumerate(axes.flatten()):\n",
    "    for line, flavor in zip(linestyles, Flavor):\n",
    "        ax.plot(E, spectra[flavor][i], lw=3, ls=line, label=flavor.to_tex())\n",
    "    ax.set(xlim=(0,100))\n",
    "    ax.set_title('$t$ = {:g}'.format(times[i]), fontsize=16)\n",
    "    ax.legend(loc='upper right', ncol=2, fontsize=12)\n",
    "\n",
    "fig.text(0.5, 0., 'energy [MeV]', ha='center')\n",
    "fig.text(0., 0.5, f'flux [{spectra[Flavor.NU_E].unit}]', va='center', rotation='vertical');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6bf5841",
   "metadata": {},
   "source": [
    "### Use Nearest-Bin \"Interpolation\" in Flux Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c797f99",
   "metadata": {},
   "outputs": [],
   "source": [
    "times = np.arange(-0.2, 3.8, 0.2) * u.s\n",
    "E = np.arange(0, 101, 1) * u.MeV\n",
    "\n",
    "fig, axes = plt.subplots(5,4, figsize=(15,12), sharex=True, sharey=True, tight_layout=True)\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "\n",
    "spectra = model.get_initial_spectra(times, E, interpolation='nearest')\n",
    "\n",
    "for i, ax in enumerate(axes.flatten()):\n",
    "    for line, flavor in zip(linestyles, Flavor):\n",
    "        ax.plot(E, spectra[flavor][i], lw=3, ls=line, label=flavor.to_tex())\n",
    "    ax.set(xlim=(0,100))\n",
    "    ax.set_title('$t$ = {:g}'.format(times[i]), fontsize=16)\n",
    "    ax.legend(loc='upper right', ncol=2, fontsize=12)\n",
    "\n",
    "fig.text(0.5, 0., 'energy [MeV]', ha='center')\n",
    "fig.text(0., 0.5, f'flux [{spectra[Flavor.NU_E].unit}]', va='center', rotation='vertical');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad9ecc5f",
   "metadata": {},
   "source": [
    "## Progenitor Mass Dependence\n",
    "\n",
    "### Luminosity vs. Time for a Selected List of Progenitor Masses\n",
    "\n",
    "Plot $L_{\\nu_{e}}(t)$ for a selection of progenitor masses to observe the dependence of the emission on mass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07270503",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3,1, figsize=(10,13), sharex=True, sharey=True,\n",
    "                         gridspec_kw = {'hspace':0.02})\n",
    "\n",
    "colors0 = mpl.cm.viridis(np.linspace(0.1,0.9, len(models)))\n",
    "colors1 = mpl.cm.inferno(np.linspace(0.1,0.9, len(models)))\n",
    "colors2 = mpl.cm.cividis(np.linspace(0.1,0.9, len(models)))\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "\n",
    "for i, model in enumerate(models.values()):\n",
    "    ax = axes[0]\n",
    "    flavor = Flavor.NU_E\n",
    "    ax.plot(model.time, model.luminosity[flavor], lw=2, color=colors0[i], ls=linestyles[i%4],\n",
    "            label='${0.value:g}$ {0.unit:latex}'.format(model.progenitor_mass))\n",
    "    ax.set(xscale='log',\n",
    "           xlim=(1e-3, 4),\n",
    "           yscale='log',\n",
    "           ylim=(0.4e52, 9e53),\n",
    "           ylabel=r'$L_{\\nu_e}(t)$ [erg s$^{-1}$]')\n",
    "    ax.grid(ls=':', which='both')\n",
    "    ax.legend(ncol=3, fontsize=12, title=r'$\\nu_e$');\n",
    "    \n",
    "    ax = axes[1]\n",
    "    flavor = Flavor.NU_E_BAR\n",
    "    ax.plot(model.time, model.luminosity[flavor], lw=2, color=colors1[i], ls=linestyles[i%4],\n",
    "        label='${0.value:g}$ {0.unit:latex}'.format(model.progenitor_mass))\n",
    "    ax.set(ylabel=r'$L_{\\bar{\\nu}_e}(t)$ [erg s$^{-1}$]')\n",
    "    ax.grid(ls=':', which='both')\n",
    "    ax.legend(ncol=3, fontsize=12, title=r'$\\bar{\\nu}_e$');\n",
    "    \n",
    "    ax = axes[2]\n",
    "    flavor = Flavor.NU_X\n",
    "    ax.plot(model.time, model.luminosity[flavor], lw=2, color=colors2[i], ls=linestyles[i%4],\n",
    "        label='${0.value:g}$ {0.unit:latex}'.format(model.progenitor_mass))\n",
    "    ax.set(xlabel='time [s]',\n",
    "           ylabel=r'$L_{\\nu_X}(t)$ [erg s$^{-1}$]')\n",
    "    ax.grid(ls=':', which='both')\n",
    "    ax.legend(ncol=3, fontsize=12, title=r'$\\nu_X$');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6cdab52",
   "metadata": {},
   "source": [
    "### Progenitor Dependence of Spectra at 70 ms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c66de569",
   "metadata": {},
   "source": [
    "#### Use Default Linear Interpolation in Flux Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "441acedf",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 70*u.ms\n",
    "E = np.arange(0, 101, 1) * u.MeV\n",
    "\n",
    "fig, axes = plt.subplots(2,4, figsize=(16,6), sharex=True, sharey=True, tight_layout=True)\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "\n",
    "for model, ax in zip(models.values(), axes.flatten()):\n",
    "    spectra = model.get_initial_spectra(t, E)\n",
    "    for line, flavor in zip(linestyles, Flavor):\n",
    "        ax.plot(E, spectra[flavor][0], lw=3, ls=line, label=flavor.to_tex())\n",
    "    ax.set(xlim=(0,100))\n",
    "    ax.set_title('${0.value:g}$ {0.unit:latex}'.format(model.progenitor_mass))\n",
    "    ax.legend(loc='upper right', ncol=2, fontsize=12)\n",
    "    ax.grid(ls=':')\n",
    "\n",
    "fig.text(0.5, 0., 'energy [MeV]', ha='center')\n",
    "fig.text(0., 0.5, f'flux [{spectra[Flavor.NU_E].unit}]', va='center', rotation='vertical');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1870b66",
   "metadata": {},
   "source": [
    "#### Use Nearest-Bin \"Interpolation\" in Flux Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e41b55f",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 70*u.ms\n",
    "E = np.arange(0, 101, 1) * u.MeV\n",
    "\n",
    "fig, axes = plt.subplots(2,4, figsize=(16,6), sharex=True, sharey=True, tight_layout=True)\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "\n",
    "for model, ax in zip(models.values(), axes.flatten()):\n",
    "    spectra = model.get_initial_spectra(t, E, interpolation='nearest')\n",
    "    for line, flavor in zip(linestyles, Flavor):\n",
    "        ax.plot(E, spectra[flavor][0], lw=3, ls=line, label=flavor.to_tex())\n",
    "    ax.set(xlim=(0,100))\n",
    "    ax.set_title('${0.value:g}$ {0.unit:latex}'.format(model.progenitor_mass))\n",
    "    ax.legend(loc='upper right', ncol=2, fontsize=12)\n",
    "    ax.grid(ls=':')\n",
    "\n",
    "fig.text(0.5, 0., 'energy [MeV]', ha='center')\n",
    "fig.text(0., 0.5, f'flux [{spectra[Flavor.NU_E].unit}]', va='center', rotation='vertical');"
   ]
  }
 ],
 "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
}