{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cython\n",
    "import timeit\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Native code compilation\n",
    "\n",
    "We will see how to convert Python code to native compiled code. We will use the example of calculating the pairwise distance between a set of vectors, a $O(n^2)$ operation. \n",
    "\n",
    "For native code compilation, it is usually preferable to use explicit for loops and minimize the use of `numpy` vectorization and broadcasting because\n",
    "\n",
    "- It makes it easier for the `numba` JIT to optimize\n",
    "- It is easier to \"cythonize\"\n",
    "- It is easier to port to C++\n",
    "\n",
    "However, use of vectors and matrices is fine especially if you will be porting to use a C++ library such as Eigen."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Timing code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def f(n=1):\n",
    "    start = time.time()\n",
    "    time.sleep(n)\n",
    "    elapsed = time.time() - start\n",
    "    return elapsed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clock time\n",
    "\n",
    "The `time` magic function calls the Unix `time` command. This returns 3 different times:\n",
    "\n",
    "- user time is time spent by user code and libraries\n",
    "- sys time is time spent on operating system calls (kernel calls)\n",
    "- wall time is time that has elapsed from your perspective\n",
    "\n",
    "For concurrent programs, user/sys time can be greater than wall time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "time.sleep(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `timeit`\n",
    "\n",
    "The `-r` argument says how many runs to average over, and `-n` says how many times to run the function in a loop per run. See `%timeit?` for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit time.sleep(0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r3 time.sleep(0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -n10 time.sleep(0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r3 -n10 time.sleep(0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `-o` flag returns an object of the time statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t1 = %timeit -n10 -o time.sleep(0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "', '.join([method for method in dir(t1) if not method.startswith('_')])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also use `timeit` as a Python module.\n",
    "\n",
    "Pass it a callable with no arguments or a string. This is not as convenient as the magic function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import timeit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timeit.timeit('time.sleep(0.01)', number=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timeit.timeit(lambda: time.sleep(0.01), number=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time unit conversions\n",
    "\n",
    "```\n",
    "1 s = 1,000 ms\n",
    "1 ms = 1,000 µs\n",
    "1 µs = 1,000 ns\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Profiling\n",
    "\n",
    "If you want to identify bottlenecks in a Python script, do the following:\n",
    "    \n",
    "- First make sure that the script is modular - i.e. it consists mainly of function calls\n",
    "- Each function should be fairly small and only do one thing\n",
    "- Then run a profiler to identify the bottleneck function(s) and optimize them\n",
    "\n",
    "See the Python docs on [profiling Python code](https://docs.python.org/3/library/profile.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Profiling can be done in a notebook with %prun, with the following readouts as column headers:\n",
    "\n",
    "- ncalls\n",
    "    - for the number of calls,\n",
    "- tottime\n",
    "    - for the total time spent in the given function (and excluding time made in calls to sub-functions),\n",
    "- percall\n",
    "    - is the quotient of tottime divided by ncalls\n",
    "- cumtime\n",
    "    - is the total time spent in this and all subfunctions (from invocation till exit). This figure is accurate even for recursive functions.\n",
    "- percall\n",
    "    - is the quotient of cumtime divided by primitive calls\n",
    "- filename:lineno(function)\n",
    "    - provides the respective data of each function \n",
    "    \n",
    "See `%prun?` for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def foo1(n):\n",
    "    return np.sum(np.square(np.arange(n)))\n",
    "\n",
    "def foo2(n):\n",
    "    return sum(i*i for i in range(n))\n",
    "\n",
    "def foo3(n):\n",
    "    [foo1(n) for i in range(10)]\n",
    "    foo2(n)\n",
    "\n",
    "def foo4(n):\n",
    "    return [foo2(n) for i in range(100)]\n",
    "    \n",
    "def work(n):\n",
    "    foo1(n)\n",
    "    foo2(n)\n",
    "    foo3(n)\n",
    "    foo4(n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "work(int(1e5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `-D` saves results in a form that the `pstats` moudle can parse.\n",
    "- `-q` suppresses output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%prun -q -D work.prof work(int(1e5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pstats\n",
    "p = pstats.Stats('work.prof')\n",
    "p.print_stats()\n",
    "pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p.sort_stats('time', 'cumulative').print_stats('foo')\n",
    "pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p.sort_stats('ncalls').print_stats(5)\n",
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can get the results object directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profile = %prun -r -q work(int(1e5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profile.sort_stats('cumtime').print_stats(5)\n",
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We may not need `pstats` for simple analysis. This limits to calls with `foo` in the name, and sorts by `cumtime`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%prun -l foo -s cumtime work(int(1e5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizing a function\n",
    "\n",
    "Our example will be to optimize a function that calculates the pairwise distance between a set of vectors.\n",
    "\n",
    "We first use a built-in function from`scipy` to check that our answers are right and also to benchmark how our code compares in speed to an optimized compiled routine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.spatial.distance import squareform, pdist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 100\n",
    "p = 100\n",
    "xs = np.random.random((n, p))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We save the result to compare with our own implementations later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sol = squareform(pdist(xs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r3 -n3 squareform(pdist(xs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simple version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdist_py(xs):\n",
    "    \"\"\"Unvectorized Python.\"\"\"\n",
    "    n, p = xs.shape\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(n):\n",
    "            for k in range(p):\n",
    "                A[i,j] += (xs[i, k] - xs[j, k])**2\n",
    "            A[i,j] = np.sqrt(A[i,j])\n",
    "    return A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we \n",
    "\n",
    "- first check that the output is **right**\n",
    "- then check how fast the code is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_py\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exploiting symmetry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdist_sym(xs):\n",
    "    \"\"\"Unvectorized Python.\"\"\"\n",
    "    n, p = xs.shape\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            for k in range(p):\n",
    "                A[i,j] += (xs[i, k] - xs[j, k])**2\n",
    "            A[i,j] = np.sqrt(A[i,j])\n",
    "    A += A.T\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_sym\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vectorizing inner loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdist_vec(xs): \n",
    "    \"\"\"Vectorize inner loop.\"\"\"\n",
    "    n, p = xs.shape\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            A[i,j] = np.sqrt(np.sum((xs[i] - xs[j])**2))\n",
    "    A += A.T\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_vec\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Broadcasting and vectorizing\n",
    "\n",
    "Note that the broadcast version does twice as much work as it does not exploit symmetry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdist_numpy(xs):\n",
    "    \"\"\"Fully vectroized version.\"\"\"\n",
    "    return np.sqrt(np.square(xs[:, None] - xs[None, :]).sum(axis=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_numpy\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 squareform(func(xs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JIT with `numba`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the `numba.jit` decorator which will trigger generation and execution of compiled code when the function is first called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `jit` as a function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdist_numba_py = jit(pdist_py, nopython=True, cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_numba_py\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `jit` as a decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True, cache=True)\n",
    "def pdist_numba_py_1(xs):\n",
    "    \"\"\"Unvectorized Python.\"\"\"\n",
    "    n, p = xs.shape\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(n):\n",
    "            for k in range(p):\n",
    "                A[i,j] += (xs[i, k] - xs[j, k])**2\n",
    "            A[i,j] = np.sqrt(A[i,j])\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_numba_py_1\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Can we make the code faster?\n",
    "\n",
    "Note that in the inner loop, we are updating a matrix when we only need to update a scalar. Let's fix this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True, cache=True)\n",
    "def pdist_numba_py_2(xs):\n",
    "    \"\"\"Unvectorized Python.\"\"\"\n",
    "    n, p = xs.shape\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(n):\n",
    "            d = 0.0\n",
    "            for k in range(p):\n",
    "                d += (xs[i, k] - xs[j, k])**2\n",
    "            A[i,j] = np.sqrt(d)\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_numba_py_2\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Can we make the code even faster?\n",
    "\n",
    "We can also try to exploit symmetry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True, cache=True)\n",
    "def pdist_numba_py_sym(xs):\n",
    "    \"\"\"Unvectorized Python.\"\"\"\n",
    "    n, p = xs.shape\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            d = 0.0\n",
    "            for k in range(p):\n",
    "                d += (xs[i, k] - xs[j, k])**2\n",
    "            A[i,j] = np.sqrt(d)\n",
    "    A += A.T\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_numba_py_sym\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Does `jit` work with vectorized code?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdist_numba_vec = jit(pdist_vec, nopython=True, cache=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Only inner loop vectorized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r3 -n3 pdist_vec(xs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_numba_vec\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Does `jit` work with broadcasting?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdist_numba_numpy = jit(pdist_numpy, nopython=True, cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r3 -n3 pdist_numpy(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This raises an error because `numba` does not know how to deal with dummy axes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    %timeit -r3 -n3 pdist_numba_numpy(xs)\n",
    "except Exception as e:\n",
    "    print('Exception raised')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### We need to use `reshape` to broadcast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdist_numpy_(xs):\n",
    "    \"\"\"Fully vectroized version.\"\"\"\n",
    "    return np.sqrt(np.square(xs.reshape(n,1,p) - xs.reshape(1,n,p)).sum(axis=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdist_numba_numpy_ = jit(pdist_numpy_, nopython=True, cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r3 -n3 pdist_numpy_(xs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_numba_numpy_\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary\n",
    "\n",
    "- `numba` appears to work best when converting fairly explicit Python code\n",
    "- This might change in the future as the `numba` JIT compiler becomes more sophisticated\n",
    "- Always check optimized code for correctness\n",
    "- We can use `timeit` magic as a simple way to benchmark functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cython\n",
    "\n",
    "Cython is an Ahead Of Time (AOT) compiler. It compiles the code and replaces the function invoked with the compiled version.\n",
    "\n",
    "In the notebook, calling `%cython -a` magic shows code colored by how many Python C API calls are being made. You want to reduce the yellow as much as possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%cython -a \n",
    "\n",
    "import numpy as np\n",
    "\n",
    "def pdist_cython_1(xs):   \n",
    "    n, p = xs.shape\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            d = 0.0\n",
    "            for k in range(p):\n",
    "                d += (xs[i,k] - xs[j,k])**2\n",
    "            A[i,j] = np.sqrt(d)\n",
    "    A += A.T\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdist_base(xs):   \n",
    "    n, p = xs.shape\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            d = 0.0\n",
    "            for k in range(p):\n",
    "                d += (xs[i,k] - xs[j,k])**2\n",
    "            A[i,j] = np.sqrt(d)\n",
    "    A += A.T\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r3 -n3 pdist_base(xs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_cython_1\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n3 func(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cython with static types\n",
    "\n",
    "- We provide types for all variables so that Cython can optimize their compilation to C code.\n",
    "- Note `numpy` functions are optimized for working with `ndarrays` and have unnecessary overhead for scalars. We therefor replace them with math functions from the C `math` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%cython -a \n",
    "\n",
    "import cython\n",
    "import numpy as np\n",
    "cimport numpy as np\n",
    "from libc.math cimport sqrt, pow\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "def pdist_cython_2(double[:, :] xs):\n",
    "    cdef int n, p\n",
    "    cdef int i, j, k\n",
    "    cdef double[:, :] A\n",
    "    cdef double d\n",
    "    \n",
    "    n = xs.shape[0]\n",
    "    p = xs.shape[1]\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            d = 0.0\n",
    "            for k in range(p):\n",
    "                d += pow(xs[i,k] - xs[j,k],2)\n",
    "            A[i,j] = sqrt(d)\n",
    "    for i in range(1, n):\n",
    "        for j in range(i):\n",
    "            A[i, j] = A[j, i]            \n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = pdist_cython_2\n",
    "print(np.allclose(func(xs), sol))\n",
    "%timeit -r3 -n1 func(xs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
