AcceptanceEffects2D_EG2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "#os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.patches import Patch\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import tensorflow as tf\n",
    "import keras\n",
    "import keras.backend as K\n",
    "from keras.layers import Dense, Input\n",
    "from keras.models import Model\n",
    "from keras.utils import to_categorical\n",
    "from keras.callbacks import EarlyStopping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import ROOT\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import root_pandas as rpd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as mticker\n",
    "from matplotlib.colors import LogNorm\n",
    "from root_pandas import read_root\n",
    "from matplotlib import rc\n",
    "from numpy import inf\n",
    "\n",
    "\n",
    "import sys,pandas as pd, matplotlib , matplotlib.pyplot as plt, matplotlib.lines , numpy as np, math, pylab\n",
    "#import ROOT\n",
    "#from ROOT import TFile\n",
    "\n",
    "import root_pandas\n",
    "%matplotlib inline\n",
    "\n",
    "rc('text', usetex=True)\n",
    "\n",
    "import matplotlib as mpl\n",
    "rc('font', family='serif')\n",
    "rc('text', usetex=True)\n",
    "rc('font', size=22)\n",
    "rc('xtick', labelsize=15)\n",
    "rc('ytick', labelsize=15)\n",
    "rc('legend', fontsize=15)\n",
    "\n",
    "#\n",
    "mpl.rcParams.update({'font.size': 19})\n",
    "#mpl.rcParams.update({'legend.fontsize': 18})\n",
    "mpl.rcParams.update({'xtick.labelsize': 18}) \n",
    "mpl.rcParams.update({'ytick.labelsize': 18}) \n",
    "mpl.rcParams.update({'text.usetex' : False})\n",
    "mpl.rcParams.update({'axes.labelsize': 18}) \n",
    "mpl.rcParams.update({'legend.frameon': False}) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Global plot settings\n",
    "from matplotlib import rc\n",
    "import matplotlib.font_manager\n",
    "\n",
    "rc('font', family='serif')\n",
    "rc('text', usetex=True)\n",
    "rc('font', size=22)\n",
    "rc('xtick', labelsize=15)\n",
    "rc('ytick', labelsize=15)\n",
    "rc('legend', fontsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define default plot styles\n",
    "plot_style_0 = {\n",
    "    'histtype': 'step',\n",
    "    'color': 'black',\n",
    "    'linewidth': 2,\n",
    "    'linestyle': '--',\n",
    "    'density': True\n",
    "}\n",
    "\n",
    "plot_style_1 = {\n",
    "    'histtype': 'step',\n",
    "    'color': 'black',\n",
    "    'linewidth': 2,\n",
    "    'density': True\n",
    "}\n",
    "\n",
    "plot_style_2 = {'alpha': 0.5, 'density': True}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check Versions\n",
    "print(tf.__version__)  # 1.15.0\n",
    "print(keras.__version__)  # 2.2.4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Obtain datasets and simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/home/miguel/EG2_DATA/sebastian/'\n",
    "data_electrons = root_pandas.read_root(path+'C.root', \"ntuple_data_electrons\")\n",
    "data_hadron = root_pandas.read_root(path+'C.root', \"ntuple_data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/home/miguel/EG2_DATA/andres/'\n",
    "mc_electrons = root_pandas.read_root(path+'MC_D.root', \"ntuple_e\")\n",
    "mc_hadron = root_pandas.read_root(path+'MC_D.root', \"ntuple_sim\")\n",
    "mc_hadron['PID'] = mc_hadron['pid']\n",
    "mc_hadron['mc_PID'] = mc_hadron['mc_pid']\n",
    "\n",
    "mc_hadron['pass_reco'] = np.where(mc_hadron['Pt2']>0, 1, 0)\n",
    "mc_hadron['pass_truth'] = np.where(mc_hadron['mc_Pt2']>0, 1, 0)\n",
    "\n",
    "data_hadron['pass_reco'] = np.where(data_hadron['Pt2']>0, 1, 0)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define cuts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def applyCut(inputDataframe, cut, text=None):\n",
    "    dataframe = inputDataframe\n",
    "    nbeforecut = dataframe.shape[0]\n",
    "    cutDataframe = dataframe.query(cut)\n",
    "    if text:\n",
    "        print (text, cutDataframe.shape[0], ' fraction kept: %2.1f'%(100.0*float(cutDataframe.shape[0])/nbeforecut))\n",
    "    return cutDataframe\n",
    "\n",
    "def applyCutsHadron(df,isMC=False):\n",
    "    temp = df\n",
    "    temp.eval('inelasticity = Nu/5.014', inplace=True)\n",
    "    temp.eval('E = sqrt(P*P+0.1396*0.1396)', inplace=True)\n",
    "    temp.eval('z = E/Nu',inplace=True)\n",
    "    temp= applyCut(temp, 'inelasticity<0.85','inelasticity < 0.85')\n",
    "    if(isMC):\n",
    "        temp = applyCut(temp,'pass_truth>0',' pass_truth>0')\n",
    "        temp = applyCut(temp,'abs(mc_pid)==211',' abs(mc_pid)=211')\n",
    "        temp.eval('mc_E = sqrt(mc_P*mc_P+0.1396*0.1396)', inplace=True)\n",
    "        temp.eval('mc_z = mc_E/mc_Nu',inplace=True)\n",
    "    else:\n",
    "        temp = applyCut(temp, 'TargType==1', 'TargType==1')\n",
    "        temp = applyCut(temp, 'abs(PID)==211','abs(PID)==211')\n",
    "\n",
    "    return temp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Apply cuts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Selecting MC events\\n')\n",
    "mc_hadron = applyCutsHadron(mc_hadron,isMC=True)\n",
    "print('Selecting data events\\n')\n",
    "data_hadron= applyCutsHadron(data_hadron,isMC=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Count the non zero pass_reco and pass_truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.count_nonzero(mc_hadron['pass_reco'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(mc_hadron['pass_reco'])\n",
    "#print(len(mc_hadron))\n",
    "#print(len(data_hadron))\n",
    "#print('A' , np.count_nonzero(mc_hadron['pass_reco'])) ##these are the ones that do have reco associated to them\n",
    "#print('B', np.count_nonzero(mc_hadron['pass_truth'])) ##these are the ones that do have truth associated with them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nevents = int(8.5e6)\n",
    "data = data_hadron.sample(n=nevents)\n",
    "mc   = mc_hadron.sample(n=nevents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "##Hadron momentum\n",
    "Px_0_G = np.array(mc['mc_Px'])\n",
    "Px_0_S =np.array(mc['Px'])\n",
    "Px_unknown_S = np.array(data['Px'])\n",
    "Px_unknown_G = np.array(data['Px'])\n",
    "\n",
    "Py_0_G = np.array(mc['mc_Py'])\n",
    "Py_0_S =np.array(mc['Py'])\n",
    "Py_unknown_S = np.array(data['Py'])\n",
    "Py_unknown_G = np.array(data['Py'])\n",
    "\n",
    "Pz_0_G = np.array(mc['mc_Pz'])\n",
    "Pz_0_S = np.array(mc['Pz'])\n",
    "Pz_unknown_S = np.array(data['Pz'])\n",
    "Pz_unknown_G = np.array(data['Pz'])\n",
    "\n",
    "Pid_0_G = np.array(mc['mc_PID'])\n",
    "Pid_0_S = np.array(mc['PID'])\n",
    "Pid_unknown_S = np.array(data['PID'])\n",
    "Pid_unknown_G = np.array(data['PID'])\n",
    "\n",
    "Q2_0_G = np.array(mc['mc_Q2'])\n",
    "Q2_0_S = np.array(mc['Q2'])\n",
    "Q2_unknown_S = np.array(data['Q2'])\n",
    "Q2_unknown_G = np.array(data['Q2'])\n",
    "\n",
    "Nu_0_G = np.array(mc['mc_Nu'])\n",
    "Nu_0_S = np.array(mc['Nu'])\n",
    "Nu_unknown_S = np.array(data['Nu'])\n",
    "Nu_unknown_G = np.array(data['Nu'])\n",
    "\n",
    "PhiPQ_0_G = np.array(mc['mc_PhiPQ'])\n",
    "PhiPQ_0_S = np.array(mc['PhiPQ'])\n",
    "PhiPQ_unknown_S = np.array(data['PhiPQ'])\n",
    "PhiPQ_unknown_G = np.array(data['PhiPQ'])\n",
    "\n",
    "Zh_0_G = np.array(mc['mc_z'])\n",
    "Zh_0_S =np.array(mc['z'])\n",
    "Zh_unknown_S = np.array(data['z'])\n",
    "Zh_unknown_G = np.array(data['z'])\n",
    "\n",
    "Pt2_0_G = np.array(mc['mc_Pt2'])\n",
    "Pt2_0_S = np.array(mc['Pt2'])\n",
    "Pt2_unknown_S = np.array(data['Pt2'])\n",
    "Pt2_unknown_G = np.array(data['Pt2'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0_G=np.c_[(Q2_0_G,Pz_0_G)]\n",
    "theta0_S=np.c_[(Q2_0_S,Pz_0_S)]\n",
    "theta_unknown_S=np.c_[(Q2_unknown_S,Pz_unknown_S)]\n",
    "theta_unknown_G=np.c_[(Q2_unknown_G,Pz_unknown_G)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0_G=np.c_[(Px_0_G,Py_0_G,Pz_0_G,Pid_0_G,Q2_0_G)]\n",
    "theta0_S=np.c_[(Px_0_S,Py_0_S,Pz_0_S,Pid_0_S,Q2_0_S)]\n",
    "theta_unknown_S=np.c_[(Px_unknown_S,Py_unknown_S,Pz_unknown_S,Pid_unknown_S,Q2_unknown_S)]\n",
    "theta_unknown_G=np.c_[(Px_unknown_G,Py_unknown_G,Pz_unknown_G,Pid_unknown_G,Q2_unknown_G)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0_G=np.c_[(Px_0_G,Py_0_G,Pz_0_G,Q2_0_G,Nu_0_G)]\n",
    "theta0_S=np.c_[(Px_0_S,Py_0_S,Pz_0_S,Q2_0_S,Nu_0_S)]\n",
    "theta_unknown_S=np.c_[(Px_unknown_S,Py_unknown_S,Pz_unknown_S,Q2_unknown_S,Nu_unknown_S)]\n",
    "theta_unknown_G=np.c_[(Px_unknown_G,Py_unknown_G,Pz_unknown_G,Q2_unknown_G,Nu_unknown_G)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0_G=np.c_[(Px_0_G,Py_0_G,Pz_0_G,Q2_0_G,Nu_0_G,PhiPQ_0_G,Pt2_0_G)]\n",
    "theta0_S=np.c_[(Px_0_S,Py_0_S,Pz_0_S,Q2_0_S,Nu_0_S,PhiPQ_0_S,Pt2_0_S)]\n",
    "theta_unknown_S=np.c_[(Px_unknown_S,Py_unknown_S,Pz_unknown_S,Q2_unknown_S,Nu_unknown_S,PhiPQ_unknown_S,Pt2_unknown_S)]\n",
    "theta_unknown_G=np.c_[(Px_unknown_G,Py_unknown_G,Pz_unknown_G,Q2_unknown_G,Nu_unknown_G, PhiPQ_unknown_G,Pt2_unknown_G)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0_G=np.c_[(Px_0_G,Py_0_G,Pz_0_G,Q2_0_G,Nu_0_G,PhiPQ_0_G,Pt2_0_G,Pid_0_G)]\n",
    "theta0_S=np.c_[(Px_0_S,Py_0_S,Pz_0_S,Q2_0_S,Nu_0_S,PhiPQ_0_S,Pt2_0_S,Pid_0_S)]\n",
    "theta_unknown_S=np.c_[(Px_unknown_S,Py_unknown_S,Pz_unknown_S,Q2_unknown_S,Nu_unknown_S,PhiPQ_unknown_S,Pt2_unknown_S,Pid_unknown_S)]\n",
    "theta_unknown_G=np.c_[(Px_unknown_G,Py_unknown_G,Pz_unknown_G,Q2_unknown_G,Nu_unknown_G, PhiPQ_unknown_G,Pt2_unknown_G, Pid_unknown_G)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0_G=np.c_[(Px_0_G,Py_0_G,Pz_0_G,Pid_0_G,Q2_0_G,Nu_0_G)]\n",
    "theta0_S=np.c_[(Px_0_S,Py_0_S,Pz_0_S,Pid_0_S,Q2_0_S,Nu_0_S)]\n",
    "theta_unknown_S=np.c_[(Px_unknown_S,Py_unknown_S,Pz_unknown_S,Pid_unknown_S,Q2_unknown_S,Nu_unknown_S)]\n",
    "theta_unknown_G=np.c_[(Px_unknown_G,Py_unknown_G,Pz_unknown_G,Pid_unknown_G,Q2_unknown_G,Nu_unknown_G)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0_G=np.c_[(Px_0_G,Py_0_G,Pz_0_G,Pid_0_G,Q2_0_G,Nu_0_G, PhiPQ_0_G, Zh_0_G)]\n",
    "theta0_S=np.c_[(Px_0_S,Py_0_S,Pz_0_S,Pid_0_S,Q2_0_S,Nu_0_S, PhiPQ_0_S, Zh_0_S)]\n",
    "theta_unknown_S=np.c_[(Px_unknown_S,Py_unknown_S,Pz_unknown_S,Pid_unknown_S,Q2_unknown_S,Nu_unknown_S, PhiPQ_unknown_S, Zh_unknown_S)]\n",
    "theta_unknown_G=np.c_[(Px_unknown_G,Py_unknown_G,Pz_unknown_G,Pid_unknown_G,Q2_unknown_G, Nu_unknown_G,PhiPQ_unknown_G, Zh_unknown_G)]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OmniFold Gaussian Toy Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-2,5, 101)\n",
    "\n",
    "fig, ax = plt.subplots(1,\n",
    "                       2,\n",
    "                       figsize=(12, 5),\n",
    "                       constrained_layout=True,\n",
    "                       sharey=True)\n",
    "\n",
    "ax[0].set_xlabel('$x_0$')\n",
    "ax[0].set_ylabel('Events per bin (normalized)')\n",
    "ax[0].hist(theta0_G[theta0_G[:,0]!=-10][:,0], bins=bins, **plot_style_2, label='Generation')\n",
    "ax[0].hist(theta0_S[theta0_S[:,0]!=-10][:,0], bins=bins, **plot_style_2, label='Simulation')\n",
    "legend = ax[0].legend(\n",
    "    loc='upper left',\n",
    "    frameon=True)\n",
    "plt.setp(legend.get_title(), multialignment='center')\n",
    "\n",
    "ax[1].set_xlabel('$x_0$')\n",
    "ax[1].hist(theta_unknown_G[theta_unknown_G[:,0]!=-10][:,0], bins=bins, **plot_style_2, label='Truth')\n",
    "ax[1].hist(theta_unknown_S[theta_unknown_S[:,0]!=-10][:,0], bins=bins, **plot_style_2, label='Detector')\n",
    "legend = ax[1].legend(\n",
    "    loc='upper left',\n",
    "    frameon=True)\n",
    "plt.setp(legend.get_title(), multialignment='center')\n",
    "\n",
    "fig.show()\n",
    "\n",
    "#bins = np.linspace(-180,180, 101)\n",
    "#bins = np.linspace(0,2, 101)\n",
    "bins = np.linspace(-300,300,4)\n",
    "fig, ax = plt.subplots(1,\n",
    "                       2,\n",
    "                       figsize=(12, 5),\n",
    "                       constrained_layout=True,\n",
    "                       sharey=True)\n",
    "\n",
    "ax[0].set_xlabel('$x_1$')\n",
    "ax[0].set_ylabel('Events per bin (normalized)')\n",
    "ax[0].hist(theta0_G[theta0_G[:,0]!=-10][:,-1], bins=bins, **plot_style_2, label='Generation')\n",
    "ax[0].hist(theta0_S[theta0_S[:,0]!=-10][:,-1], bins=bins, **plot_style_2, label='Simulation')\n",
    "legend = ax[0].legend(\n",
    "    loc='upper left',\n",
    "    frameon=True)\n",
    "plt.setp(legend.get_title(), multialignment='center')\n",
    "\n",
    "ax[1].set_xlabel('$x_1$')\n",
    "ax[1].hist(theta_unknown_G[theta_unknown_G[:,0]!=-10][:,-1], bins=bins, **plot_style_2, label='Truth')\n",
    "ax[1].hist(theta_unknown_S[theta_unknown_S[:,0]!=-10][:,-1], bins=bins, **plot_style_2, label='Detector')\n",
    "legend = ax[1].legend(\n",
    "    loc='upper left',\n",
    "    frameon=True)\n",
    "plt.setp(legend.get_title(), multialignment='center')\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##theta0_S[pass_reco==0][:,0] = -10\n",
    "##theta0_G[pass_truth==0][:,0] = -10\n",
    "#theta_unknown_S[:,0][data['pass_reco']!=0] = -10\n",
    "#theta_unknown_G[:,0][data['pass_reco']!=0] = -10\n",
    "pass_reco = np.array(mc['pass_reco'])\n",
    "pass_truth = np.array(mc['pass_truth'])\n",
    "theta0_S[:,0][pass_reco==0] = -10\n",
    "theta0_G[:,0][pass_truth==0] = -10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## testing stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(mc['pass_reco'])\n",
    "#print(np.array(mc['pass_reco']))\n",
    "#pass_reco = np.array(mc['pass_reco'])\n",
    "#pass_truth = np.array(mc['pass_truth'])\n",
    "#theta0_S[pass_reco==0][:,0] = -10\n",
    "#theta0_G[pass_truth==0][:,0] = -10\n",
    "\n",
    "#print('A' , np.count_nonzero(mc['pass_reco'])) ##these are the ones that do have reco associated to them\n",
    "#print('B', np.count_nonzero(mc['pass_truth'])) ##these are the ones that do have truth associated with them\n",
    "\n",
    "#theta0_G = np.c_[np.random.normal(1.0, 1.2,10),np.random.normal(1.0, 1.2,10)]  # Generator-level synthetic sample\n",
    "#print(theta0_G)\n",
    "##theta0_S = np.c_[np.array([(x + np.random.normal(0, 0.1)) for x in theta0_G[:,0]]),np.array([(x + np.random.normal(0, 0.1)) for x in theta0_G[:,1]])]  # Detector smearing for synthetic sample\n",
    "#print (theta0_S)\n",
    "##pass_reco = np.random.binomial(1,1.-0.3,len(theta0_G))\n",
    "#pass_truth = np.random.binomial(1,1.-0.0,len(theta0_G))\n",
    "#print(pass_truth )\n",
    "#print(pass_reco )\n",
    "#theta0_S[:,0]\n",
    "#theta0_S[:,0][pass_reco==0] = -10\n",
    "#theta0_G[:,0][pass_truth==0] = -10\n",
    "#theta0_S[:,0]\n",
    "#theta0_G[:,0]\n",
    "#theta0_G[theta0_S[:,0]!=-10][:,1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unfold Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0 = np.stack([theta0_G, theta0_S], axis=1)\n",
    "labels0 = np.zeros(len(theta0))\n",
    "\n",
    "\n",
    "theta_unknown = np.stack([theta_unknown_G, theta_unknown_S], axis=1)\n",
    "labels_unknown = np.ones(len(theta_unknown))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterations = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_observables= 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xvals_1 = np.concatenate((theta0_S, theta_unknown_S))\n",
    "yvals_1 = np.concatenate((labels0, labels_unknown))\n",
    "\n",
    "xvals_2 = np.concatenate((theta0_G, theta0_G))\n",
    "yvals_2 = np.concatenate((labels0, labels_unknown))\n",
    "\n",
    "weights = np.empty(shape=(iterations, 2, len(theta0)))\n",
    "# shape = (iteration, step, event)\n",
    "\n",
    "inputs = Input((num_observables, ))\n",
    "#inputs = Input((2, ))\n",
    "hidden_layer_1 = Dense(50, activation='relu')(inputs)\n",
    "hidden_layer_2 = Dense(50, activation='relu')(hidden_layer_1)\n",
    "hidden_layer_3 = Dense(50, activation='relu')(hidden_layer_2)\n",
    "outputs = Dense(1, activation='sigmoid')(hidden_layer_3)\n",
    "\n",
    "model = Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "earlystopping = EarlyStopping(patience=10,\n",
    "                              verbose=1,\n",
    "                              restore_best_weights=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from NN (DCTR)\n",
    "def reweight(events):\n",
    "    f = model.predict(events, batch_size=10000)\n",
    "    weights = f / (1. - f)\n",
    "    return np.squeeze(np.nan_to_num(weights))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial iterative weights are ones\n",
    "\n",
    "weights_pull = np.ones(len(theta0_S))\n",
    "weights_push = np.ones(len(theta0_S))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K.clear_session()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(iterations):\n",
    "    print(\"\\nITERATION: {}\\n\".format(i + 1))\n",
    "\n",
    "    # STEP 1: classify Sim. (which is reweighted by weights_push) to Data\n",
    "    # weights reweighted Sim. --> Data\n",
    "    print(\"STEP 1\\n\")\n",
    "\n",
    "    weights_1 = np.concatenate((weights_push, np.ones(len(theta_unknown_S))))\n",
    "    # actual weights for Sim., ones for Data (not MC weights)\n",
    "\n",
    "    X_train_1, X_test_1, Y_train_1, Y_test_1, w_train_1, w_test_1 = train_test_split(\n",
    "        xvals_1, yvals_1, weights_1)\n",
    "\n",
    "    #theta0_G[theta0_S[:,0]!=-10][:,1]\n",
    "    model.compile(loss='binary_crossentropy',\n",
    "                  optimizer='Adam',\n",
    "                  metrics=['accuracy'])\n",
    "    model.fit(X_train_1[X_train_1[:,0]!=-10],\n",
    "              Y_train_1[X_train_1[:,0]!=-10],\n",
    "              sample_weight=w_train_1[X_train_1[:,0]!=-10],\n",
    "              epochs=200,\n",
    "              batch_size=10000,\n",
    "              validation_data=(X_test_1[X_test_1[:,0]!=-10], Y_test_1[X_test_1[:,0]!=-10], w_test_1[X_test_1[:,0]!=-10]),\n",
    "              callbacks=[earlystopping],\n",
    "              verbose=1)\n",
    "\n",
    "    weights_pull = weights_push * reweight(theta0_S)\n",
    "    weights_pull[theta0_S[:,0]==-10] = 1. #these are events that don't pass reco; take the prior.\n",
    "    weights[i, :1, :] = weights_pull\n",
    "\n",
    "    # STEP 2: classify Gen. to reweighted Gen. (which is reweighted by weights_pull)\n",
    "    # weights Gen. --> reweighted Gen.\n",
    "    print(\"\\nSTEP 2\\n\")\n",
    "\n",
    "    weights_2 = np.concatenate((np.ones(len(theta0_G)), weights_pull))\n",
    "    # ones for Gen. (not MC weights), actual weights for (reweighted) Gen.\n",
    "\n",
    "    X_train_2, X_test_2, Y_train_2, Y_test_2, w_train_2, w_test_2 = train_test_split(\n",
    "        xvals_2, yvals_2, weights_2)\n",
    "\n",
    "    model.compile(loss='binary_crossentropy',\n",
    "                  optimizer='Adam',\n",
    "                  metrics=['accuracy'])\n",
    "    model.fit(X_train_2,\n",
    "              Y_train_2,\n",
    "              sample_weight=w_train_2,\n",
    "              epochs=200,\n",
    "              batch_size=10000,\n",
    "              validation_data=(X_test_2, Y_test_2, w_test_2),\n",
    "              callbacks=[earlystopping],\n",
    "              verbose=1)\n",
    "\n",
    "    weights_push = reweight(theta0_G)\n",
    "    weights[i, 1:2, :] = weights_push"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot data and MC reweighted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## This is what in Ben's notebook. \n",
    "## The _S is RECO level. \n",
    "## You require that the RECO dataframe has a valid pass reco.\n",
    "## The G is the truth level\n",
    "## You required that the truth level has pass_truth valid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(weights)):\n",
    "\n",
    "    print(\"ITERATION: {}\".format(i + 1))\n",
    "    bins = np.linspace(0.0, 4.0, 40)\n",
    "    \n",
    "    fig, ax = plt.subplots(3,\n",
    "                           3,\n",
    "                           figsize=(16, 16),\n",
    "                           constrained_layout=True)\n",
    "    ax[0,0].set_xlabel(r'$Q^{2}$')\n",
    "    ax[0,0].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[0,0].hist(Q2_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[0,0].hist(Q2_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[0,0].hist(Q2_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    ax[0,0].legend(frameon=False,loc='best')\n",
    "    \n",
    "    \n",
    "    bins = np.linspace(2.0, 5, 20)\n",
    "\n",
    "    ax[0,1].set_xlabel(r'$\\nu$')\n",
    "    ax[0,1].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[0,1].hist(Nu_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[0,1].hist(Nu_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[0,1].hist(Nu_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    \n",
    "    \n",
    "    \n",
    "    bins = np.linspace(-1.5, 1.5, 50)\n",
    "\n",
    "    ax[0,2].set_xlabel(r'$P_{x}$')\n",
    "    ax[0,2].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[0,2].hist(Px_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[0,2].hist(Px_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[0,2].hist(Px_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    \n",
    "    \n",
    "    bins = np.linspace(-1.5, 1.5, 20)\n",
    "\n",
    "    ax[1,0].set_xlabel(r'$P_{y}$')\n",
    "    ax[1,0].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[1,0].hist(Py_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[1,0].hist(Py_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[1,0].hist(Py_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    \n",
    "    bins = np.linspace(-0.5, 3.0, 50)\n",
    "\n",
    "    ax[1,1].set_xlabel(r'$P_{z}$')\n",
    "    ax[1,1].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[1,1].hist(Pz_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[1,1].hist(Pz_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[1,1].hist(Pz_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    \n",
    "    \n",
    "    bins = np.linspace(-180, 180, 20)\n",
    "\n",
    "    ax[1,2].set_xlabel(r'$\\phi_{PQ}$')\n",
    "    ax[1,2].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[1,2].hist(PhiPQ_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[1,2].hist(PhiPQ_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[1,2].hist(PhiPQ_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    \n",
    "    bins = np.linspace(0, 1.2, 50)\n",
    "\n",
    "    ax[2,0].set_xlabel(r'$z$')\n",
    "    ax[2,0].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[2,0].hist(Zh_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[2,0].hist(Zh_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[2,0].hist(Zh_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    \n",
    "    \n",
    "    bins = np.linspace(0, 2, 50)\n",
    "\n",
    "    ax[2,1].set_xlabel(r'$p_{T}^{2}$')\n",
    "    ax[2,1].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[2,1].hist(Pt2_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[2,1].hist(Pt2_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[2,1].hist(Pt2_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    \n",
    "    bins = np.linspace(-300, 300, 4)\n",
    "\n",
    "    ax[2,2].set_xlabel('PID')\n",
    "    ax[2,2].set_ylabel('Events per bin (Normalized)')\n",
    "    ax[2,2].hist(Pid_0_S[theta0_S[:,0]!=0], bins=bins, label=r'MC .', **plot_style_2)\n",
    "    ax[2,2].hist(Pid_0_S[theta0_S[:,0]!=0],\n",
    "               bins=bins,\n",
    "               weights=weights[i, 0, :][theta0_S[:,0]!=0],\n",
    "               label=r'MC . DCTR wgt.',\n",
    "               **plot_style_1)\n",
    "\n",
    "    ax[2,2].hist(Pid_unknown_S,\n",
    "               bins=bins,\n",
    "               label=r'Data (reco)',\n",
    "               **plot_style_2)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Put weights obtained with the DNN back into the mc dataframe:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mc['weight_reco'] = weights[-1, 0, :]\n",
    "mc['weight_truth'] = weights[-1, 1, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transform dataframe back into ROOT file "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import root_pandas\n",
    "from root_pandas import read_root\n",
    "from root_pandas import to_root "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "to_root(mc,'D_weighted.root', key='ntuple_sim')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute efficiency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,2,figsize=(16,5),sharey=True)\n",
    "\n",
    "mask = np.logical_and(theta0_S[:,0]!=-10,theta0_G[:,0]!=-10)\n",
    "\n",
    "bins = np.linspace(0,1.0,40)\n",
    "axs[0].set_title('Non weighted')\n",
    "axs[1].set_title('Weighted')\n",
    "#axs[0].hist(Zh_0_G[theta0_S[:,0]==-10],alpha=0.3,bins=100, label='all generated')\n",
    "axs[0].hist(Zh_0_G[theta0_G[:,0]!=-10],alpha=0.3,bins=bins, label='all generated')\n",
    "axs[0].hist(Zh_0_G[mask], alpha=0.3,bins=bins, label='pass reco')\n",
    "#axs[1].hist(Zh_0_G[theta0_S[:,0]==-10],alpha=0.3,bins=100,weights=weights[-1, 1, :][theta0_S[:,0]==-10],label='all generated')\n",
    "axs[1].hist(Zh_0_G[theta0_G[:,0]!=-10],alpha=0.3,bins=bins,weights=weights[-1, 1, :][theta0_G[:,0]!=-10],label='all generated')\n",
    "axs[1].hist(Zh_0_G[mask],alpha=0.3,bins=bins,weights=weights[-1, 1, :][mask],label='pass reco')\n",
    "\n",
    "\n",
    "axs[0].set(xlabel=r\"$z$\", ylabel=\"Events\")\n",
    "axs[1].set(xlabel=r\"$z$\", ylabel=\"Events\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#numerator\n",
    "mask = np.logical_and(theta0_S[:,0]!=-10,theta0_G[:,0]!=-10)\n",
    "print (mask) \n",
    "num,x,_  = plt.hist(Zh_0_G[mask], weights=weights[-1, 1, :][mask],bins=20,label='numerator, weighted',alpha=0.3)\n",
    "den,x,_ = plt.hist(Zh_0_G[theta0_G[:,0]!=-10], weights=weights[-1, 1, :][theta0_G[:,0]!=-10],bins=20,label='numerator, weighted',alpha=0.3)\n",
    "x= (x[1:] + x[:-1])/2.0\n",
    "\n",
    "#plt.hist(theta0_G[theta0_S[:,0]!=-10][:,1])\n",
    "plt.show()\n",
    "eff_weighted= np.true_divide(num,den)\n",
    "print(sum(num))\n",
    "print(sum(den))\n",
    "print(sum(num)/sum(den))\n",
    "#plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot efficiency with no weighting (i.e. out of the box MC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#numerator\n",
    "mask = np.logical_and(theta0_S[:,0]!=-10,theta0_G[:,0]!=-10)\n",
    "print (mask) \n",
    "num,x,_  = plt.hist(Zh_0_G[mask], bins=20,label='numerator, weighted',alpha=0.3)\n",
    "den,x,_ = plt.hist(Zh_0_G[theta0_G[:,0]!=-10], bins=20,label='numerator, weighted',alpha=0.3)\n",
    "x= (x[1:] + x[:-1])/2.0\n",
    "\n",
    "#plt.hist(theta0_G[theta0_S[:,0]!=-10][:,1])\n",
    "plt.show()\n",
    "eff= np.true_divide(num,den)\n",
    "print(sum(num))\n",
    "print(sum(den))\n",
    "print(sum(num)/sum(den))\n",
    "#plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x,eff,label='MC')\n",
    "plt.plot(x,eff_weighted, label='MC weighted')\n",
    "\n",
    "plt.ylabel('Acceptance')\n",
    "plt.xlabel('z')\n",
    "plt.legend(loc='best',frameon=False)\n",
    "plt.ylim([0.,0.5])\n",
    "plt.xlim([0.,1.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##Plot the unfolded distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-2, 2, 101)\n",
    "\n",
    "for i in range(1):\n",
    "    print(\"ITERATION {}:\".format(i + 1))\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    #fig, (ax1, ax2) = plt.subplots(1, 2)\n",
    "    \n",
    "    gs = plt.GridSpec(1, 2)\n",
    "    gs.update(left=0.08, right=0.925,\n",
    "              top=0.95, bottom=0.05,\n",
    "              hspace=0.3, wspace=0.5)\n",
    "\n",
    "    ax1 = plt.subplot(gs[0, 0])\n",
    "    ax2 = plt.subplot(gs[0, 1])     \n",
    "    \n",
    "    hist0 = ax1.hist(theta0_G[theta0_G[:,0]!=-10][:,0],\n",
    "                     bins=bins,\n",
    "                     **plot_style_2,label='Gen')\n",
    "\n",
    "    hist1 = ax1.hist(\n",
    "        theta0_G[theta0_G[:,0]!=-10][:,0],\n",
    "        bins=bins,\n",
    "        label='Gen weighted',\n",
    "        weights=weights[i, 1, :][theta0_G[:,0]!=-10],\n",
    "        **plot_style_1)\n",
    "\n",
    "    ax1.legend(frameon=False)\n",
    "    ax1.set(xlabel=r\"$x_{1,G}$\", ylabel=\"Events per bin (normalized)\")\n",
    "    #ax1.set_ylim([0,0.6])\n",
    "    bins = np.linspace(-2, 5, 101)\n",
    "\n",
    "    hist0 = ax2.hist(theta0_G[theta0_G[:,0]!=-10][:,2],\n",
    "                     bins=bins,\n",
    "                     **plot_style_2)\n",
    "\n",
    "    hist1 = ax2.hist(\n",
    "        theta0_G[theta0_G[:,0]!=-10][:,2],\n",
    "        bins=bins,\n",
    "        weights=weights[i, 1, :][theta0_G[:,0]!=-10],\n",
    "        **plot_style_1)\n",
    "\n",
    "    ax2.legend(frameon=False)\n",
    "    ax2.set(xlabel=r\"$x_{2,G}$\", ylabel=\"Events per bin (normalized)\")\n",
    "    #ax2.set_ylim([0,0.6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(weights[-1, 0, :]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(weights[i, 1, :]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(mc))"
   ]
  }
 ],
 "metadata": {
  "gist": {
   "data": {
    "description": "OmniFoldAcceptance2D",
    "public": false
   },
   "id": ""
  },
  "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.8"
  },
  "notify_time": "0",
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}