{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Flavor Transformations\n",
    "\n",
    "This notebook produces the figures in the SNEWPY paper showing the effect of the flavor transformation prescriptions upon the nakazato-shen-z0.004-t_rev100ms-s20.0 model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from snewpy.neutrino import MassHierarchy, MixingParameters, ThreeFlavorMixingParameters, FourFlavorMixingParameters\n",
    "import snewpy.flavor_transformation as xforms\n",
    "from snewpy.flavor import ThreeFlavor, FourFlavor\n",
    "\n",
    "from snewpy.models.ccsn import Nakazato_2013\n",
    "\n",
    "from astropy import units as u\n",
    "from astropy import constants as c\n",
    "from astropy.coordinates import SkyCoord, EarthLocation, AltAz\n",
    "from astropy.time import Time                                       \n",
    "\n",
    "mpl.rc('font', size=18)\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Nakazato_2013(progenitor_mass=20*u.solMass, revival_time=100*u.ms, metallicity=0.004, eos='shen')\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformed and untransformed energy integrated flux at Earth\n",
    "\n",
    "Compute and plot the flux at Earth for a SN at 10 kpc with no flavor transformation, and with the chosen flavor transformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_total_flux(model, xform_nmo, xform_imo):\n",
    "    \"\"\"Plot initial and oscillated neutrino luminosities.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model : SupernovaModel\n",
    "        An input model from a CCSN simulation.\n",
    "    flav_xform : FlavorTransformation\n",
    "        A FlavorTransformation subclass; used to create an instance.\n",
    "    \"\"\"\n",
    "    \n",
    "    energies = np.linspace(1,60,119) * u.MeV\n",
    "    d = 10*u.kpc # distance to SN\n",
    "        \n",
    "    times = model.get_time()\n",
    "    burst_epoch = times <= 0.1*u.s\n",
    "    accretion_epoch = (times > 0.1*u.s) & (times <= 0.5*u.s)\n",
    "    cooling_epoch = (times > 0.5*u.s) & (times <= 10*u.s)\n",
    "    \n",
    "    ispec = model.get_flux(times, energies, d)\n",
    "    ospec_nmo = model.get_flux(times, energies, d, xform_nmo)\n",
    "    ospec_imo = model.get_flux(times, energies, d, xform_imo)\n",
    "    \n",
    "    ilum = ispec.integrate_or_sum('energy') \n",
    "    olum_nmo = ospec_nmo.integrate_or_sum('energy') \n",
    "    olum_imo = ospec_imo.integrate_or_sum('energy') \n",
    "    \n",
    "    # make the figures \n",
    "    fig, axes = plt.subplots(3,3, figsize=(20,12), tight_layout=True)\n",
    "    \n",
    "    smax = [0.,0.,0.]\n",
    "    titles = ['Untransformed', 'Transformed (NMO)', 'Transformed (IMO)']\n",
    "    for i, spec in enumerate([ilum, olum_nmo, olum_imo]):   \n",
    "        for j, phase in enumerate([burst_epoch, accretion_epoch, cooling_epoch]):\n",
    "            ax = axes[i,j]\n",
    "            timeunits = 'ms' if j==0 else 's'\n",
    "            x_values = times[phase].to(timeunits)\n",
    "            for flavor in ThreeFlavor:\n",
    "                y_values = spec.array.squeeze()[flavor][phase].to('1/(cm2 s)')\n",
    "                if i == 0:\n",
    "                    smax[j] = np.maximum(smax[j], 1.1*np.max(y_values))\n",
    "                    \n",
    "                ax.plot(x_values, y_values, label=flavor.to_tex(), lw=3,\n",
    "                        color={'E':'C0','MU':'C1','TAU':'C2'}[flavor.lepton],\n",
    "                        ls='-' if flavor.is_neutrino else ':')\n",
    "            \n",
    "            ax.set(xlim=(x_values[0].value, x_values[-1].value),\n",
    "                   ylim=(0, smax[j].value))\n",
    "            \n",
    "            if j==0:\n",
    "                ax.set(ylabel=f'flux [{y_values.unit._repr_latex_()}]')\n",
    "                ax.legend(loc='upper right', ncol=1, fontsize=18)\n",
    "            if j==1:\n",
    "                ax.set(title=titles[i])\n",
    "            if i < 2:\n",
    "                ax.set(xticklabels=[])\n",
    "            else:\n",
    "                ax.set(xlabel='time [{}]'.format(timeunits))\n",
    "            \n",
    "            ax.grid(ls=':')\n",
    "\n",
    "      \n",
    "    return fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Untransformed and Transformed Spectra at Earth\n",
    "\n",
    "Compute and plot the spectra at Earth for a SN at 10 kpc with no flavor transformation, and with the chosen flavor transformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_spectra(model, xform_nmo, xform_imo, t):\n",
    "    \"\"\"Plot initial and oscillated neutrino luminosities.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model : SupernovaModel\n",
    "        An input model from a CCSN simulation.\n",
    "    flav_xform : FlavorTransformation\n",
    "        A FlavorTransformation subclass; used to create an instance.\n",
    "    t : astropy.Quantity\n",
    "        Time to compute the spectrum.\n",
    "    \"\"\"\n",
    "\n",
    "    energies = np.linspace(1,60,119) * u.MeV\n",
    "    d = 10*u.kpc # distance to SN\n",
    "\n",
    "    #get the spectra\n",
    "    ispec = model.get_flux(t, energies, d)\n",
    "    ospec_nmo = model.get_flux(t, energies, d, xform_nmo)\n",
    "    ospec_imo = model.get_flux(t, energies, d, xform_imo)\n",
    "    \n",
    "    fig, axes = plt.subplots(2,2, figsize=(18,15), sharex=True, sharey=True, tight_layout=True)\n",
    "\n",
    "    for idx_mixing, ospec in enumerate([ospec_nmo, ospec_imo]):\n",
    "        for idx_nu, flavors in enumerate([ThreeFlavor['NU'], ThreeFlavor['NU_BAR']]):\n",
    "            for idx_spec, spec in enumerate([ispec, ospec_nmo]):\n",
    "                spec_val = spec.array.squeeze()\n",
    "                spec_val = spec_val.to('1e16/(erg s cm2)')\n",
    "                for flavor in flavors:\n",
    "                    ax = axes[idx_nu][idx_mixing]\n",
    "                    ax.plot(energies, spec_val[flavor],\n",
    "                            label = ('Untransformed','Transformed')[idx_spec]+' '+flavor.to_tex(),\n",
    "                            color={'E':'C0','MU':'C1','TAU':'C2'}[flavor.lepton],\n",
    "                            ls=('-',':')[idx_spec],\n",
    "                            lw=2,  alpha=0.7\n",
    "                           )\n",
    "                \n",
    "        ax.set(title='{Nu}eutrinos in the {mixing}: $t = ${t:.1f}'.format(Nu=['N','Antin'][idx_nu],\n",
    "                                                                          mixing=['NMO','IMO'][idx_mixing],\n",
    "                                                                          t=t)\n",
    "              )\n",
    "    for ax in axes.flatten():\n",
    "        ax.grid()\n",
    "        ax.legend(loc='upper right', ncol=2, fontsize=16)\n",
    "    for ax in axes[:,0]:\n",
    "        ax.set(ylabel=f'flux {spec_val.unit._repr_latex_()}')#[$10^{16}$ erg$^{-1}$ cm$^{-2}$ s$^{-1}$]')\n",
    "    \n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixpars_nmo = MixingParameters('NORMAL')\n",
    "mixpars_imo = MixingParameters('INVERTED')\n",
    "\n",
    "xform_nmo = xforms.AdiabaticMSW(mixpars_nmo)\n",
    "xform_imo = xforms.AdiabaticMSW(mixpars_imo) \n",
    "\n",
    "fig = plot_total_flux(model, xform_nmo, xform_imo)\n",
    "fig.show()\n",
    "# fig.savefig('flux_adiabaticmsw.pdf')\n",
    "fig = plot_spectra(model, xform_nmo, xform_imo, 100*u.ms)\n",
    "fig.show()\n",
    "# fig.savefig('spectra_adiabaticmsw.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set 1 degree as the mixing angle theta14\n",
    "mixpars_nmo = FourFlavorMixingParameters(**MixingParameters('NORMAL'), theta14=1*u.deg)\n",
    "xf_nmo = xforms.AdiabaticMSWes(mixpars_nmo)\n",
    "\n",
    "mixpars_imo = FourFlavorMixingParameters(**MixingParameters('INVERTED'), theta14=1*u.deg)\n",
    "xf_imo = xforms.AdiabaticMSWes(mixpars_imo)\n",
    "\n",
    "fig = plot_total_flux(model, xf_nmo, xf_imo)\n",
    "fig.show()\n",
    "# fig.savefig('flux_adiabaticmswes.pdf')\n",
    "fig = plot_spectra(model, xf_nmo, xf_imo, 100*u.ms)\n",
    "fig.show()\n",
    "# fig.savefig('spectra_adiabaticmswes.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}