{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e8dbfcc",
   "metadata": {},
   "source": [
    "# Fornax 2022 Models\n",
    "\n",
    "Neutrino spectra from the 2D long-duraction models produced for 100 progenitors (1/3 form black holes).\n",
    "\n",
    "Data take from the [Princeton group webpage](https://www.astro.princeton.edu/~burrows/nu-emissions.2d.large/) and converted to HDF5 format for use in SNEWPY."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3dcba05",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snewpy.neutrino import Flavor\n",
    "from snewpy.models.ccsn import Fornax_2022\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69e2e682",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rc('font', size=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2c2ac5d",
   "metadata": {},
   "source": [
    "## Initialize Models\n",
    "\n",
    "To start, see what progenitors are avaialble for the `Fornax_2022` model. Use the `param` property to view all physics parameters and their possible values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e311770",
   "metadata": {},
   "outputs": [],
   "source": [
    "Fornax_2022.param"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e5f0052",
   "metadata": {},
   "source": [
    "Initialize some progenitors and plot the luminosity of different neutrino flavors for two of them. Note that the `Fornax_2022` set of models do not distinguish between $\\nu_x$ and $\\bar{\\nu}_x$ so both have the same luminosity. If this is the first time you are using a progenitor model, `snewpy` will download the data files for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0750b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = {}\n",
    "for m in Fornax_2022.param['progenitor_mass'][::19]:\n",
    "    # Initialize every 20th progenitor.\n",
    "    print(m)\n",
    "    models[m] = Fornax_2022(progenitor_mass=m)\n",
    "    \n",
    "models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d7c2a2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "models[23.43*u.solMass].metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ded8b91",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3, 2, figsize=(12, 15), sharex=True, sharey=True, tight_layout=True)\n",
    "axes = axes.flatten()\n",
    "\n",
    "for i, (key, model) in enumerate(models.items()):\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",
    "        \n",
    "    modtitle = rf\"{model.metadata['Progenitor mass'].value} $M_\\odot$\"\n",
    "    if model.metadata['Black hole']:\n",
    "        modtitle += ' (BH)'\n",
    "        \n",
    "    ax.set(xlim=(-0.05, 1.5),\n",
    "           xlabel=r'$t-t_{\\rm bounce}$ [s]',\n",
    "           title=modtitle)\n",
    "    ax.grid()\n",
    "    ax.legend(loc='upper right', ncol=2, fontsize=18)\n",
    "\n",
    "axes[0].set(ylabel=r'luminosity [foe s$^{-1}$]');\n",
    "# axes[5].set_axis_off();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78d70089",
   "metadata": {},
   "source": [
    "## Spectra of All Flavors vs. Time for the 19.02 $M_\\odot$ Model\n",
    "\n",
    "### Use Default Linear Interpolation in Flux Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a301327",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models[23.43*u.solMass]\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",
    "model.interpolation='linear'\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].array.squeeze(), 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": "27ef0892",
   "metadata": {},
   "source": [
    "### Use Nearest-Bin \"Interpolation\" in Flux Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cea3d5b",
   "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",
    "model.interpolation='nearest'\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].array.squeeze(), 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": "c9dc1934",
   "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": "ab022bc4",
   "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}{1}'.format(model.progenitor_mass, ' (BH)' if 'bh' in model.progenitor else ''))\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}{1}'.format(model.progenitor_mass, ' (BH)' if 'bh' in model.progenitor else ''))\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_MU\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_\\mu}(t)$ [erg s$^{-1}$]')\n",
    "    ax.grid(ls=':', which='both')\n",
    "    ax.legend(ncol=3, fontsize=12, title=r'$\\nu_\\mu$');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89a17c23",
   "metadata": {},
   "source": [
    "## Progenitor Dependence of Spectra at 70 ms\n",
    "\n",
    "### Use Default Linear Interpolation in Flux Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f0eba54",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 70*u.ms\n",
    "E = np.arange(0, 101, 1) * u.MeV\n",
    "\n",
    "fig, axes = plt.subplots(2,3, figsize=(12,6), sharex=True, sharey=True, tight_layout=True)\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "\n",
    "for model, ax in zip(models.values(), axes.flatten()):\n",
    "    model.interpolation='linear'\n",
    "    spectra = model.get_initial_spectra(t, E)\n",
    "    for line, flavor in zip(linestyles, Flavor):\n",
    "        ax.plot(E, spectra[flavor,0].array.squeeze(), lw=3, ls=line, label=flavor.to_tex())\n",
    "    ax.set(xlim=(0,100))\n",
    "    ax.set_title('${0.value:g}$ {0.unit:latex}{1}'.format(model.progenitor_mass, ' (BH)' if 'bh' in model.progenitor else ''))\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": "456450c5",
   "metadata": {},
   "source": [
    "### Use Nearest-Bin \"Interpolation\" in Flux Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52dca852",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 70*u.ms\n",
    "E = np.arange(0, 101, 1) * u.MeV\n",
    "\n",
    "fig, axes = plt.subplots(2,3, figsize=(12,6), sharex=True, sharey=True, tight_layout=True)\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "for model, ax in zip(models.values(), axes.flatten()):\n",
    "    model.interpolation='nearest'\n",
    "    spectra = model.get_initial_spectra(t, E)\n",
    "    for line, flavor in zip(linestyles, Flavor):\n",
    "        ax.plot(E, spectra[flavor,0].array.squeeze(), lw=3, ls=line, label=flavor.to_tex())\n",
    "    ax.set(xlim=(0,100))\n",
    "    ax.set_title('${0.value:g}$ {0.unit:latex}{1}'.format(model.progenitor_mass, ' (BH)' if 'bh' in model.progenitor else ''))\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 (ipykernel)",
   "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}