00_numpy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NumPy: Numeric computing\n",
    "\n",
    "The core of all modern scientific computing. You should know this one! This is NumPy + SciPy + Matplotlib, which form the core of the modern computing stack for [Scientific Python](https://scientific-python.org)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is used for N-dimensional arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = np.array([1, 2, 3])\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The clever trick is that computations on an array are pre-compiled, and can be much faster than Python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that it was designed for large arrays; around 10 elements or less you _might_ be faster using plain Python (though it's still more expressive). You should \"vectorize\" your code (by making the arrays bigger with more dimensions) if your arrays are very small and you care about performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the spirit of the course, I'll avoid covering the basics of NumPy here (arrays, simple dtypes, operations), and instead focus on intermediate features."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Slicing\n",
    "\n",
    "NumPy slicing may not make copies of arrays, but can just adjust the start and strides of the data. This is usually true of reshaping and adding empty (length 1) dimensions, as well. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.zeros([2, 3, 4], dtype=int)\n",
    "print(arr.flags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr[:1, :1, :1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Slicing produces a copy; we can see the effect of setting values, for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.zeros([2, 3, 4], dtype=int)\n",
    "x = arr[:, :2, 2:]\n",
    "x[...] = 1\n",
    "print(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.strides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.flags"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Flattening also avoids a copy if the data is contagious:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.zeros([2, 3, 4], dtype=int)\n",
    "y = arr.ravel()\n",
    "y[:3] = 2\n",
    "print(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, flattening / reshaping an array without a copy may not always be possible, such as if the data is non-contagious:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = arr.T.ravel()\n",
    "y[:3] = 3\n",
    "print(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Selections\n",
    "\n",
    "You can also use a very similar syntax to select items from an array. Selections always *return* copies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.zeros([12], dtype=int)\n",
    "v = arr[[1, 2, 3]]\n",
    "v[...] = [1, 2, 3]\n",
    "print(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the right of an equal sign, they are a totally different overload, so that is supported (nothing is \"returned\", however, so there's no issue here)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.zeros([12], dtype=int)\n",
    "arr[[1, 2, 3]] = [1, 2, 3]\n",
    "print(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also use a boolean mask:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.arange(4, dtype=int)\n",
    "arr[[True, False, True, False]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.zeros([4], dtype=int)\n",
    "arr[[True, False, True, False]] = [1, 2]\n",
    "print(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An easy way to make a boolean array of the same shape? Use a comparison on the original array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.arange(12, dtype=int)\n",
    "arr[arr > 5] = -1\n",
    "print(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Logical operators on arrays are actually overloaded bitwise operators (`&`, `|`, `~`), not actual logical operators (`and`, `or`, `not`)! This causes several problems for arrays:\n",
    "\n",
    "* `a < b < c` expands to `a < b and b < c`, which doesn't work on arrays\n",
    "* `a < b & b < c` is actually `a < (b & b) < c`, which is not what you wanted\n",
    "* `(a < b) & (b < c)` parentheses are required to make it work on arrays!\n",
    "\n",
    "There are some ideas to fix this, but nothing has made it into Python yet. So for now, parentheses it up!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.arange(12, dtype=int)\n",
    "print(f\"{arr[(arr>5) & (arr<9)] = }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Random Numbers\n",
    "\n",
    "The simple random number interface `np.random.<distribution>` should only be used for highly interactive work. Normal usage should use a more modern, object oriented approach; first construct a random number generator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can pass an explicit seed if you want reproducibility. Then this generator has distributions as methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dist = rng.normal(0, 1, size=1_000_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fix, ax = plt.subplots()\n",
    "ax.hist(dist, bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "The benefits include explicit control over the generator, reproducibility, and support for multiple independent generators."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Broadcasting\n",
    "\n",
    "One benefit of true 1D arrays (vs. languages that don't have 1D arrays) is that NumPy can support broadcasting. Broadcasting occurs whenever an operation (element-wise or matrix multiplication) encounters an array that is the wrong shape. The rules can be viewed two ways:\n",
    "\n",
    "1. If the number of dimensions does not match, prepend length-1 dimensions until they do.\n",
    "2. If the size of a dimension does not match and one of the dimensions is 1, duplicate along that axis to make it match.\n",
    "\n",
    "Or:\n",
    "1. Moving right to left, if a dimension is missing or 1, set the stride to 0 on that array, so the iteration pointer does not change along that dimension.\n",
    "\n",
    "The former is often easier to reason about, while the later is how it's implemented (so it is very efficient)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(np.ones((2, 3, 4)) * np.ones((1, 4))).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can predict the final shape with `broadcast_shapes`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.broadcast_shapes((2, 3, 4), (1, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also explicitly expand with `np.broadcast_arrays`.\n",
    "\n",
    "Broadcasting is extremely useful for array-at-a-time programming, such as for coordinate arrays and the like. If you organize your calculations to delay expansion, you can optimize a lot of computation out."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, we can use `ogrid` or `meshgrid(..., sparse=True)` to replace `mgrid` but with compressed arrays ready for broadcasting. Here we make a grid with 101 points on x from 1-2 and 501 points on y from 0 to 5 (ends inclusive):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = np.mgrid[2.5:3.5:101j, 0:5:501j]\n",
    "radius = np.sqrt(x**2 + y**2)\n",
    "print(f\"{x.shape=}, {y.shape=}\")\n",
    "print(f\"{radius[50,400] = }\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = np.ogrid[2.5:3.5:101j, 0:5:501j]\n",
    "radius = np.sqrt(x**2 + y**2)\n",
    "print(f\"{x.shape=}, {y.shape=}\")\n",
    "print(f\"{radius[50,400] = }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Masked arrays\n",
    "\n",
    "A powerful and under supported feature is masked arrays. If values are missing, you can \"mask\" them.\n",
    "\n",
    "This is quite nice logically, but the downside is not all interfaces support masked arrays. You also use extra space for the boolean mask. Another trick is to use NaN's in a floating point array, or to use Pandas's support for None's in all arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = np.ogrid[-5:5:101j, -5:5:101j]\n",
    "radius = np.sqrt(x**2 + y**2)\n",
    "mradius = np.ma.masked_where(radius > 5, radius)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.pcolormesh(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.pcolormesh(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fix, ax = plt.subplots()\n",
    "ax.pcolormesh(x.T, y.T, mradius.T)\n",
    "ax.set_aspect(\"equal\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Record arrays (AKA preview for Pandas!)\n",
    "\n",
    "NumPy has support for complex DTypes. While Pandas or xarray handles structured dtypes more elegantly (and in a columnar form), but if you actually have data from some source that is structured and in an array form, this can be incredibly useful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.array([(1, 2.0), (3, 4.0)], dtype=[(\"id\", int), (\"val\", float)])\n",
    "arr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can select out a single named dtype (without copy) as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr[\"val\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Smaller features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Buffer protocol\n",
    "\n",
    "There's a buffer protocol in Python 3 that allows different libraries to communicate this sort of data without copies. This is also useful to communicate things like shared memory or existing memory to NumPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "python_array = array.array(\"d\", (1, 2, 3, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can explicitly convert from a buffer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.frombuffer(python_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or most NumPy functions also work directly on buffers by converting them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.asarray(python_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This \"conversion\" does not copy! You still are looking at the buffer's original memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr[1] = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "python_array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can verify that the NumPy array doesn't own the memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr.flags.owndata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Custom arrays\n",
    "\n",
    "### UFuncs\n",
    "\n",
    "NumPy has the concept of UFuncs; functions that can take array arguments (broadcastable) and will return a broadcast result. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.sin(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.sin(np.array([0, 1, 2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.add(7, np.arange(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "UFuncs have several standard keyword arguments:\n",
    "\n",
    "* `out`: Output to existing array, skip allocation\n",
    "* `where`: Mask computation\n",
    "* `axes`, `axis`: axes to operate on (generalized UFuncs only)\n",
    "* `keepdims`: To remove reduced dims (some generalized UFuncts only)\n",
    "* `casting='same_kind'`: rules for casting different DTypes\n",
    "* `order='K'`: memory layout (\"K\"eep)\n",
    "* `dtype=None`: Dtype for output array\n",
    "* `subok=True`: Output can be a subclass\n",
    "* `signature`/`extobj`: Exact control over dtypes/buffers\n",
    "\n",
    "(A generalized UFunc supports specific patterns, like matrix multiplication, rather than being element-wise)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "UFuncs also support several methods and have a few properties. There are exactly six methods:\n",
    "\n",
    "* `__call__`: Elementwise computation\n",
    "* `at`: Local elementwise computation (provide indices)\n",
    "* `reduce`: A reduction\n",
    "* `reduceat`: A local reduction (provide indices)\n",
    "* `accumulate`: An accumulation\n",
    "* `outer`: An outer apply (`np.multiply.outer` is identical to `np.tensordot`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NEP 13/18\n",
    "\n",
    "If you have a custom class, you can also customize UFuncs via NEP 13!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Aside: let's add pytest's raises decorator\n",
    "import pytest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleUnc(np.ndarray):\n",
    "    pass\n",
    "\n",
    "\n",
    "def simple_unc(val, unc):\n",
    "    arr = np.array(list(zip(val, unc)), dtype=[(\"val\", float), (\"unc\", float)])\n",
    "    return arr.view(SimpleUnc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = simple_unc([1, 2, 3], [0.1, 0.1, 0.1])\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a == a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that this does raise an error if you try to add it, though:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pytest.raises(TypeError):\n",
    "    a + a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's provide `__array_ufunc__`, which will allow us to customize the behavior of UFuncs. This will use Python 3.10's pattern matching for simplicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleUnc(np.ndarray):\n",
    "    def __array_ufunc__(self, ufunc, method: str, *inputs, **kwargs):\n",
    "        # Avoid infinite recursion\n",
    "        raw_inputs = [np.asarray(x) for x in inputs]\n",
    "\n",
    "        match (ufunc, method, raw_inputs):\n",
    "            # Custom add / subtract\n",
    "            case np.add | np.subtract, \"__call__\", (a, b):\n",
    "                # This can waste an allocation\n",
    "                (result,) = kwargs.pop(\"out\", [np.empty(self.shape, self.dtype)])\n",
    "\n",
    "                ufunc(a[\"val\"], b[\"val\"], out=result[\"val\"], **kwargs)\n",
    "                np.add(a[\"unc\"] ** 2, b[\"unc\"] ** 2, out=result[\"unc\"], **kwargs)\n",
    "                np.sqrt(result[\"unc\"], out=result[\"unc\"], **kwargs)\n",
    "\n",
    "                return result.view(self.__class__)\n",
    "\n",
    "            # Fall back on whatever it would have done before - do not return this subclass\n",
    "            case _:\n",
    "                return super().__array_ufunc__(ufunc, method, *raw_inputs, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = simple_unc([1, 2, 3], [0.1, 0.1, 0.1])\n",
    "a + a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a - a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is very simple and doesn't account for correlations, but it does show that it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a == a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's a lot more to NumPy; it's worth looking around in the docs. An example of a function you might find is `np.diff(a)`, which is equivalent to `a[:-1] - a[1:]` but reads better without repeating the array name. Also look at the huge number of useful methods in SciPy."
   ]
  }
 ],
 "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.11.4"
  },
  "vscode": {
   "interpreter": {
    "hash": "2d4e9b9c84dab3e1662173f95b81bd7f8a551068d04f5f3c42d164db7312a928"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}