05-rdataframe-collections.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6c20dcc6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Working with collections and object selections\n",
    "\n",
    "- RDataFrame reads collections as the special type [ROOT::RVec](https://root.cern/doc/master/classROOT_1_1VecOps_1_1RVec.html) - e.g. a branch containing an array of floating point numbers can be read as a `ROOT::RVec<float>`.\n",
    "\n",
    "- C-style arrays (with variable or static size), `std::vectors` and many other collection types can be read this way. \n",
    "\n",
    "- When reading ROOT data, column values of type `ROOT::RVec<T>` perform no copy of the underlying array.\n",
    "\n",
    "- `RVec` is a container similar to `std::vector` (and can be used just like a `std::vector`) but it also offers a rich interface to operate on the array elements in a vectorised fashion, similarly to Python's NumPy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25b99067",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "import ROOT\n",
    "\n",
    "treename = \"myDataset\"\n",
    "filename = \"data/collections_dataset.root\"\n",
    "df = ROOT.RDataFrame(treename, filename)\n",
    "\n",
    "print(f\"Columns in the dataset: {df.GetColumnNames()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f3633ae",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "To quickly inspect the data we can export it as a dictionary of `numpy` arrays thanks to the `AsNumpy` RDataFrame method. \n",
    "\n",
    "Note that for each row, `E` is an array of values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4664f42",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "npy_dict = df.AsNumpy([\"E\"])\n",
    "\n",
    "for row, vec in enumerate(npy_dict[\"E\"]):\n",
    "    print(f\"\\nRow {row} contains:\\n{vec}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b896f202",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Define a new column with operations on RVecs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f140a78",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "df1 = df.Define(\"good_pt\", \"sqrt(px*px + py*py)[E>100]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ca3b396",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "`sqrt(px*px + py*py)[E>100]`:\n",
    "- `px`, `py` and `E` are the columns, the elements of those columns are `RVec`s\n",
    "\n",
    "- Operations on `RVec`s, such as sum, product, sqrt, preserve the dimensionality of the array\n",
    "\n",
    "- `[E>100]` selects the elements of the array that satisfy the condition\n",
    "\n",
    "- `E > 100`: boolean expressions on `RVec`s such as `E > 100` return a mask, that is an array with information which values pass the selection (e.g. `[0, 1, 0, 0]` if only the second element satisfies the condition)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e283340e",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Now we can plot the newly defined column values in a histogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9795d4a7",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "c = ROOT.TCanvas()\n",
    "h = df1.Histo1D((\"pt\", \"pt\", 16, 0, 4), \"good_pt\")\n",
    "h.Draw()\n",
    "c.Draw()"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}