07_callables.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LowLevelCallables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simple question: Let's say you have an array (nicely laid out in memory), but you *have* to loop over it. Which is faster, `np.log` or `math.log`?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ctypes\n",
    "import math\n",
    "\n",
    "import numba\n",
    "import numpy as np\n",
    "import scipy.integrate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.random.rand(10_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "for item in arr:\n",
    "    np.log(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "for item in arr:\n",
    "    math.log(item)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course, if we use array processing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "np.log(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ## Aside\n",
    ">\n",
    "> NumPy arrays were designed to work at larger scales. Under about 10 elements, you may even get higher performance from Python lists! Normally, the nice syntax is still worth it, and you can usually find a way to scale out to array processing by adding another dimension, but something to keep in mind."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's cast the above problem in a slightly different use case, one you are much more likely to run into: functions that take a function. Let's say we have a processing function that takes an array and a function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def array_map(array, func):\n",
    "    out_array = np.empty_like(array)\n",
    "\n",
    "    # Note the comma\n",
    "    (size,) = array.shape\n",
    "\n",
    "    for i in range(size):\n",
    "        out_array[i] = func(array[i])\n",
    "\n",
    "    return out_array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quick check to see if the above effect is still true:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "array_map(arr, np.log)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "array_map(arr, math.log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's imagine that array_map actually contained a compiled loop. And our function was a compiled function. What will happen *inside the loop*, though, is:\n",
    "\n",
    "```\n",
    "Compiled -> Python -> Compiled\n",
    "```\n",
    "\n",
    "\n",
    "Which kills our performance. What we'd like to do is skip the Python middle man in this case.\n",
    "\n",
    "Hey, we have Numba, we don't have to imagine:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.njit\n",
    "def numba_array_map(array, func):\n",
    "    out_array = np.empty_like(array)\n",
    "\n",
    "    # Note the comma\n",
    "    (size,) = array.shape\n",
    "\n",
    "    for i in range(size):\n",
    "        out_array[i] = func(array[i])\n",
    "\n",
    "    return out_array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba helpfully will not allow us to pass in a Python function. But if we pass in a numba function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.njit\n",
    "def numba_log(v):\n",
    "    return math.log(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "numba_array_map(arr, numba_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Success! We can pass in a jit function into a jit function and they call each other *without* going back to Python! But wouldn't it be nice to be able to do this without a jit master function?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Callbacks and Scipy\n",
    "\n",
    "Let's look at Scipy. It is a large library with lots of routines that can take functions and iterate with them. They have implemented a standard way to interact with compiled function pointers through what they call a LowLevelCallable interface; if you have a function pointer, you can completely skip the Python middle man!\n",
    "\n",
    "> Note that the LowLevelCallable is just a standard interface they proposed inside Scipy to handle callables from three different sources (PyCapsule, ctypes, and cffi), and to bundle in the idea of user data (absent in ctypes and cffi). The idea can be used in other places, usually with just the ctypes interface."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try the following integral:\n",
    "\n",
    "$$\n",
    "\\int _{-\\infty} ^ {\\infty} e^{-a x ^2} dx = \\sqrt{\\frac{\\pi}{a}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @numba.vectorize([numba.double(numba.double, numba.double)])\n",
    "def integrand(x, a):\n",
    "    return np.exp(-a * x**2)\n",
    "\n",
    "\n",
    "@np.vectorize\n",
    "def gauss_py(a):\n",
    "    y, abserr = scipy.integrate.quad(integrand, -np.inf, np.inf, (a,))\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> #### Note:\n",
    ">\n",
    "> Since you may not have seen it before, `np.vectorize` is a Python version of `numba.vectorize`; you don't get a performance benefit from it, but it simplifies calling this on an array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "a = np.linspace(0.1, 10, 10_000)\n",
    "\n",
    "py_result = gauss_py(a)\n",
    "\n",
    "print(py_result)\n",
    "print(np.sqrt(np.pi / a))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Results are not bad, but the performance is not great, for just 10K points. Even if we add numba, not much changes. This is because we are calling integrand through Python in a loop inside the quad routine.\n",
    "\n",
    "Let's check the LowLevelCallable signature:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# scipy.integrate.quad?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the key part:\n",
    "\n",
    "\n",
    "```\n",
    "func : {function, scipy.LowLevelCallable}\n",
    "    A Python function or method to integrate.  If `func` takes many\n",
    "    arguments, it is integrated along the axis corresponding to the\n",
    "    first argument.\n",
    "\n",
    "    If the user desires improved integration performance, then `f` may\n",
    "    be a `scipy.LowLevelCallable` with one of the signatures::\n",
    "\n",
    "        double func(double x)\n",
    "        double func(double x, void *user_data)\n",
    "        double func(int n, double *xx)\n",
    "        double func(int n, double *xx, void *user_data)\n",
    "\n",
    "    The ``user_data`` is the data contained in the `scipy.LowLevelCallable`.\n",
    "    In the call forms with ``xx``,  ``n`` is the length of the ``xx``\n",
    "    array which contains ``xx[0] == x`` and the rest of the items are\n",
    "    numbers contained in the ``args`` argument of quad.\n",
    "\n",
    "    In addition, certain ctypes call signatures are supported for\n",
    "    backward compatibility, but those should not be used in new code.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, double(double) sounds easy - but we need to pass in one more bit of information, the value of `a`. Let's try making that first using args (nicer), and then using user data (ugly):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the signature we expect:\n",
    "\n",
    "```c\n",
    "double func(int n, double *xx)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.cfunc(numba.double(numba.int32, numba.types.CPointer(numba.double)))\n",
    "def integrand(n, x_ptr):\n",
    "    x, a = numba.carray(x_ptr, (n,), np.double)  # Fails if n != 2, but that's good\n",
    "    return np.exp(-a * x**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the numba function provides a ctypes interface through the `.ctypes` property, so we can use that in LowLevelCallable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = scipy.LowLevelCallable(integrand.ctypes)\n",
    "\n",
    "\n",
    "@np.vectorize\n",
    "def gauss_py(a):\n",
    "    y, abserr = scipy.integrate.quad(c, -np.inf, np.inf, (a,))\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "a = np.linspace(0.1, 10, 10_000)\n",
    "\n",
    "py_result = gauss_py(a)\n",
    "\n",
    "print(py_result)\n",
    "print(np.sqrt(np.pi / a))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Much better! We've now avoided calling Python at all once we enter the integrate loop. This should perform close to a full Fortran or C implementation, and it just took adding 2-3 lines of code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### User Data\n",
    "\n",
    "> Included as an example. Don't do it this way. Just don't."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the signature we expect:\n",
    "\n",
    "```c\n",
    "double func(double x, void *user_data)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.cfunc(numba.double(numba.double, numba.types.voidptr))\n",
    "def integrand(x, user_ptr):\n",
    "    (a,) = numba.carray(user_ptr, (1,), np.double)\n",
    "    return np.exp(-a * x**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_array = np.array([0.0])\n",
    "c = scipy.LowLevelCallable(integrand.ctypes, a_array.ctypes.data_as(ctypes.c_void_p))\n",
    "\n",
    "\n",
    "@np.vectorize\n",
    "def gauss_py(a):\n",
    "    a_array[0] = a\n",
    "    y, abserr = scipy.integrate.quad(c, -np.inf, np.inf)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "a = np.linspace(0.1, 10, 10_000)\n",
    "\n",
    "py_result = gauss_py(a)\n",
    "\n",
    "print(py_result)\n",
    "print(np.sqrt(np.pi / a))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note to self: you can get the address of the callable directly from the ctypes object with:\n",
    "\n",
    "```python\n",
    "ctypes.cast(integrand.ctypes, ctypes.c_void_p)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Further reading\n",
    "\n",
    "* [Cython example](https://tjol.eu/blog/lowlevelcallable-magic.html) (this lesson was based generally on this)\n",
    "* [Numba example](https://ilovesymposia.com/2017/03/12/scipys-new-lowlevelcallable-is-a-game-changer/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:performance-minicourse] *",
   "language": "python",
   "name": "conda-env-performance-minicourse-py"
  },
  "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.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "2d4e9b9c84dab3e1662173f95b81bd7f8a551068d04f5f3c42d164db7312a928"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}