root_file_reading.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:24:31.955984Z",
     "start_time": "2020-10-26T15:24:30.942205Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Welcome to JupyROOT 6.18/04\n"
     ]
    }
   ],
   "source": [
    "import uproot4\n",
    "import uproot\n",
    "import numpy as np\n",
    "from root_numpy import tree2array, root2array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:23:57.030791Z",
     "start_time": "2020-10-26T15:23:57.027885Z"
    }
   },
   "source": [
    "## Baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:24:32.702377Z",
     "start_time": "2020-10-26T15:24:31.958559Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 719 ms, sys: 22.7 ms, total: 742 ms\n",
      "Wall time: 733 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "nparray = root2array(\n",
    "                ['output_9.root']*10,\n",
    "                treename = 'deepntuplizer/tree', \n",
    "                branches = ['fj_pt', 'fj_eta'],#, 'fj_sdmass'],\n",
    "                )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uproot3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:24:35.541641Z",
     "start_time": "2020-10-26T15:24:32.704957Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.83 s, sys: 33.4 ms, total: 2.86 s\n",
      "Wall time: 2.82 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "filenames = ['output_9.root']*10\n",
    "treename = 'deepntuplizer/tree'\n",
    "branches = ['fj_pt', 'fj_eta']\n",
    "\n",
    "totsize = np.sum([uproot.open(file)[treename].numentries for file in filenames])\n",
    "\n",
    "tree = uproot.open(filenames[0])[treename]\n",
    "dtypes = []\n",
    "for b in branches:\n",
    "    dtypes.append((b, str(tree[b].interpretation).split(\"'\")[1]))\n",
    "    \n",
    "big_array = np.empty(totsize, dtype=dtypes)\n",
    "\n",
    "last_entry = {name: 0 for name in branches}\n",
    "for file in filenames:\n",
    "    tree = uproot.open(file)['deepntuplizer/tree']\n",
    "    for name in branches:\n",
    "        branch = tree[name]\n",
    "        for basket in branch.iterate_baskets():\n",
    "            start = last_entry[name]\n",
    "            stop = start + len(basket)\n",
    "            big_array[name][start:stop] = basket\n",
    "            last_entry[name] = stop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:24:16.865289Z",
     "start_time": "2020-10-26T15:24:16.862626Z"
    }
   },
   "source": [
    "## Uproot 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:24:50.600566Z",
     "start_time": "2020-10-26T15:24:35.544208Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 15.3 s, sys: 290 ms, total: 15.5 s\n",
      "Wall time: 15 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "filenames = ['output_9.root']*10\n",
    "treename = 'deepntuplizer/tree'\n",
    "branches = ['fj_pt', 'fj_eta']\n",
    "\n",
    "totsize = np.sum([uproot4.open(file)[treename].num_entries for file in filenames])\n",
    "\n",
    "tree = uproot.open(filenames[0])[treename]\n",
    "dtypes = []\n",
    "for b in branches:\n",
    "    dtypes.append((b, str(tree[b].interpretation).split(\"'\")[1]))\n",
    "    \n",
    "big_array = np.empty(totsize, dtype=dtypes)\n",
    "\n",
    "last_entry = {name: 0 for name in branches}\n",
    "for file in filenames:\n",
    "    tree = uproot4.open(file)['deepntuplizer/tree']\n",
    "    for name in branches:\n",
    "        branch = tree[name]\n",
    "        for basket in branch.iterate():\n",
    "            start = last_entry[name]\n",
    "            stop = start + len(basket)\n",
    "            big_array[name][start:stop] = basket\n",
    "            last_entry[name] = stop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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
}