{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Welcome to JupyROOT 6.16/00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/disk/lhcb_data2/davide/miniconda3/envs/root_env/lib/python3.7/site-packages/root_numpy/__init__.py:46: RuntimeWarning: numpy 1.17.2 is currently installed but you installed root_numpy against numpy 1.9.3. Please consider reinstalling root_numpy for this numpy version.\n",
      "  RuntimeWarning)\n"
     ]
    }
   ],
   "source": [
    "import root_numpy as rn\n",
    "#import pandas as pd\n",
    "import numpy as np\n",
    "import ROOT as r\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "file_name='B2Dmunu'\n",
    "file_path='/disk/lhcb_data2/davide/HCAL_project_full_event/'+file_name+'.root'\n",
    "tree_name='Bd2Dmu/DecayTree'\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#b=r.TBrowser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = r.TFile(file_path)\n",
    "t = f.Get(tree_name)\n",
    "N=t.GetEntries()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size=20000\n",
    "\n",
    "n_batches= N//batch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "93510"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "pion1_X_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_X_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi1_L0Calo_HCAL_xProjection'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion1_X = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_X[j]= np.array([pion1_X_temp[j][i][0] for i in range(pion1_X_temp[j].shape[0])])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "pion1_X_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_X_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi1_L0Calo_HCAL_xProjection'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion1_X = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_X[j]= np.array([pion1_X_temp[j][i][0] for i in range(pion1_X_temp[j].shape[0])])\n",
    "    \n",
    "    \n",
    "pion1_Y_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_Y_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi1_L0Calo_HCAL_yProjection'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion1_Y = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_Y[j]= np.array([pion1_Y_temp[j][i][0] for i in range(pion1_Y_temp[j].shape[0])])\n",
    "\n",
    "pion1_realET_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_realET_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi1_L0Calo_HCAL_realET'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion1_realET = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_realET[j]= np.array([pion1_realET_temp[j][i][0] for i in range(pion1_realET_temp[j].shape[0])])\n",
    "    \n",
    "pion1_region_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_region_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi1_L0Calo_HCAL_region'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion1_region = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion1_region[j]= np.array([pion1_region_temp[j][i][0] for i in range(pion1_region_temp[j].shape[0])])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "pion2_X_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion2_X_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi2_L0Calo_HCAL_xProjection'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion2_X = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion2_X[j]= np.array([pion2_X_temp[j][i][0] for i in range(pion2_X_temp[j].shape[0])])\n",
    "    \n",
    "    \n",
    "pion2_Y_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion2_Y_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi2_L0Calo_HCAL_yProjection'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion2_Y = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion2_Y[j]= np.array([pion2_Y_temp[j][i][0] for i in range(pion2_Y_temp[j].shape[0])])\n",
    "\n",
    "pion2_realET_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion2_realET_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi2_L0Calo_HCAL_realET'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion2_realET = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion2_realET[j]= np.array([pion2_realET_temp[j][i][0] for i in range(pion2_realET_temp[j].shape[0])])\n",
    "    \n",
    "pion2_region_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion2_region_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['pi2_L0Calo_HCAL_region'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "pion2_region = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    pion2_region[j]= np.array([pion2_region_temp[j][i][0] for i in range(pion2_region_temp[j].shape[0])])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_X_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    K_X_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['K_L0Calo_HCAL_xProjection'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "K_X = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    K_X[j]= np.array([K_X_temp[j][i][0] for i in range(K_X_temp[j].shape[0])])\n",
    "    \n",
    "    \n",
    "K_Y_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    K_Y_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['K_L0Calo_HCAL_yProjection'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "K_Y = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    K_Y[j]= np.array([K_Y_temp[j][i][0] for i in range(K_Y_temp[j].shape[0])])\n",
    "\n",
    "K_realET_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    K_realET_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['K_L0Calo_HCAL_realET'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "K_realET = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    K_realET[j]= np.array([K_realET_temp[j][i][0] for i in range(K_realET_temp[j].shape[0])])\n",
    "    \n",
    "K_region_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    K_region_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['K_L0Calo_HCAL_region'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "K_region = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    K_region[j]= np.array([K_region_temp[j][i][0] for i in range(K_region_temp[j].shape[0])])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "L0Hadron_TIS_dict_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    L0Hadron_TIS_dict_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['B0_L0HadronDecision_TIS'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "L0Hadron_TIS_dict = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    L0Hadron_TIS_dict[j]= np.array([L0Hadron_TIS_dict_temp[j][i][0] for i in range(L0Hadron_TIS_dict_temp[j].shape[0])])\n",
    "\n",
    "L0Hadron_TOS_dict_temp={}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    L0Hadron_TOS_dict_temp[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=['B0_L0HadronDecision_TOS'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "L0Hadron_TOS_dict = {}\n",
    "\n",
    "for j in range(n_batches):\n",
    "    L0Hadron_TOS_dict[j]= np.array([L0Hadron_TOS_dict_temp[j][i][0] for i in range(L0Hadron_TOS_dict_temp[j].shape[0])])\n",
    "    \n",
    " \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "xProjections_dict={}\n",
    "particle='pi1'\n",
    "\n",
    "for j in range(n_batches):\n",
    "    xProjections_dict[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=[particle+'_L0Calo_HCAL_xProjections'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "yProjections_dict={}\n",
    "particle='pi1'\n",
    "\n",
    "for j in range(n_batches):\n",
    "    yProjections_dict[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=[particle+'_L0Calo_HCAL_yProjections'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n",
    "    \n",
    "\n",
    "realETs={}\n",
    "particle='pi1'\n",
    "\n",
    "for j in range(n_batches):\n",
    "    realETs[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=[particle+'_L0Calo_HCAL_realETs'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_events={\n",
    "            'xProjections':xProjections_dict, \n",
    "            'yProjections':yProjections_dict, \n",
    "            'realETs': realETs,\n",
    "    \n",
    "            'L0HadronDec_TIS':L0Hadron_TIS_dict,\n",
    "            'L0HadronDec_TOS':L0Hadron_TOS_dict, \n",
    "    \n",
    "            'pi1_xProjection':pion1_X,\n",
    "            'pi1_yProjection':pion1_Y,\n",
    "            'pi1_realET':pion1_realET,\n",
    "            'pi1_region':pion1_region,\n",
    "    \n",
    "            'pi2_xProjection':pion2_X,\n",
    "            'pi2_yProjection':pion2_Y,\n",
    "            'pi2_realET':pion2_realET,\n",
    "            'pi2_region':pion2_region,\n",
    "   \n",
    "            'K_xProjection':K_X,\n",
    "            'K_yProjection':K_Y,\n",
    "            'K_realET':K_realET,\n",
    "            'K_region':K_region,\n",
    "    \n",
    "            }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(true_events['xProjections'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/disk/lhcb_data/davide/HCAL_project_full_event/csv/MCtracker_info.pickle', 'wb') as f:\n",
    "    pickle.dump(true_events, f)\n"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}