{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f2414af0-d2cf-472d-baf7-b67abb62796b",
   "metadata": {},
   "source": [
    "## Interactive OscProb oscillogram Plotter\n",
    "\n",
    "This is an example showing ipywidgets to interactively plot a function; in this case\n",
    "neutrino oscillation probabilities.\n",
    "\n",
    "To run this notebook with the central lyon jupyter service, set up the\n",
    "km3net-env kernel as explained in the wiki \n",
    "\n",
    "https://wiki.km3net.de/index.php/Computing/KM3NeT_Computing_at_CC_Lyon#Jupyter_Notebook_Server \n",
    "\n",
    "then open this notebook via\n",
    "\n",
    "https://notebook.cc.in2p3.fr\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "482dea2b-abb2-4333-beaf-61b49d48616c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ROOT, IPython, numpy, inspect\n",
    "from math import *\n",
    "from ipywidgets import interact, interactive, VBox, HBox\n",
    "\n",
    "ROOT.gSystem.Load(\"/sps/km3net/users/heijboer/OscProb/libOscProb.so\") # the oscprob module needs updating -- so use this\n",
    "ROOT.gStyle.SetOptStat(0)\n",
    "ROOT.TH1.AddDirectory(0)\n",
    "#ROOT.gInterpreter->AddIncludePath(\"/sps/km3net/users/heijboer/OscProb/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "66d21ed7-5b6e-4d93-9047-11bba1017b57",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# helpers to convert ROOT canvas to Ipython-widgets compatble image, and to display them.\n",
    "\n",
    "ROOT.gInterpreter.Declare(\"\"\"\n",
    "#include \"TImage.h\"\n",
    "\n",
    "struct gbx\n",
    "{\n",
    "    unsigned char* buffer;\n",
    "    int l;\n",
    "\n",
    "    gbx( TASImage* img ) {\n",
    "        img->GetImageBuffer( (char**) &buffer, &l, TImage::kPng );\n",
    "    }\n",
    "};\n",
    "\"\"\")\n",
    "\n",
    "\n",
    "def to_img( pad ) :\n",
    "    \n",
    "    \"transform a root TPad to an Ipython.display image\"\n",
    "    \n",
    "    image = ROOT.TImage.Create()\n",
    "    image.FromPad( pad )\n",
    "    Y = ROOT.gbx( image )\n",
    "    x = numpy.recarray( (Y.l,), dtype='c', buf=Y.buffer )\n",
    "    img = IPython.display.Image( x , format='png' )\n",
    "    return img \n",
    "\n",
    "\n",
    "def show( f ) :\n",
    "    \n",
    "    'f should be a function returning a histogram; show the histogram in the out display'\n",
    "\n",
    "    global out\n",
    "    \n",
    "    def r ( *args, **kwargs ) :\n",
    "        with out:\n",
    "           h = f( *args, **kwargs )\n",
    "           h.Draw(\"colz\")\n",
    "           img = to_img( ROOT.gPad )\n",
    "           display(img)\n",
    "           del h\n",
    "    \n",
    "    # give r the same signature\n",
    "    r.__signature__ = inspect.signature( f )\n",
    "    return r\n",
    "\n",
    "# make a convas\n",
    "c = ROOT.TCanvas(\"c\",\"c\",500,500)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c9810aa3-1338-4910-abda-c980794b4431",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def oscillogram( initial_flav=1, final_flav=1, mh = 1, \n",
    "                dm21 = 7.5, \n",
    "                dm31 = 2.457,\n",
    "                th12 = asin(sqrt(0.304)), \n",
    "                th13 = asin(sqrt(0.0218)), \n",
    "                th23 = asin(sqrt(0.452 )), \n",
    "                dcp =0 ,\n",
    "                nbinsx = 50, nbinsy = 50 ) :\n",
    "    \n",
    "    P = ROOT.OscProb.PMNS_Fast()\n",
    "    prem = ROOT.OscProb.PremModel()\n",
    "    \n",
    "    P.SetDm(2, dm21 * 1e-5);\n",
    "    P.SetDm(3, dm31 * 1e-3);\n",
    "    P.SetAngle(1,2, th12);\n",
    "    P.SetAngle(1,3, th13);\n",
    "    P.SetAngle(2,3, th23);\n",
    "    P.SetDelta(1,3, dcp );\n",
    "    \n",
    "    h = ROOT.TH2D(\"h\",\"oscillogram;Energy (GeV);cos zenith\",nbinsx,0,20,nbinsy,-1,0);\n",
    "    \n",
    "    for iy in range( 1, h.GetNbinsY()+1) :    \n",
    "    \n",
    "        cos_theta = h.GetYaxis().GetBinCenter(iy)\n",
    "        prem.FillPath( cos_theta )\n",
    "        P.SetPath(prem.GetNuPath())\n",
    "    \n",
    "        for ix in range ( 1, h.GetNbinsX()+1) :   \n",
    "            E = h.GetXaxis().GetBinCenter(ix)\n",
    "            p = P.Prob( initial_flav, final_flav, E);\n",
    "            h.SetBinContent(ix,iy,p)\n",
    "    \n",
    "    return h\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0b161656-24d3-47c8-ab0a-d6803cac637a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d140e3cbdc464341b02f4bd7510c6ede",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(IntSlider(value=1, description='initial_flav', max=3, min=1), IntSlider(value=1,…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee6b45088c1346b0b3c59214acd3e9dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output(layout=Layout(border='3px solid blue', height='500px', width='600px'))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "iplot = interactive( show( oscillogram ),\n",
    "                     initial_flav=(1,3,1), final_flav=(1,3,1), mh = (-1,1),       \n",
    "                     dm21 = ( 6, 9 ,0.1),\n",
    "                     dm31 = ( 2, 3, 0.1 ) )\n",
    "\n",
    "leftbox  = VBox( iplot.children[: int(len(iplot.children) / 2)])\n",
    "rightbox = VBox( iplot.children[int(len(iplot.children) / 2) : -1] )\n",
    "display( HBox([leftbox, rightbox]) )\n",
    "\n",
    "out = iplot.children[-1]\n",
    "out.layout = {\n",
    "            'width': '600px',\n",
    "            'height': '500px',\n",
    "            'border': '3px solid blue'\n",
    "            }\n",
    "iplot\n",
    "display(out)\n",
    "show(oscillogram)() # once now 'out' exists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9057eb6a-52d6-497e-89b3-477b8d5b3d21",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading root....  /pbs/throng/km3net/software/root/6.22.06\n",
      "loading /pbs/throng/km3net/software/aanet/2.2.6/lib/libaaevt.so ... ok (0.15 s.)\n",
      "loading /pbs/throng/km3net/software/aanet/2.2.6/lib/libaastro.so ... ok (0.04 s.)\n",
      " welcome to aanet v2.2.6\n",
      "reducing __dir__s for ipython tab completion..\n",
      " new timer python-based\n",
      " new timer c++-based\n",
      "1000.5\n",
      "calibration took 0.0113809 27528999 4.13415e-10 \n",
      "-------------------------------------------\n",
      " TickTimer Report \n",
      "-------------------------------------------\n",
      " timer        ncalls total(s) time/call(s) \n",
      "-------------------------------------------\n",
      " c++-based    10     0.791905 0.0791905    \n",
      " python-based 10     1.02792  0.102792     \n",
      "-------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# curious about speedup, when filling the histogram in C++\n",
    "import aa\n",
    "\n",
    "ROOT.cxx(\"\"\"\n",
    "\n",
    "TH2D oscillogram( int initial_flav=1, int final_flav=1, int mh = 1, \n",
    "                double dm21 = 7.5, \n",
    "                double dm31 = 2.457,\n",
    "                double th12 = asin(sqrt(0.304)), \n",
    "                double th13 = asin(sqrt(0.0218)), \n",
    "                double th23 = asin(sqrt(0.452 )), \n",
    "                double dcp =0 ,\n",
    "                int nbinsx = 50, int nbinsy = 50 ) \n",
    "{\n",
    "    OscProb::PMNS_Fast P;\n",
    "    OscProb::PremModel prem;\n",
    "    \n",
    "    P.SetDm(2, dm21 * 1e-5);\n",
    "    P.SetDm(3, dm31 * 1e-3);\n",
    "    P.SetAngle(1,2, th12);\n",
    "    P.SetAngle(1,3, th13);\n",
    "    P.SetAngle(2,3, th23);\n",
    "    P.SetDelta(1,3, dcp );\n",
    "    \n",
    "    TH2D h (\"h\",\"oscillogram;Energy (GeV);cos zenith\",nbinsx,0,20,nbinsy,-1,0);\n",
    "    h.SetDirectory(0);\n",
    "    \n",
    "    for ( int iy = 1; iy <= h.GetNbinsY()+1; iy++ )\n",
    "    {   \n",
    "        const double cos_theta = h.GetYaxis()->GetBinCenter(iy);\n",
    "        prem.FillPath( cos_theta );\n",
    "        P.SetPath(prem.GetNuPath());\n",
    "    \n",
    "        for ( int ix = 1 ; ix <= h.GetNbinsX()+1; ix ++ )\n",
    "        {\n",
    "            double E = h.GetXaxis()->GetBinCenter(ix);\n",
    "            double p = P.Prob( initial_flav, final_flav, E);\n",
    "            h.SetBinContent(ix,iy,p);\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    return h;\n",
    "}\n",
    "\"\"\")\n",
    "\n",
    "cpposcillogram = ROOT.oscillogram\n",
    "\n",
    "ROOT.Timer.reset_all()\n",
    "\n",
    "for i in range(10) :\n",
    "\n",
    "    with ROOT.Timer(\"python-based\") :\n",
    "        h = oscillogram()\n",
    "\n",
    "    with ROOT.Timer(\"c++-based\") :\n",
    "        h = cpposcillogram()\n",
    "\n",
    "ROOT.Timer.report()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "km3net-env-1.9",
   "language": "python",
   "name": "km3net-env-1.9"
  },
  "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": 5
}