{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Earth Survival Probability from USSR\n",
    "\n",
    "Check the Earth-crossing survival probability computed in the old USSR code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from importlib.resources import files\n",
    "\n",
    "from snewpy.neutrino import Flavor\n",
    "from asteria import set_rcparams"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Styles for Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_rcparams(verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Helper Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getUSSRdata(filename):\n",
    "    file = files('asteria.data.USSR').joinpath(f'earth_survival_probability/{filename}')\n",
    "    if not file.exists():\n",
    "        raise FileNotFoundError(f'Could not find earth_survival_probability/{filename} in asteria.')\n",
    "        \n",
    "    with open(file, 'r') as infile:\n",
    "        # Get Size of data from first line and initialize.\n",
    "        line = infile.readline().strip('#').replace(',','').split()\n",
    "        Enu = np.zeros( int(line[1]) )\n",
    "        data = np.zeros( shape=( int(line[0]), int(line[1])) )\n",
    "        \n",
    "        # Clear other lines of header\n",
    "        infile.readline() \n",
    "        infile.readline() \n",
    "        i = 0\n",
    "        for line in infile:\n",
    "            line = line.split()\n",
    "            Enu[i] = float(line[0])\n",
    "            for nu, item in enumerate(line[1:]):\n",
    "                data[nu][i] = float(item)\n",
    "            i+=1\n",
    "    return Enu, data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Earth Survival Probability from USSR\n",
    "Commands used: `$ ./sim_test -mn 404 -oem 33 -os {os} -osm {osm}` with...\n",
    " - `{osm} = 0,1` For Star oscillations off, on\n",
    " - `{os} = 1,2,3 / 4,5,6` For $sin^2(2\\theta_{13})$ = 1e-2, 5e-4, 1e-6. (1,2,3 are NH / 4,5,6 are IH)\n",
    " \n",
    " In all cases $\\delta_{cp} = 0$. $\\theta_n$ is the nadir angle, which at the time of writing is always 33 $^{\\circ}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Enu, Pe = getUSSRdata( 'nadir-33__sqsin2theta13-1E-2__NH__StarOsc_Off.txt' )\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "for nu, flavor in enumerate(Flavor):\n",
    "    ax.plot(Enu, Pe[nu], lw=2, alpha=0.7, label=flavor.to_tex())\n",
    "    \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1e-2$' +' | '+'NH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "ax.plot(Enu, Pe[0] + Pe[2], lw=2, alpha=0.7, label=r'$\\nu$')\n",
    "ax.plot(Enu, Pe[1] + Pe[3], lw=2, alpha=0.7, label=r'$\\bar{\\nu}$') \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1E-6$' +' | '+'IH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Enu, Pe = getUSSRdata( 'nadir-33__sqsin2theta13-1E-2__IH__StarOsc_Off.txt' )\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "for nu, flavor in enumerate(Flavor):\n",
    "    ax.plot(Enu, Pe[nu], lw=2, alpha=0.7, label=flavor.to_tex())\n",
    "    \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1e-2$' +' | '+'IH' + ' | Star Osc. OFF' )\n",
    "ax.legend(loc='upper right');\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "ax.plot(Enu, Pe[0] + Pe[2], lw=2, alpha=0.7, label=r'$\\nu$')\n",
    "ax.plot(Enu, Pe[1] + Pe[3], lw=2, alpha=0.7, label=r'$\\bar{\\nu}$') \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1E-6$' +' | '+'IH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Enu, Pe = getUSSRdata( 'nadir-33__sqsin2theta13-5E-4__NH__StarOsc_Off.txt' )\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "for nu, flavor in enumerate(Flavor):\n",
    "    ax.plot(Enu, Pe[nu], lw=2, alpha=0.7, label=flavor.to_tex())\n",
    "    \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 5e-4$' +' | '+'NH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "ax.plot(Enu, Pe[0] + Pe[2], lw=2, alpha=0.7, label=r'$\\nu$')\n",
    "ax.plot(Enu, Pe[1] + Pe[3], lw=2, alpha=0.7, label=r'$\\bar{\\nu}$') \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1E-6$' +' | '+'IH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Enu, Pe = getUSSRdata( 'nadir-33__sqsin2theta13-5E-4__IH__StarOsc_Off.txt' )\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "for nu, flavor in enumerate(Flavor):\n",
    "    ax.plot(Enu, Pe[nu], lw=2, alpha=0.7, label=flavor.to_tex())\n",
    "    \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 5e-4$' +' | '+'IH' + ' | Star Osc. OFF' )\n",
    "ax.legend(loc='upper right');\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "ax.plot(Enu, Pe[0] + Pe[2], lw=2, alpha=0.7, label=r'$\\nu$')\n",
    "ax.plot(Enu, Pe[1] + Pe[3], lw=2, alpha=0.7, label=r'$\\bar{\\nu}$') \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1E-6$' +' | '+'IH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Enu, Pe = getUSSRdata( 'nadir-33__sqsin2theta13-1E-6__NH__StarOsc_Off.txt' )\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "for nu, flavor in enumerate(Flavor):\n",
    "    ax.plot(Enu, Pe[nu], lw=2, alpha=0.7, label=flavor.to_tex())\n",
    "    \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1E-6$' +' | '+'NH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "ax.plot(Enu, Pe[0] + Pe[2], lw=2, alpha=0.7, label=r'$\\nu$')\n",
    "ax.plot(Enu, Pe[1] + Pe[3], lw=2, alpha=0.7, label=r'$\\bar{\\nu}$') \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1E-6$' +' | '+'IH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Enu, Pe = getUSSRdata( 'nadir-33__sqsin2theta13-1E-6__IH__StarOsc_Off.txt' )\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "for nu, flavor in enumerate(Flavor):\n",
    "    ax.plot(Enu, Pe[nu], lw=2, alpha=0.7, label=flavor.to_tex())\n",
    "    \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1E-6$' +' | '+'IH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');\n",
    "                       \n",
    "fig, ax = plt.subplots(1,1, figsize=(9,5))\n",
    "ax.plot(Enu, Pe[0] + Pe[2], lw=2, alpha=0.7, label=r'$\\nu$')\n",
    "ax.plot(Enu, Pe[1] + Pe[3], lw=2, alpha=0.7, label=r'$\\bar{\\nu}$') \n",
    "ax.set_xlabel( r'E$_\\nu$ [MeV]', horizontalalignment='right', x = 1)\n",
    "ax.set_ylabel( 'Probability' , horizontalalignment='right', y = 1)\n",
    "ax.set_title(r'$\\theta_n = 33$'+' | '+ r'$sin^2 (2 \\theta_{13}) = 1E-6$' +' | '+'IH' + ' | Star Osc. Off' )\n",
    "ax.legend(loc='upper right');"
   ]
  }
 ],
 "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}