{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hep/davide/miniconda3/envs/root_env/lib/python2.7/site-packages/root_numpy/_tree.py:5: RuntimeWarning: numpy.dtype size changed, may indicate binary incompatibility\n",
      "  from . import _librootnumpy\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": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "file_name='B2Dmunu_full'\n",
    "file_path='/disk/lhcb_data/davide/HCAL_project_full_event/'+file_name+'.root'\n",
    "tree_name='Bd2Dmu/DecayTree'\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "b=r.TBrowser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = r.TFile(file_path)\n",
    "t = f.Get(tree_name)\n",
    "j=t.GetEntries()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size=20000\n",
    "#n_batches=2\n",
    "\n",
    "N = j\n",
    "n_batches= N//batch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "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",
    "realETs_dict={}\n",
    "particle='pi1'\n",
    "\n",
    "for j in range(n_batches):\n",
    "    realETs_dict[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",
    "    \n",
    "regions_dict={}\n",
    "particle='pi1'\n",
    "\n",
    "for j in range(n_batches):\n",
    "    regions_dict[j]=rn.root2array(\n",
    "        filenames=file_path, \n",
    "        treename=tree_name,\n",
    "        branches=[particle+'_L0Calo_HCAL_region'],\n",
    "        start=j*batch_size,\n",
    "        stop=(j+1)*batch_size,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_events={\n",
    "            'xProjections':xProjections_dict, \n",
    "            'yProjections':yProjections_dict, \n",
    "            'realETs':realETs_dict,\n",
    "            'regions':regions_dict,\n",
    "            }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(true_events['xProjections'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "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 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}