root_file_reading.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:59:30.436876Z",
     "start_time": "2020-10-26T15:59:29.380548Z"
    }
   },
   "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": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:59:39.298642Z",
     "start_time": "2020-10-26T15:59:39.211317Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 68.2 ms, sys: 0 ns, total: 68.2 ms\n",
      "Wall time: 67 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "nparray = root2array(\n",
    "                ['output_9.root'],#*10,\n",
    "                treename = 'deepntuplizer/tree', \n",
    "                branches = ['fj_pt', 'fj_eta'],\n",
    "                )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uproot3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T15:59:40.226469Z",
     "start_time": "2020-10-26T15:59:39.851857Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 354 ms, sys: 15 ms, total: 369 ms\n",
      "Wall time: 361 ms\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": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-26T16:00:12.778585Z",
     "start_time": "2020-10-26T16:00:12.172751Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 623 ms, sys: 2.17 ms, total: 625 ms\n",
      "Wall time: 580 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "filenames = ['output_9.root']#*10\n",
    "treename = 'deepntuplizer/tree'\n",
    "branches = ['fj_pt', 'fj_eta']\n",
    "formats = ['f8','f8']\n",
    "\n",
    "arrays = uproot4.concatenate(filenames, branches, library='np')\n",
    "\n",
    "big_array4 = np.rec.fromarrays(list(arrays.values()), names=branches, formats=formats)"
   ]
  },
  {
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}