vertical.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "06e249a0-a8e2-4def-b4ca-d713d6744fb5",
   "metadata": {},
   "source": [
    "# Vertical scaling of columnar analysis: Numba and C++ integration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88373c9c-16fd-4b66-be4a-4e91fbc50ef6",
   "metadata": {},
   "source": [
    "(no-frills notebook)\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07942888-3bc5-42b2-b210-8f5d1c477937",
   "metadata": {},
   "source": [
    "## Motivation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c5df422-ca87-4ace-808f-acb9e0cbe6d7",
   "metadata": {},
   "source": [
    "There are two kinds of languages, and Python is one of the slow kind.\n",
    "\n",
    "<img src=\"img/benchmark-games-2023.svg\" width=\"100%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebd398e1-bb29-4d48-99ad-eb54ecc01ca4",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79cd37c9-8fe6-4beb-a098-f0a219d6807c",
   "metadata": {},
   "source": [
    "But there are a lot of ways to connect Python with compiled languages, usually C++ (increasingly, Rust).\n",
    "\n",
    "<img src=\"img/history-of-bindings-2.svg\" width=\"100%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "480816ba-bc9a-41f7-80fa-ab77c6cc68be",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bd0218c-8684-4ae9-acb4-1b76b671917e",
   "metadata": {},
   "source": [
    "## Setting up an example problem"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ee0c9ea-26aa-41ec-9652-26ef594ecc2a",
   "metadata": {},
   "source": [
    "First, let's get a biggish dataset..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32ca9b1e-bdf7-4e9a-bc0d-831f3df0f79e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import awkward as ak\n",
    "import vector\n",
    "vector.register_awkward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a12c3857-1396-43c7-98d0-7feb3b087df2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "events = ak.from_parquet(\"data/SMHiggsToZZTo4L.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fc90a60-39c1-4b98-a62c-126b60c7c22a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "events.muon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9846b3ec-0f5d-4b1e-ac64-035fafc4126f",
   "metadata": {},
   "source": [
    "...and then let's give it some work to do. How about dimuon pairs? (The \"hello world\" of HEP.)\n",
    "\n",
    "This uses the columnar methods from last week (CoDaS-HEP)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "709707bb-8f54-4aa7-b72d-aa8fb5462b62",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "muplus = events.muon[events.muon.charge > 0]\n",
    "muminus = events.muon[events.muon.charge < 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e12ee04-c80c-42e4-89f1-27b9935ebc9c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mu1, mu2 = ak.unzip(ak.cartesian((muplus, muminus)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acca539e-526e-4b62-ba7f-771546350a9e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.hist(ak.ravel(\n",
    "    \n",
    "    (mu1 + mu2).mass\n",
    "\n",
    "), bins=100, range=(0, 100));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d0e305c-ba75-4b07-b4dd-b8e2d05d7393",
   "metadata": {},
   "source": [
    "It returns the right answer, but maybe we don't like the performance. The sum\n",
    "\n",
    "```python\n",
    "(mu1 + mu2)\n",
    "```\n",
    "\n",
    "converts cylindrical coordinates into Cartesian, for all four coordinates, and then\n",
    "\n",
    "```python\n",
    "(mu1 + mu2).mass\n",
    "```\n",
    "\n",
    "computes $m = \\sqrt{E^2 - {p_x}^2 - {p_y}^2 - {p_z}^2}$ by computing arrays for each intermediate step. We can do better than that by\n",
    "\n",
    "- not computing unnecessary coordinates\n",
    "- not creating and filling unnecessary arrays.\n",
    "\n",
    "We can do the first part with another formula (which ignores the mass of the muon):\n",
    "\n",
    "$$m \\approx \\sqrt{2 \\, {p_T}_1 \\, {p_T}_2 \\, \\left(\\cosh\\left({\\eta}_1 - {\\eta}_2\\right) - \\cos\\left({\\phi}_1 - {\\phi}_2\\right) \\right)}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6099b632-4af8-4d2c-aad9-22ee2706d55c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fast_mass(mu1, mu2):\n",
    "    return np.sqrt(\n",
    "        2 * mu1.pt * mu2.pt * (np.cosh(mu1.eta - mu2.eta) - np.cos(mu1.phi - mu2.phi))\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6d368ba-a4d0-4ab9-a93a-a0fed6c77a14",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.hist(ak.ravel(\n",
    "    \n",
    "    fast_mass(mu1, mu2)\n",
    "\n",
    "), bins=100, range=(0, 100));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "358bc7fe-881e-4991-a495-a7b8517c537b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fast_mass_pedantic(mu1, mu2):\n",
    "    tmp1 = mu1.eta - mu2.eta\n",
    "    tmp2 = np.cosh(tmp1)\n",
    "    del tmp1\n",
    "    tmp3 = mu1.phi - mu2.phi\n",
    "    tmp4 = np.cos(tmp3)\n",
    "    del tmp3\n",
    "    tmp5 = tmp2 - tmp4\n",
    "    del tmp2\n",
    "    del tmp4\n",
    "    tmp6 = mu2.pt * tmp5\n",
    "    del tmp5\n",
    "    tmp7 = mu1.pt * tmp6\n",
    "    del tmp6\n",
    "    tmp8 = 2 * tmp7\n",
    "    del tmp7\n",
    "    tmp9 = np.sqrt(tmp8)\n",
    "    del tmp8\n",
    "    return tmp9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72066e9d-662f-4b31-a322-f7a54c2b547a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.hist(ak.ravel(\n",
    "    \n",
    "    fast_mass_pedantic(mu1, mu2)\n",
    "\n",
    "), bins=100, range=(0, 100));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b9e1dd4-57df-4d0c-9485-6ef2e5159048",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "fast_mass(mu1, mu2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da46fa22-5b12-498b-b555-df980e49c852",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "fast_mass_pedantic(mu1, mu2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f66d12d0-6c6b-4bf7-99e8-814f8956a62a",
   "metadata": {},
   "source": [
    "There are lots of reasons why allocating new arrays and copying results to them for each _step_ in a formula is slower than iterating over the input and computing the _whole_ formula before moving on to the next input.\n",
    "\n",
    "- Allocating memory is expensive (a _search_ through fragmented RAM).\n",
    "- Accessing new areas of RAM is expensive because it bypasses the CPU caching mechanism.\n",
    "- Most mathematical calculations are faster than fetching data from RAM (into CPU caches)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d016d53-0fa1-460e-93f9-2199aa18386e",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0714215b-88e5-484c-8810-c696ac40e145",
   "metadata": {},
   "source": [
    "## The example problem in C++"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "152bee11-e24e-4658-9710-862ac9a68ee2",
   "metadata": {},
   "source": [
    "See [this minicourse](https://github.com/henryiii/python-compiled-minicourse) (from [iscinumpy.dev](https://iscinumpy.dev/)) for simple ways and robust ways to wrap C++ code.\n",
    "\n",
    "This is a simple way: [pybind11](https://pybind11.readthedocs.io/) is good, but if we want this to compile everywhere (including Windows), we should build makefiles with [CMake](https://cliutils.gitlab.io/modern-cmake/) and [scikit-build](https://scikit-build.readthedocs.io/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "639cf4ea-06ce-44eb-a521-ec08997edfed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile test1_pybind11.cpp\n",
    "\n",
    "#include <pybind11/pybind11.h>\n",
    "namespace py = pybind11;\n",
    "\n",
    "float run(float mu1pt, float mu2pt, float mu1eta, float mu2eta, float mu1phi, float mu2phi) {\n",
    "    return sqrt(\n",
    "        2 * mu1pt * mu2pt * (cosh(mu1eta - mu2eta) - cos(mu1phi - mu2phi))\n",
    "    );\n",
    "}\n",
    "\n",
    "PYBIND11_MODULE(test1_pybind11, m) {\n",
    "    m.def(\"run\", &run);\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e46307a4-006e-45a6-b71b-4fba34ba96c1",
   "metadata": {},
   "source": [
    "The next two cells find compiler arguments (in Python) and then run the compiler (in the shell)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f768e4c7-0b60-478f-bb09-81723e3163ed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from pybind11 import get_include\n",
    "\n",
    "inc = \"-I \" + get_include()\n",
    "plat = \"-undefined dynamic_lookup\" if \"darwin\" in sys.platform else \"-fPIC\"\n",
    "pyinc = !python3-config --cflags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc9cc1a2-af64-4d63-ba83-4f6f013db91b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!c++ -std=c++11 test1_pybind11.cpp -shared {inc} {pyinc.s} -o test1_pybind11.so {plat}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2799a140-2626-4f4e-bf34-f0bc0322c534",
   "metadata": {},
   "source": [
    "And the next cell loads the library and runs it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "730c2016-8fec-48ee-a3a9-8bf66b14b386",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import test1_pybind11\n",
    "\n",
    "test1_pybind11.run(mu1.pt[0, 0], mu2.pt[0, 0], mu1.eta[0, 0], mu2.eta[0, 0], mu1.phi[0, 0], mu2.phi[0, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00341042-bd26-4550-9777-8e2809f3ec89",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = []\n",
    "for i in range(len(mu1)):\n",
    "    for j in range(len(mu1[i])):\n",
    "        results.append(\n",
    "            test1_pybind11.run(mu1.pt[i, j], mu2.pt[i, j], mu1.eta[i, j], mu2.eta[i, j], mu1.phi[i, j], mu2.phi[i, j])\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ce1c619-e752-4328-b8d6-fa1c485425fe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.hist(results, bins=100, range=(0, 100));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb483dc7-2f42-4f86-a5d2-ad7b27fe8268",
   "metadata": {},
   "source": [
    "That took much too long! I'm not even going to measure it.\n",
    "\n",
    "Why? **(Don't scroll down!)**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7af14adc-423a-4053-a8f9-cc5f5d85c707",
   "metadata": {},
   "source": [
    "<br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "996a88e7-5356-4356-bb24-c7f20dd36c59",
   "metadata": {},
   "source": [
    "Instead of calling the C++ function in a Python loop, let's put the loop in C++."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7779a4f8-1efc-46db-826b-8b55d9b8cb31",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile test2_pybind11.cpp\n",
    "\n",
    "#include <pybind11/pybind11.h>\n",
    "#include <pybind11/numpy.h>\n",
    "namespace py = pybind11;\n",
    "\n",
    "void run(\n",
    "    py::array_t<float> results,\n",
    "    py::array_t<float, py::array::c_style | py::array::forcecast> mu1pt,\n",
    "    py::array_t<float, py::array::c_style | py::array::forcecast> mu2pt,\n",
    "    py::array_t<float, py::array::c_style | py::array::forcecast> mu1eta,\n",
    "    py::array_t<float, py::array::c_style | py::array::forcecast> mu2eta,\n",
    "    py::array_t<float, py::array::c_style | py::array::forcecast> mu1phi,\n",
    "    py::array_t<float, py::array::c_style | py::array::forcecast> mu2phi\n",
    ") {\n",
    "    float* results_ = results.mutable_data();\n",
    "    const float* mu1pt_ = mu1pt.data();\n",
    "    const float* mu2pt_ = mu2pt.data();\n",
    "    const float* mu1eta_ = mu1eta.data();\n",
    "    const float* mu2eta_ = mu2eta.data();\n",
    "    const float* mu1phi_ = mu1phi.data();\n",
    "    const float* mu2phi_ = mu2phi.data();\n",
    "\n",
    "    for (int i = 0;  i < results.size();  i++) {\n",
    "        results_[i] = sqrt(\n",
    "            2 * mu1pt_[i] * mu2pt_[i] * (cosh(mu1eta_[i] - mu2eta_[i]) - cos(mu1phi_[i] - mu2phi_[i]))\n",
    "        );\n",
    "    }\n",
    "}\n",
    "\n",
    "PYBIND11_MODULE(test2_pybind11, m) {\n",
    "    m.def(\"run\", &run);\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de4372a8-68c0-46cf-bd1f-37b7c0825071",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!c++ -std=c++11 test2_pybind11.cpp -shared {inc} {pyinc.s} -o test2_pybind11.so {plat}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0abef945-287d-46a0-a7ab-ac64e1121085",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import test2_pybind11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcbd32f6-0466-46f6-8ac4-c600f3f8e84e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mu1pt = ak.to_numpy(ak.ravel(mu1.pt))\n",
    "mu2pt = ak.to_numpy(ak.ravel(mu2.pt))\n",
    "mu1eta = ak.to_numpy(ak.ravel(mu1.eta))\n",
    "mu2eta = ak.to_numpy(ak.ravel(mu2.eta))\n",
    "mu1phi = ak.to_numpy(ak.ravel(mu1.phi))\n",
    "mu2phi = ak.to_numpy(ak.ravel(mu2.phi))\n",
    "assert len(mu1pt) == len(mu2pt) == len(mu1eta) == len(mu2eta) == len(mu1phi) == len(mu2phi)\n",
    "\n",
    "results = np.empty(len(mu1pt), dtype=np.float32)\n",
    "test2_pybind11.run(results, mu1pt, mu2pt, mu1eta, mu2eta, mu1phi, mu2phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31ca7533-73b5-4c32-ab67-9d744e8695a8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.hist(results, bins=100, range=(0, 100));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08fe89c8-61d1-43f3-8a09-a45629612ec3",
   "metadata": {},
   "source": [
    "Now _that_ was fast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d46bc721-d746-49f2-945f-b8eccd7e98ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "(mu1 + mu2).mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f2c049a-ce69-48b1-8d5e-50233e690246",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "results = np.empty(len(mu1pt), dtype=np.float32)\n",
    "test2_pybind11.run(results, mu1pt, mu2pt, mu1eta, mu2eta, mu1phi, mu2phi)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aae71b4-dabb-4f66-b7d3-4cb5d92546dd",
   "metadata": {},
   "source": [
    "It was fast, but required a lot of work.\n",
    "\n",
    "It's the best way to wrap C++ code as a library, but if you want to interleave snippets of C++ in a mostly Python script/notebook, you can use\n",
    "\n",
    "- [ROOT](https://root.cern/)'s [gInterpreter.Declare](https://root.cern.ch/doc/master/classTInterpreter.html#a84a4890123faea8e740cd9f8d690e1e3)\n",
    "- [ROOT](https://root.cern/)'s [RDataFrame](https://root.cern/doc/master/classROOT_1_1RDataFrame.html)\n",
    "- [cppyy](https://cppyy.readthedocs.io/) (see [demo](https://github.com/henryiii/python-compiled-minicourse/blob/master/06-cppyy/06-cppyy.ipynb))\n",
    "\n",
    "Any method that integrates Python and C++ has to translate data structures between the two languages. C++ has no Awkward Array, but NumPy arrays can be interpreted as `std::array` of numeric types.\n",
    "\n",
    "The problem of integrating Python and C++ often comes down to deconstructing Python objects into NumPy arrays, passing them over, and building NumPy output back into Python objects."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc7f1e5-703a-4673-a56e-abff145ddea7",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f58679db-79c1-4df7-8720-9a27e325dad8",
   "metadata": {},
   "source": [
    "## The example problem in Numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d891aa4-54c3-465b-a421-c314ec2319fa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numba as nb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd525029-4721-4efd-9f4b-41c97ee698a2",
   "metadata": {},
   "source": [
    "Also for interleaving compiled code in a mostly Python script/notebook, Numba does two things differently:\n",
    "\n",
    "- it uses Python syntax instead of C++ syntax\n",
    "- it translates objects for you/presents an illusion that you're using Python objects throughout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64034786-c764-4805-862c-e89bd87586dc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def this_is_slow():\n",
    "    out = 0.0\n",
    "    for i in range(1000000):\n",
    "        out += i\n",
    "        out *= 0.5\n",
    "    return out\n",
    "\n",
    "this_is_slow()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d990c1d3-c5d6-4ab4-a22a-6e3cce4066e6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "this_is_slow()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d5b0413-660b-4c12-9357-256908062965",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def this_is_fast():\n",
    "    out = 0.0\n",
    "    for i in range(1000000):\n",
    "        out += i\n",
    "        out *= 0.5\n",
    "    return out\n",
    "\n",
    "this_is_fast()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4995d87-c069-4fb4-8a7d-ac1ee6198273",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "this_is_fast()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2da7bd2-4485-460e-a510-625a32ccc909",
   "metadata": {},
   "source": [
    "The `@nb.njit` (no-Python Just-In-Time compilation) looks like magic, but there's a catch:\n",
    "\n",
    "- not all Python code can be accelerated\n",
    "- when it can't be accelerated, the error message is often formidable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6977f71e-6cf7-4057-93db-85f53c95e6f4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def this_is_not_working():\n",
    "    out = \"bad\"\n",
    "    for i in range(1000000):\n",
    "        if i == 0:\n",
    "            out = 0.0\n",
    "        out += i\n",
    "        out *= 0.5\n",
    "    return out\n",
    "\n",
    "this_is_not_working()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "208e22e0-d0fc-4cf2-9564-bcb6caa8b04f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class SomethingIMadeUp:\n",
    "    pass\n",
    "\n",
    "something_I_made_up = SomethingIMadeUp()\n",
    "\n",
    "@nb.njit\n",
    "def this_is_not_working_either(x):\n",
    "    return x\n",
    "\n",
    "this_is_not_working_either(something_I_made_up)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a369214-4271-4609-8f01-6c5ef3a4b643",
   "metadata": {},
   "source": [
    "Numba only knows about [these Python features](https://numba.readthedocs.io/en/stable/reference/pysupported.html) and [these NumPy functions](https://numba.readthedocs.io/en/stable/reference/numpysupported.html).\n",
    "\n",
    "And any that someone has written an explicit extension for."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3677478-610a-499e-b754-a66e58812013",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0960fe5-4a36-4d38-a950-6c4fedd1e20c",
   "metadata": {},
   "source": [
    "Awkward Array has been extended to Numba.\n",
    "\n",
    "None of the `ak.*` functions or fancy slices are supported, but basic iteration is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ebdae28-819b-464e-8072-2f6b04c5a68b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def mass_in_numba(results, mu1_everything, mu2_everything):\n",
    "    k = 0\n",
    "    for mu1_event, mu2_event in zip(mu1_everything, mu2_everything):\n",
    "        for m1, m2 in zip(mu1_event, mu2_event):\n",
    "            results[k] = np.sqrt(\n",
    "                2 * m1.pt * m2.pt * (np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi))\n",
    "            )\n",
    "            k += 1\n",
    "\n",
    "results = np.zeros(len(mu1pt), dtype=np.float32)\n",
    "mass_in_numba(results, mu1, mu2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86b081e0-208f-44c1-9530-e815ca4aa7d7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.hist(results, bins=100, range=(0, 100));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fb6f204-ed7c-44f0-b799-86b3d3bf9f83",
   "metadata": {},
   "source": [
    "Compare:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a90436d1-311f-4631-9019-e1157c00eeca",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "(mu1 + mu2).mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb83952-882f-4a85-b18a-46ccfdbf9d75",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "results = np.empty(len(mu1pt), dtype=np.float32)\n",
    "test2_pybind11.run(results, mu1pt, mu2pt, mu1eta, mu2eta, mu1phi, mu2phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a79e68da-1038-4239-86cc-b950c0ba53dc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "results = np.empty(len(mu1pt), dtype=np.float32)\n",
    "mass_in_numba(results, mu1, mu2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf1f769f-a696-4bd7-95bb-d0583ac67677",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "604e4149-bf27-4b05-90fb-09ccf18d109a",
   "metadata": {},
   "source": [
    "### Numba vectorize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f62d6eab-4e0d-4499-9488-6d2d2e43fe52",
   "metadata": {},
   "source": [
    "Numba has several ways to compile code. (https://numba.readthedocs.io/en/stable/user/vectorize.html) makes a function in which `output[i] = f(input[i])` for a given `f`.\n",
    "\n",
    "(The signature in `@nb.vectorize(...)` will not be required for long: [numba/numba#8995](https://github.com/numba/numba/pull/8995).)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c871aa30-ea8a-4b40-8fe1-f9a16b60652b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@nb.vectorize([nb.float32(nb.float32, nb.float32, nb.float32, nb.float32, nb.float32, nb.float32)])\n",
    "def mass_in_numba_vectorize(mu1pt, mu2pt, mu1eta, mu2eta, mu1phi, mu2phi):\n",
    "    return np.sqrt(\n",
    "        2 * mu1pt * mu2pt * (np.cosh(mu1eta - mu2eta) - np.cos(mu1phi - mu2phi))\n",
    "    )\n",
    "\n",
    "mass_in_numba_vectorize(mu1.pt, mu2.pt, mu1.eta, mu2.eta, mu1.phi, mu2.phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "214849a3-0e41-4bf8-a0e9-17174c746798",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "results = mass_in_numba_vectorize(mu1.pt, mu2.pt, mu1.eta, mu2.eta, mu1.phi, mu2.phi)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f724c2b-2070-4285-87ea-30d5f7db01eb",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51261036-2d91-4f9a-a72c-339b38fae7bd",
   "metadata": {},
   "source": [
    "### Creating an Awkward Array in a Numba-compiled function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bef8f10-a007-443d-baec-c69077b8e28f",
   "metadata": {},
   "source": [
    "[ArrayBuilder](https://awkward-array.org/doc/main/reference/generated/ak.ArrayBuilder.html) makes Awkward Arrays, discovering the type from the order in which methods are called. It can also be used in Numba-compiled functions.\n",
    "\n",
    "- `ArrayBuilder` must be constructed outside the Numba-compiled function.\n",
    "- Its `snapshot` method must be called outside the Numba-compiled function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7677d09b-a902-49c5-85a8-9168c1ae1251",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def mass_in_numba_awkward_output(builder, mu1_everything, mu2_everything):\n",
    "    for mu1_event, mu2_event in zip(mu1_everything, mu2_everything):\n",
    "        builder.begin_list()\n",
    "        for m1, m2 in zip(mu1_event, mu2_event):\n",
    "            builder.append(\n",
    "                np.sqrt(\n",
    "                    2 * m1.pt * m2.pt * (np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi))\n",
    "                )\n",
    "            )\n",
    "        builder.end_list()\n",
    "\n",
    "builder = ak.ArrayBuilder()\n",
    "mass_in_numba_awkward_output(builder, mu1, mu2)\n",
    "builder.snapshot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3016d99-673e-4250-864a-bf14c0f963d3",
   "metadata": {},
   "source": [
    "Dynamically discovering types is a bit slow (even though it's implemented in C++)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31cad68d-ad87-477d-a3fd-53228dc9f82e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "builder = ak.ArrayBuilder()\n",
    "mass_in_numba_awkward_output(builder, mu1, mu2)\n",
    "builder.snapshot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14d70d84-c2e2-41d8-97f9-9a286cf96f04",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9d3e71d-4009-4e11-a778-4d4cafc82c82",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Using the new LayoutBuilder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "570bceb8-01dd-490b-b8af-427d46c9a7eb",
   "metadata": {},
   "source": [
    "`LayoutBuilder` is a brand-new (but lower-level) interface that lets you specify the data type upfront, gaining back some speed.\n",
    "\n",
    "(There's also a [LayoutBuilder in C++](https://github.com/scikit-hep/awkward/blob/main/header-only/layout-builder/awkward/LayoutBuilder.h) for building Awkward Arrays from C++.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea70d125-22b0-4c08-b4a5-9197cc9ec3cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import awkward.numba.layoutbuilder as lb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12c2c1a7-a23d-4861-aac4-f32f3c0148e6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def mass_in_numba_awkward_output2(builder, mu1_everything, mu2_everything):\n",
    "    for mu1_event, mu2_event in zip(mu1_everything, mu2_everything):\n",
    "        subbuilder = builder.begin_list()\n",
    "        for m1, m2 in zip(mu1_event, mu2_event):\n",
    "            subbuilder.append(\n",
    "                np.sqrt(\n",
    "                    2 * m1.pt * m2.pt * (np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi))\n",
    "                )\n",
    "            )\n",
    "        builder.end_list()\n",
    "\n",
    "builder = lb.ListOffset(np.int64, lb.Numpy(np.float32))\n",
    "mass_in_numba_awkward_output2(builder, mu1, mu2)\n",
    "ak.Array(builder.snapshot())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89fb2bc9-a4f0-42a4-9cdd-e7c53704240b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "builder = lb.ListOffset(np.int64, lb.Numpy(np.float32))\n",
    "mass_in_numba_awkward_output2(builder, mu1, mu2)\n",
    "ak.Array(builder.snapshot())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fab4ffff-90ad-4678-a7ae-94f1274c3d89",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b8f6f6-508e-4eef-a1cc-959a4f7198a4",
   "metadata": {},
   "source": [
    "### Using Vector in Numba-compiled functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93323c70-a837-4088-a71e-b67a2a90d395",
   "metadata": {},
   "source": [
    "The [Vector](https://vector.readthedocs.io/) library is also a Numba extension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ac4d6dd-3324-4abe-a415-231f59caae56",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def mass_in_numba_with_vector(results, mu1_everything, mu2_everything):\n",
    "    k = 0\n",
    "    for mu1_event, mu2_event in zip(mu1_everything, mu2_everything):\n",
    "        for m1, m2 in zip(mu1_event, mu2_event):\n",
    "            results[k] = (m1 + m2).mass\n",
    "            k += 1\n",
    "\n",
    "results = np.zeros(len(mu1pt), dtype=np.float32)\n",
    "mass_in_numba_with_vector(results, mu1, mu2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49106f60-b75d-4050-9c5c-6dedc2d0f486",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.hist(results, bins=100, range=(0, 100));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "228d03fd-f7c2-455b-8892-a99db7bd8ba7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "results = np.empty(len(mu1pt), dtype=np.float32)\n",
    "mass_in_numba_with_vector(results, mu1, mu2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a486a2-3517-4024-8b19-93b908ef7708",
   "metadata": {},
   "source": [
    "But now we're back to converting the vectors to Cartesian coordinates and calculating $m = \\sqrt{E^2 - {p_x}^2 - {p_y}^2 - {p_z}^2}$ (though it's compiled and optimized by LLVM).\n",
    "\n",
    "However, compare to Vector without Numba:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1826265a-dd0c-4e50-a32f-74cfc004fa2c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "(mu1 + mu2).mass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fde540b1-9853-4ef2-916f-b4449b9d6c84",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38e2c35d-9d54-4e42-8525-ebc36dc477c6",
   "metadata": {},
   "source": [
    "### Filling histograms (hist) in Numba-compiled functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4b066a0-bfcd-4bcc-82ff-4dc339a72b9e",
   "metadata": {},
   "source": [
    "Coming ~~soon~~ someday..."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}