{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "89257e22",
   "metadata": {},
   "source": [
    "# Kato 2017 Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3434b585",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from astropy import units as u\n",
    "\n",
    "# import the snewpy modules\n",
    "from snewpy.models import presn\n",
    "from snewpy.flavor_transformation import NoTransformation, AdiabaticMSW\n",
    "from snewpy.neutrino import Flavor, MixingParameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b73c22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the plot parameters\n",
    "plt.rc('grid', ls=':')\n",
    "plt.rc('axes', grid=True)\n",
    "plt.rc('legend', fontsize=12, loc='upper right')\n",
    "\n",
    "# define drawing styles for each flavor\n",
    "styles = {f: dict(color='C0' if f.is_electron else 'C1',\n",
    "                  ls='-' if f.is_neutrino else ':',\n",
    "                  label=f.to_tex()) for f in Flavor}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "212af44c-6498-4a2c-b1e2-4bc98a1e4645",
   "metadata": {},
   "source": [
    "## Initialize the model and calculate the flux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e56a06b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See what progenitors are available …\n",
    "print(f\"Available progenitors:\\n{presn.Kato_2017.param}\")\n",
    "\n",
    "# … and select one of them to initialise the model\n",
    "model = presn.Kato_2017(progenitor_mass=15*u.Msun)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b629c538-005d-43c5-b244-bafc9dcb78dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Energy array and time to compute spectra.\n",
    "# Note that any convenient units can be used and the calculation will remain internally consistent.\n",
    "E = np.linspace(0, 25, 201) * u.MeV\n",
    "t = np.geomspace(-2*u.day, -1*u.s, 101)\n",
    "distance = 1*u.kpc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d0a2218-5955-4d20-89ee-f6de8181bb4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "flux = model.get_flux(t, E,\n",
    "                      distance=distance,\n",
    "                      flavor_xform=NoTransformation())\n",
    "flux"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88e8ce95-6705-4282-82c6-ca797234ca2f",
   "metadata": {},
   "source": [
    "## Plotting the integral neutrino fluence and rates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "522e70c0-e6a3-4d5f-9a80-4f970c23d69b",
   "metadata": {},
   "source": [
    "### Integral neutrino fluence vs. Energy for the last hour before collapse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2be68b2-bb67-470e-bc01-74e4e48df195",
   "metadata": {},
   "outputs": [],
   "source": [
    "#integrate the flux over the last hour before the collapse\n",
    "fluence = flux.integrate('time', limits=[-1, 0]<<u.hour)\n",
    "#get the relevant arrays for plotting\n",
    "x = fluence.energy\n",
    "y = fluence.array.squeeze()\n",
    "#plot each flavor\n",
    "for flv in fluence.flavor:\n",
    "    plt.plot(x,y[flv], **styles[flv])\n",
    "#add the legend\n",
    "plt.legend(ncol=2, loc='lower center')\n",
    "#adjust the scales\n",
    "plt.autoscale(tight=True,axis='x')\n",
    "plt.yscale('log')\n",
    "plt.ylim(1)\n",
    "#define the labels\n",
    "xlabel = 'Neutrino energy'\n",
    "ylabel = 'Neutrino fluence'\n",
    "plt.xlabel(f'{xlabel} [{x.unit._repr_latex_()}]')\n",
    "plt.ylabel(f'{ylabel} [{y.unit._repr_latex_()}]')\n",
    "#add the plot title\n",
    "plt.title('Neutrinos in the last 1 hour before collapse')\n",
    "#add the model parameters\n",
    "plt.annotate(str(model) + f'\\ndistance : {distance}',\n",
    "             xy=(0.98,0.98),\n",
    "             xycoords='axes fraction',\n",
    "             va='top', ha='right'\n",
    "            )\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ffa39ae",
   "metadata": {},
   "source": [
    "### Integral neutrino rate vs. time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ce19434-8e82-4876-8bd0-ac695df339fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#integrate the flux over all energy bins\n",
    "rate = flux.integrate('energy')\n",
    "#get the relevant arrays for plotting\n",
    "x = rate.time\n",
    "y = rate.array.squeeze()\n",
    "#we can also convert to a more convenient unit\n",
    "y = y.to('1/(hour*m**2)')\n",
    "#plot each flavor\n",
    "for flv in rate.flavor:\n",
    "    plt.plot(x,y[flv], **styles[flv])\n",
    "#add the legend\n",
    "plt.legend(ncol=2, loc='lower center')\n",
    "#adjust the scales\n",
    "plt.autoscale(tight=True,axis='x')\n",
    "plt.yscale('log')\n",
    "#plt.ylim(1)\n",
    "#define the labels\n",
    "xlabel = 'Time before collapse'\n",
    "ylabel = 'Neutrino rate'\n",
    "plt.xlabel(f'{xlabel} [{x.unit._repr_latex_()}]')\n",
    "plt.ylabel(f'{ylabel} [{y.unit._repr_latex_()}]')\n",
    "#add the plot title\n",
    "plt.title('Integral neutrino rate')\n",
    "#add the model parameters\n",
    "plt.annotate(str(model) + f'\\ndistance : {distance}',\n",
    "             xy=(0.02,0.98),\n",
    "             xycoords='axes fraction',\n",
    "             va='top', ha='left'\n",
    "            )\n",
    "plt.xscale('symlog')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9769409-4520-4e7e-a6ea-79941d260b76",
   "metadata": {},
   "source": [
    "## Plot energy spectrum with neutrino oscillations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41d18d10-5cc9-4e24-bd7c-3fe991c17b97",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Energy array and time to compute spectra.\n",
    "# Note that any convenient units can be used and the calculation will remain internally consistent.\n",
    "E = np.linspace(0,25,201) * u.MeV\n",
    "t = -1*u.s\n",
    "\n",
    "xforms = [NoTransformation(), AdiabaticMSW(MixingParameters()), AdiabaticMSW(MixingParameters('INVERTED'))]\n",
    "labels = ['No transformation', 'AdiabaticMSW NMO', 'AdiabaticMSW IMO']\n",
    "xform_styles = ['-','--',':']\n",
    "flavor_colors = ['C0','C1','C2','C3']\n",
    "\n",
    "#calculate the flux for different transformations\n",
    "fluxes  = {xform: model.get_flux(t, E, distance, flavor_xform=xform) for xform in xforms}\n",
    "#make a big figure\n",
    "fig, axes = plt.subplots(1,4,figsize=(16,4), sharex=True, sharey=True)\n",
    "\n",
    "for flv, ax in zip(Flavor,axes):\n",
    "    plt.sca(ax)\n",
    "    for xform, line_style, label in zip(xforms,xform_styles, labels):\n",
    "        flux = fluxes[xform]\n",
    "        #get the relevant arrays for plotting\n",
    "        x = flux.energy\n",
    "        y = flux.array[flv].squeeze()\n",
    "        #plot each flavor\n",
    "        plt.plot(x,y, \n",
    "                 ls=line_style, \n",
    "                 color=flavor_colors[flv], \n",
    "                 label=f'{Flavor(flv).to_tex()} {label}'\n",
    "                )\n",
    "        \n",
    "    #add the legend\n",
    "    plt.legend(ncol=1, loc='lower left', fontsize=10)\n",
    "    \n",
    "    #define the labels\n",
    "    xlabel = 'Neutrino energy'\n",
    "    ylabel = 'Neutrino flux'\n",
    "    plt.xlabel(f'{xlabel} [{x.unit._repr_latex_()}]')\n",
    "    #add the plot title\n",
    "\n",
    "    #add the model parameters\n",
    "    plt.annotate(str(model) + f'\\ndistance : {distance}',\n",
    "                 xy=(0.98,0.98),\n",
    "                 xycoords='axes fraction',\n",
    "                 va='top', ha='right'\n",
    "                )\n",
    "    axes[0].set_ylabel(f'{ylabel} [{y.unit._repr_latex_()}]')\n",
    "    #adjust the scales\n",
    "    plt.autoscale(tight=True,axis='x')\n",
    "    plt.autoscale(tight=False,axis='y')\n",
    "    plt.yscale('log')\n",
    "    plt.ylim(1)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "snewpy",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}