{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KDE PDFs Via GenericPDFRatio\n",
    "\n",
    ".. contents:: :local:\n",
    "\n",
    "This notebook will walk you through using the generic PDF implementation in csky to implement the KDE spatial PDFs that were used in the most [recent NorthernTracks time integrated analysis](https://user-web.icecube.wisc.edu/~tglauch/PS_Analysis/_build/html/index.html). In addition to reproducing the time integrated results for NGC 1068, we'll additionally show how this can be combined with the untriggered flare likelihood, allowing the KDE spatial PDF description to be used for transient analyses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Imports we'll need\n",
    "import numpy as np\n",
    "import csky as cy\n",
    "import os\n",
    "import photospline as psp\n",
    "from csky import cext, hyp\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll create our function that calculates the generic PDF ratio, as well as some helper functions to evaluate the provded KDEs and mask out events that are far away (to save computation time). We also need to load up the spline tables that the improved point source people have generated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_path='/data/ana/analyses/northern_tracks/version-005-p00/KDE_PDFs_v007/IC86_pass2/'\n",
    "\n",
    "spatial_pdf = psp.SplineTable(os.path.join(base_path, 'sig_E_psi_photospline_v006_4D.fits'))\n",
    "energy_pdf = psp.SplineTable(os.path.join(base_path, 'E_dec_photospline_v006_3D.fits'))\n",
    "bkg_pdf = psp.SplineTable(os.path.join(base_path, 'bg_2d_photospline.fits'))\n",
    "\n",
    "\n",
    "def GreatCircleDistance(ra_1, dec_1, ra_2, dec_2):\n",
    "        '''Compute the great circle distance between two events'''\n",
    "        '''All coordinates must be given in radians'''\n",
    "        delta_dec = np.abs(dec_1 - dec_2)\n",
    "        delta_ra = np.abs(ra_1 - ra_2)\n",
    "        x = (np.sin(delta_dec / 2.))**2. + np.cos(dec_1) *\\\n",
    "            np.cos(dec_2) * (np.sin(delta_ra / 2.))**2.\n",
    "        return 2. * np.arcsin(np.sqrt(x))\n",
    "\n",
    "def spatial_kde(log_psi, log_ang_err, log_energy, gamma):\n",
    "    '''Evaluates the spatial kde for an event and spectral index hypothesis'''\n",
    "    return spatial_pdf.evaluate_simple([log_ang_err, log_energy, log_psi, np.full(len(log_psi), gamma)], 0)\n",
    "\n",
    "def energy_kde(log_energy, src_dec, gamma):\n",
    "    '''Evaluates the energy kde for an event and spectral index hypothesis'''\n",
    "    sin_src_dec = np.sin(src_dec)\n",
    "    return energy_pdf.evaluate_simple([log_energy, np.full(len(log_energy), sin_src_dec), np.full(len(log_energy), gamma)], 0)\n",
    "\n",
    "def calc_sb_ratio(ev_ra, ev_dec, ev_angErr, ev_logE, src, gamma, **kwargs):\n",
    "    '''Calculates the (space x energy) S/B ratio using spatial and energy kde pdfs'''\n",
    "    '''There are 2 cases, one where the mask is provided by the trial runner args, and one \n",
    "        where no mask has been provided'''\n",
    "    '''This is the function that we will be pointing our trial runner to to use as a \"generic\" pdf ratio'''\n",
    "    if '_mask' in kwargs:\n",
    "        mask = kwargs['_mask']\n",
    "        antimask = mask[mask==False]\n",
    "        src_ra = src['ra']\n",
    "        src_dec = src['dec']\n",
    "\n",
    "        ev_ra_mask = ev_ra[mask]\n",
    "        ev_dec_mask = ev_dec[mask]\n",
    "        ev_angErr_mask = ev_angErr[mask]\n",
    "        ev_logE_mask = ev_logE[mask]\n",
    "\n",
    "        psi = GreatCircleDistance(ev_ra_mask, ev_dec_mask, src_ra, src_dec)\n",
    "\n",
    "        log_psi = np.log10(psi)\n",
    "        log_ang_err = np.log10(ev_angErr_mask)\n",
    "        log_energy = ev_logE_mask\n",
    "        masked_ev_dec = ev_dec_mask\n",
    "\n",
    "        spdfs = spatial_kde(log_psi, log_ang_err, log_energy, gamma)\n",
    "        spatial_norm = np.log(10) * psi * np.sin(psi)\n",
    "        spdfs/= spatial_norm\n",
    "\n",
    "        epdfs = energy_kde(log_energy, src_dec, gamma)\n",
    "\n",
    "        bg_pdfs = bkg_pdf.evaluate_simple([log_energy, np.sin(masked_ev_dec)])\n",
    "        sb_ratio_close = spdfs*epdfs/bg_pdfs\n",
    "\n",
    "        far_evts_sb = np.zeros(len(antimask))\n",
    "        sb_ratio = np.zeros(len(ev_ra))\n",
    "        sb_ratio[mask] = sb_ratio_close\n",
    "\n",
    "    else:\n",
    "        src_ra = src['ra']\n",
    "        src_dec = src['dec']\n",
    "        psi = GreatCircleDistance(ev_ra, ev_dec, src_ra, src_dec)\n",
    "\n",
    "        log_psi = np.log10(psi)\n",
    "        log_ang_err = np.log10(ev_angErr)\n",
    "        log_energy = ev_logE\n",
    "        masked_ev_dec = ev_dec\n",
    "\n",
    "        spdfs = spatial_kde(log_psi, log_ang_err, log_energy, gamma)\n",
    "        spatial_norm = np.log(10) * psi* np.sin(psi)\n",
    "        spdfs/= spatial_norm\n",
    "\n",
    "        epdfs = energy_kde(log_energy, src_dec, gamma)\n",
    "\n",
    "        bg_pdfs = bkg_pdf.evaluate_simple([log_energy, np.sin(masked_ev_dec)])\n",
    "        sb_ratio = spdfs*epdfs/bg_pdfs\n",
    "\n",
    "    return sb_ratio\n",
    "\n",
    "def dist_mask(arr, src, cut_deg):\n",
    "    '''A simple distance mask that masks out events (arr) farther than cut_deg away from src'''\n",
    "    out = np.zeros(len(arr), dtype=bool)\n",
    "    ras = arr['ra']\n",
    "    decs = arr['dec']\n",
    "    psi = GreatCircleDistance(ras, decs, src.ra, src.dec)\n",
    "    out = psi < np.radians(cut_deg)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do the usual setup of our analysis object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting up Analysis for:\n",
      "NT86v5\n",
      "Setting up NT86v5...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_pass2_MC.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2011_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2012_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2013_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2014_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2015_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2016_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2017_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2018_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/IC86_2019_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2011_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2012_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2013_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2014_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2015_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2016_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2017_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2018_exp.npy ...\n",
      "Reading /data/ana/analyses/northern_tracks/version-005-p00/GRL/IC86_2019_exp.npy ...\n",
      "Energy PDF Ratio Model...\n",
      "  * gamma = 4.0000 ...\n",
      "Signal Acceptance Model...\n",
      "  * gamma = 4.0000 ...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "#Set up our analysis objects to use NTv5\n",
    "ana_dir = cy.utils.ensure_dir('/data/ana/analyses/')\n",
    "ana = cy.get_analysis(cy.selections.repo, 'version-005-p00', cy.selections.NTDataSpecs.nt_v5[0:1], dir=ana_dir, min_sigma=0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can set up our first trial runner. This trial runner will be used to generate trials, and we'll create another one to actually evaluate the likelihood. We do this to ensure that \n",
    "\n",
    "I'll be setting this up to point to NGC 1068 (ra=40.667, dec =-0.0069) so we can check against the results that have already been obtained using SkyLLH. In our trial runner setup, we choose our space pdf ratio to be `'generic'`, and point the function to be used to `calc_sb_ratio`. `features` is a dictionary of the args that `calc_sb_ratio` will need, and `fits` are the parameters we'd like the likelihood to fit for (in addition to `ns`). Note that we turn `energy=False` here, since our pdf ratio function `calc_sb_ratio` will be handling both the space and energy components of the likelihood. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "srcs = cy.sources(np.radians(40.667), np.radians(-0.0069), name=0)\n",
    "gamma = 2.0\n",
    "\n",
    "tr = cy.get_trial_runner(ana=ana, src=srcs, space='generic', func=calc_sb_ratio, \n",
    "                         features={'ev_ra':'ra', 'ev_dec':'dec', 'ev_angErr':'sigma', 'ev_logE':'log10energy'}, \n",
    "                         fits={'gamma':(1.0,2.0,3.0,4.0)}, extra_keep=['dec', 'sigma', 'sindec', 'event'], \n",
    "                         energy=False, \n",
    "                         cut_n_sigma=np.inf, \n",
    "                         concat_evs=True, \n",
    "                         flux=hyp.PowerLawFlux(gamma))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can generate and evaluate a trial. Note that here we introduce a distance mask so we don't have to evaluate the KDE splines for events farther than 15 degrees. This significantly speeds up the calculation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(29.204403027486393, {'gamma': 3.1586663282907503, 'ns': 76.96023713945011})\n",
      "fit took 1.4317007064819336 seconds\n"
     ]
    }
   ],
   "source": [
    "trialseed = 0\n",
    "trial = tr.get_one_trial(n_sig=0, TRUTH=True, seed=trialseed)\n",
    "for src in srcs:\n",
    "\n",
    "    t1 = time.time()\n",
    "\n",
    "    trial_mask = dist_mask(trial[0][0][0], srcs, 15.)\n",
    "    fit = tr.get_one_fit_from_trial(trial,flat=False,_masks=[trial_mask], TRUTH=True)\n",
    "    print(fit)\n",
    "    t2 = time.time()\n",
    "    print(\"fit took\", t2-t1, \"seconds\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time Dependent fits with KDE PDFs\n",
    "\n",
    "We can combine the new KDE spatial/energy PDFs with the untriggered flare likelihood by using a `MultiflareTrialRunner`, set up like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "mtr = cy.conf.get_multiflare_trial_runner(ana=ana, src=srcs, threshold=1000., space='generic', func=calc_sb_ratio, \n",
    "                                          features={'ev_ra':'ra', 'ev_dec':'dec', 'ev_angErr':'sigma', 'ev_logE':'log10energy'}, \n",
    "                                          fits={'gamma':(1.0,2.0,3.0,4.0)}, extra_keep=['dec', 'sigma', 'event'], \n",
    "                                          energy=False, \n",
    "                                          concat_evs=True, \n",
    "                                          flux=hyp.PowerLawFlux(gamma), \n",
    "                                          muonflag=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finished making trial\n",
      "type of evaluator is <class 'csky.pdf.SpaceTimePDFRatioEvaluator'>\n",
      "this evaluator does not have separate space and energy components\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/cvmfs/icecube.opensciencegrid.org/py3-v4.1.1/RHEL_7_x86_64/lib/python3.7/site-packages/ipykernel_launcher.py:37: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.\n",
      "/cvmfs/icecube.opensciencegrid.org/py3-v4.1.1/RHEL_7_x86_64/lib/python3.7/site-packages/ipykernel_launcher.py:38: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.\n",
      "/cvmfs/icecube.opensciencegrid.org/py3-v4.1.1/RHEL_7_x86_64/lib/python3.7/site-packages/ipykernel_launcher.py:39: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.\n",
      "/cvmfs/icecube.opensciencegrid.org/py3-v4.1.1/RHEL_7_x86_64/lib/python3.7/site-packages/ipykernel_launcher.py:40: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.\n",
      "/cvmfs/icecube.opensciencegrid.org/py3-v4.1.1/RHEL_7_x86_64/lib/python3.7/site-packages/ipykernel_launcher.py:60: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 590 windows\n"
     ]
    }
   ],
   "source": [
    "trialseed = 0\n",
    "mtr_trial = mtr.get_one_trial(seed=trialseed, TRUTH=False)\n",
    "mtr_fit = mtr.get_one_fit_from_trial(mtr_trial, flat=False, _fmin_epsilon=1e-3)\n",
    "print(\"mtr_fit is\", mtr_fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}