{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a Supernova Model with an Analytic Spectrum\n",
    "\n",
    "This notebook demonstrates how to use the `Analytic3Species` class from `snewpy.models` to create an analytic supernova model.\n",
    "The neutrino spectrum of this model follows a Gamma distribution (see [arXiv:1211.3920](https://arxiv.org/abs/1211.3920)) with user-selected parameters.\n",
    "\n",
    "In this notebook, we first create a model file, then visualize the spectral parameters and finally use SNOwGLoBES to determine the number of events expected in a detector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from astropy.table import Table\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import numpy as np\n",
    "\n",
    "from snewpy import snowglobes, model_path\n",
    "from snewpy.neutrino import Flavor\n",
    "from snewpy.models.ccsn import Analytic3Species\n",
    "\n",
    "mpl.rc('font', size=14)\n",
    "\n",
    "SNOwGLoBES_path = None  # change to SNOwGLoBES directory if using a custom detector configuration\n",
    "\n",
    "model_folder = f\"{model_path}/AnalyticFluence/\"\n",
    "os.makedirs(model_folder, exist_ok=True)\n",
    "print(f\"Using folder `{model_folder}`.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a SN model file modelled after the Livermore model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These numbers _almost_ reproduce the Livermore model included in the SNOwGLoBES repository.\n",
    "# They are obtained by calculating the total L, <E> and <E^2> from the livermore.dat\n",
    "# fluence file (which is modelled after a 10kpc supernova).\n",
    "total_energy = (5.478e+52, 5.485e+52, 4 * 5.55e+52)\n",
    "mean_energy = (11.5081, 15.4678, 21.0690)\n",
    "rms_or_pinch = \"rms\"\n",
    "rms_energy = (12.8788, 17.8360, 24.3913)\n",
    "\n",
    "# Make an astropy table with two times, 0s and 1s, with constant neutrino properties\n",
    "table = Table()\n",
    "table['TIME'] = np.linspace(0,1,2)\n",
    "table['L_NU_E'] =  np.linspace(1,1,2)*total_energy[0]\n",
    "table['L_NU_E_BAR'] = np.linspace(1,1,2)*total_energy[1]\n",
    "table['L_NU_X'] = np.linspace(1,1,2)*total_energy[2]/4. #Note, L_NU_X is set to 1/4 of the total NU_X energy\n",
    "        \n",
    "table['E_NU_E'] = np.linspace(1,1,2)*mean_energy[0]\n",
    "table['E_NU_E_BAR'] = np.linspace(1,1,2)*mean_energy[1]\n",
    "table['E_NU_X'] = np.linspace(1,1,2)*mean_energy[2]\n",
    "\n",
    "if rms_or_pinch == \"rms\":\n",
    "    table['RMS_NU_E'] = np.linspace(1,1,2)*rms_energy[0]\n",
    "    table['RMS_NU_E_BAR'] = np.linspace(1,1,2)*rms_energy[1]\n",
    "    table['RMS_NU_X'] = np.linspace(1,1,2)*rms_energy[2]\n",
    "    table['ALPHA_NU_E'] = (2.0 * table['E_NU_E'] ** 2 - table['RMS_NU_E'] ** 2) / (\n",
    "        table['RMS_NU_E'] ** 2 - table['E_NU_E'] ** 2)\n",
    "    table['ALPHA_NU_E_BAR'] = (2.0 * table['E_NU_E_BAR'] ** 2 - table['RMS_NU_E_BAR'] ** 2) / (\n",
    "        table['RMS_NU_E_BAR'] ** 2 - table['E_NU_E_BAR'] ** 2)\n",
    "    table['ALPHA_NU_X'] = (2.0 * table['E_NU_X'] ** 2 - table['RMS_NU_X'] ** 2) / (\n",
    "        table['RMS_NU_X'] ** 2 - table['E_NU_X'] ** 2)\n",
    "elif rms_or_pinch == \"pinch\":\n",
    "    table['ALPHA_NU_E'] = np.linspace(1,1,2)*pinch_values[0]\n",
    "    table['ALPHA_NU_E_BAR'] = np.linspace(1,1,2)*pinch_values[1]\n",
    "    table['ALPHA_NU_X'] = np.linspace(1,1,2)*pinch_values[2]\n",
    "    table['RMS_NU_E'] = np.sqrt((2.0 + table['ALPHA_NU_E'])/(1.0 + table['ALPHA_NU_E'])*table['E_NU_E']**2)\n",
    "    table['RMS_NU_E_BAR'] =  np.sqrt((2.0 + table['ALPHA_NU_E_BAR'])/(1.0 + table['ALPHA_NU_E_BAR'])*table['E_NU_E_BAR']**2)\n",
    "    table['RMS_NU_X'] = np.sqrt((2.0 + table['ALPHA_NU_X'])/(1.0 + table['ALPHA_NU_X'])*table['E_NU_X']**2 )\n",
    "else:\n",
    "    print(\"incorrect second moment method: rms or pinch\")\n",
    "\n",
    "filename = \"AnalyticFluence_demo.dat\"\n",
    "table.write(model_folder + filename, format='ascii', overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the Analytic Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "filename = \"AnalyticFluence_demo.dat\"\n",
    "model = Analytic3Species(model_folder + filename)\n",
    "flavors = [Flavor.NU_E,Flavor.NU_E_BAR,Flavor.NU_X]\n",
    "\n",
    "fig,axes = plt.subplots(1,3,figsize=(16,3))\n",
    "plt.subplots_adjust(wspace=0.3)\n",
    "for flavor in flavors:\n",
    "    axes[0].plot(model.time,model.luminosity[flavor],label=flavor.to_tex())\n",
    "axes[0].set_ylabel(\"Luminosity [erg/s]\")\n",
    "axes[0].set_xlabel(\"time [s]\")\n",
    "axes[0].legend(frameon=False)\n",
    "\n",
    "for flavor in flavors:\n",
    "    axes[1].plot(model.time,model.meanE[flavor],label=flavor.to_tex())\n",
    "axes[1].set_ylabel(\"Mean Energy [MeV]\")\n",
    "axes[1].set_xlabel(\"time [s]\")\n",
    "axes[1].legend(frameon=False)\n",
    "\n",
    "for flavor in flavors:\n",
    "    axes[2].plot(model.time,model.pinch[flavor],label=flavor.to_tex())\n",
    "axes[2].set_ylabel(\"pinch parameter\")\n",
    "axes[2].set_xlabel(\"time [s]\")\n",
    "axes[2].legend(frameon=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating Number of Events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#set model type as \n",
    "modeltype = 'Analytic3Species'\n",
    "#set model file we just created above\n",
    "filename = \"AnalyticFluence_demo.dat\"\n",
    "\n",
    "#set desired Oscillation scenario, snowglobes default is\n",
    "#NoTransformation, so that is used here to match the output of\n",
    "#./supernova.pl livermore water wc100kt30prct\n",
    "transformation = \"NoTransformation\"\n",
    "\n",
    "#set distance in kpc\n",
    "distance=10\n",
    "\n",
    "#set desired detector\n",
    "detector='wc100kt30prct'\n",
    "\n",
    "#Running the SNEWPY/SNOwGLoBES modules\n",
    "outfile = \"Analytic3Species_demo\"\n",
    "\n",
    "#first generated integrated fluence files for SNOwGLoBES\n",
    "print(\"Preparing fluences ...\")\n",
    "tarredoutfile = snowglobes.generate_fluence(model_folder + filename, modeltype, transformation, distance, outfile)\n",
    "\n",
    "#run the fluence file through SNOwGLoBES \n",
    "print(\"Running SNOwGLoBES ...\")\n",
    "snowglobes.simulate(SNOwGLoBES_path, tarredoutfile, detector_input=detector)\n",
    "\n",
    "print(\"Collating...\")\n",
    "tables = snowglobes.collate(SNOwGLoBES_path, tarredoutfile, skip_plots=True)\n",
    "\n",
    "totalcounts = 0\n",
    "for i in range(1,6):\n",
    "    print(\"Number of events of type\",tables['Collated_'+outfile+'_'+detector+'_events_smeared_weighted.dat']['header'].split()[i],end=': ')\n",
    "    counts = sum(tables['Collated_'+outfile+'_'+detector+'_events_smeared_weighted.dat']['data'][i])\n",
    "    totalcounts += counts\n",
    "    print(counts)\n",
    "print(\"Total number of events: \",totalcounts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for smear in [\"smeared\", \"unsmeared\"]:\n",
    "    energy = tables['Collated_'+outfile+'_'+detector+'_events_'+smear+'_weighted.dat']['data'][0]*1000.\n",
    "    nc = tables['Collated_'+outfile+'_'+detector+'_events_'+smear+'_weighted.dat']['data'][1]\n",
    "    escattering = tables['Collated_'+outfile+'_'+detector+'_events_'+smear+'_weighted.dat']['data'][2]\n",
    "    ibd = tables['Collated_'+outfile+'_'+detector+'_events_'+smear+'_weighted.dat']['data'][3]\n",
    "    nueO16 = tables['Collated_'+outfile+'_'+detector+'_events_'+smear+'_weighted.dat']['data'][4]\n",
    "    nuebarO16 = tables['Collated_'+outfile+'_'+detector+'_events_'+smear+'_weighted.dat']['data'][5]\n",
    "    plt.plot(energy,nc,label=\"NC\")\n",
    "    plt.plot(energy,escattering,label=\"escattering\")\n",
    "    plt.plot(energy,ibd,label=\"ibd\")\n",
    "    plt.plot(energy,nueO16,label=\"nueO16\")\n",
    "    plt.plot(energy,nuebarO16,label=\"nuebarO16\")\n",
    "    plt.legend()\n",
    "    plt.title(smear)\n",
    "    plt.xlabel(\"Energy [MeV]\")\n",
    "    plt.ylabel(\"Count per bin\")\n",
    "    plt.yscale('log')\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "12a4e164d15418f42fe0584c567a58ace9669f8a11b564e22ebd17e6959ef919"
  },
  "kernelspec": {
   "display_name": "Python 3.9.5 64-bit ('snews': conda)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}