{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3fab591f",
   "metadata": {},
   "source": [
    "# Yoshida 2016 Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83d2d170-a5b3-44d9-ad05-762448b08c80",
   "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": "f2ad86ec-ff7e-4ef6-818e-f244ac8112f4",
   "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={'E':'C0','MU':'C1','TAU':'C2'}[f.lepton],\n",
    "                  ls='-' if f.is_neutrino else ':',\n",
    "                  label=f.to_tex()) for f in Flavor}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f1b6439-6702-417a-a9b6-0ca0d5798f47",
   "metadata": {},
   "source": [
    "## Initialize the model and calculate the flux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffff2732-a87a-401c-b85d-a0f0d5598657",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See what progenitors are available …\n",
    "print(f\"Available progenitors:\\n{presn.Yoshida_2016.param}\")\n",
    "\n",
    "# … and select one of them to initialise the model\n",
    "model = presn.Yoshida_2016(progenitor_mass=15*u.Msun)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d8e7b15-0a7e-49f3-b6bd-9463e8b371a5",
   "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": "efb466d7-1b4c-4949-a983-5134fba3f35b",
   "metadata": {},
   "outputs": [],
   "source": [
    "flux = model.get_flux(t, E,\n",
    "                      distance=distance,\n",
    "                      flavor_xform=NoTransformation())\n",
    "flux"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f06fc5ed-e39f-4f91-a8c6-09a671a7f6d0",
   "metadata": {},
   "source": [
    "## Plotting the integral neutrino fluence and rates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34cd5637-201f-493c-9aa8-10a653d4158b",
   "metadata": {},
   "source": [
    "### Integral neutrino fluence vs. Energy for the last hour before collapse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f8f5b53-257d-43b7-b933-415738e8b4a1",
   "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": "8347961e",
   "metadata": {},
   "source": [
    "### Integral neutrino rate vs. time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dba0da41-6061-4c7a-b1fc-4b73f29ea960",
   "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": "2dfec9e5-2819-4993-b352-819224b65279",
   "metadata": {},
   "source": [
    "## Plot energy spectrum with neutrino oscillations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94e4605b-fcee-4d3a-8344-caa4100e3ee1",
   "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(), \n",
    "          AdiabaticMSW(MixingParameters(\"NORMAL\")),\n",
    "          AdiabaticMSW(MixingParameters(\"INVERTED\"))\n",
    "         ]\n",
    "labels = ['No transformation', 'AdiabaticMSW NMO', 'AdiabaticMSW IMO']\n",
    "xform_styles = ['-','--',':']\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(3,2,figsize=(8,12), sharex=True, sharey=True)\n",
    "\n",
    "for flv, ax in zip(Flavor,axes.flatten()):\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",
    "        style = styles[flv]\n",
    "        style.update(ls=line_style,\n",
    "                     label=f'{Flavor(flv).to_tex()} {label}'\n",
    "                    )\n",
    "        plt.plot(x,y, **style)\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",
    "    for ax in axes[:,0]:\n",
    "        ax.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": "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}