Extra High performace python [answers].ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Day 3 - Performant Python\n",
    "\n",
    "Python has a reputation for being efficient in terms of programmer calories, i.e., Python allows you to do a lot with little code and succinctly express your ideas, but ineficient in terms of clock cycles. However, it turns out that this isn't necessarily true. In fact, there are many ways to make Python go fast, some of which we'll cover today.\n",
    "\n",
    "We'll start by talking about why Python is slow by default. Next, we'll talk about how to use numpy and numba to make it go fast. Over the course of the session you'll be given a few exercises."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## numpy\n",
    "In most cases you should first try numpy if you need your code to go faster.\n",
    "* numpy is fast because it's implemented in a manner optimized for numerical computing, whereas vanilla Python is designed for flexibility. For example, Python lists may contain any data types whereas numpy arrays have a fixed data type, e.g., float64.\n",
    "* numpy is written in C (mostly). Hence, working with numpy arrays means you're calling C code behind the covers.\n",
    "* The difference is even more pronounced when you aren't using the built-in Python functions, in this case `sum()`, since the built-in functions often are more optimized than what is easy to write yourself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# built-in sum\n",
    "import random\n",
    "n = 100000\n",
    "\n",
    "# vanilla Python sum\n",
    "data = [random.random() for _ in range(n)] # n random floats\n",
    "%timeit sum(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# numpy sum\n",
    "import numpy as np\n",
    "data = np.random.random(n) # n random floats\n",
    "%timeit data.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Python diff\n",
    "n = 100000\n",
    "def mydiff(l):\n",
    "    '''compute the derivative of l in-place.'''\n",
    "    if not len(l):\n",
    "        return l\n",
    "    for i in range(1, len(l)):\n",
    "        l[i] = l[i] - l[i-1]\n",
    "    l[0] = 0\n",
    "    return l\n",
    "\n",
    "data = [random.random() for _ in range(n)] # n random floats\n",
    "%timeit mydiff(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# numpy diff\n",
    "import numpy as np\n",
    "data = np.random.random(n) # n random floats\n",
    "%timeit np.diff(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Array Operations\n",
    "Like with Matlab we need to use array operations instead of loops for numpy to be fast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def diff(l):\n",
    "    '''compute the derivative of l in-place.'''\n",
    "    n = len(l)\n",
    "    rv = np.zeros(n-1)\n",
    "    if not len(l):\n",
    "        return rv\n",
    "    for i in range(n-1):\n",
    "        rv[i] = l[i+1] - l[i]\n",
    "    return rv\n",
    "\n",
    "def diff_array(l):\n",
    "    '''compute the derivative of l in-place using array operations.'''\n",
    "    n = len(l)\n",
    "    rv = np.zeros(n-1)\n",
    "    if not n:\n",
    "        return rv\n",
    "    rv[:] = l[1:] - l[:-1]\n",
    "    return rv\n",
    "\n",
    "n = 10000\n",
    "data = np.random.random(n) # n random floats\n",
    "%timeit diff(data)\n",
    "%timeit diff_array(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise: k-th order diff\n",
    "We've given you a recursive implementation of the k-th derivative. It relies on indexing individual elements and recursion. It's not very fast. Can you make it go faster? Make a note of the speed of our implementation and compute the speedup factor you got. Use array operations and possible drop the recursive call in favour of a loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kdiff(l, k):\n",
    "    '''compute the k-th derivative of l.'''\n",
    "    assert k % 1 == 0 and k >= 0, f'k must be a non-negative integer, but is {k}'\n",
    "    if k == 0:\n",
    "        return l \n",
    "    n = len(l)\n",
    "    rv = np.zeros(n-1)\n",
    "    if not len(l):\n",
    "        return rv\n",
    "    for i in range(n-1):\n",
    "        rv[i] = l[i+1] - l[i]\n",
    "    return kdiff(rv, k-1)\n",
    "\n",
    "def kdiff_fast(l, k):\n",
    "    '''compute the k-th derivative of l.'''\n",
    "    assert k % 1 == 0 and k >= 0, f'k must be a non-negative integer, but is {k}'\n",
    "    if k == 0:\n",
    "        return l\n",
    "    rv = np.array(l)\n",
    "    for i in range(1, k+1):\n",
    "        rv[i:] -= rv[i-1:-1]\n",
    "    return rv[k:]\n",
    "\n",
    "# test our naive implementation\n",
    "ans = kdiff(np.array([1, 2, 3]), 1)\n",
    "correct = np.array([1, 1])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = kdiff(np.array([1, 2, 3]), 2)\n",
    "correct = np.array([0])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = kdiff(np.array([1, 3, 4, 10]), 2)\n",
    "correct = np.array([-1, 5])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "# test the fast implementation\n",
    "ans = kdiff_fast(np.array([1, 2, 3]), 1)\n",
    "correct = np.array([1, 1])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = kdiff_fast(np.array([1, 2, 3]), 2)\n",
    "correct = np.array([0])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = kdiff_fast(np.array([1, 3, 4, 10]), 2)\n",
    "correct = np.array([-1, 5])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "print('all tests pass')\n",
    "\n",
    "# benchmark performance\n",
    "n = 10000\n",
    "data = np.random.random(n) # n random floats\n",
    "%timeit kdiff(data, 4)\n",
    "%timeit kdiff_fast(data, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## numba\n",
    "As long as your needs are covered by the numpy and scipy libraries you don't need to worry about performance. This is great news since this will probably be true for most of what you do. But we need a way to make Python go fast also in cases where numpy isn't enough! One common strategy is to write your performance-critical functions in C/Fortran and call those functions from Python via a wrapper. But, there's a better way.\n",
    "\n",
    "In this session, we're going to cover numba, a just-in-time compiler for Python. Essentially its a system that compiles Python code into machine code similar to what you would get from C. This means you can write Python code that runs almost as fast as C code.\n",
    "\n",
    "* Compiling functions can bring big performance gains.\n",
    "* Traditional languages like C are compiled ahead of time before the program can be started. Just-in-time compuling (jitting) is the process of compiling functions the first time a function is called during program execution.\n",
    "* Jitting can bring the benefits of compilation that languages like C enjoy without sacrificing the convenience of Python.\n",
    "* numba compiles optimized code for whatever data types your arguments are.\n",
    "\n",
    "### JIT-ing Arbitrary Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit, njit\n",
    "n = 100000\n",
    "\n",
    "@njit\n",
    "def diff(l):\n",
    "    if not len(l):\n",
    "        return l\n",
    "    for i in range(1, len(l)):\n",
    "        l[i] = l[i] - l[i-1]\n",
    "    l[0] = 0\n",
    "    return l\n",
    "\n",
    "data = np.random.random(n) # n random floats\n",
    "diff(data)\n",
    "%timeit diff(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise: cumsum\n",
    "Write a function `cumsum` that takes an array as its single argument and returns a new array of the same length, where the i-th element is the sum of all elements with indices smaller than or equal to i. Time it with an without the `@njit` decorator. How fast can you make it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def cumsum(a):\n",
    "    '''return the cumulative sum of a'''\n",
    "    rv = np.zeros(len(a))\n",
    "    if not len(a):\n",
    "        return rv\n",
    "    rv[0] = a[0]\n",
    "    for i in range(1, len(a)):\n",
    "        rv[i] = rv[i-1] + a[i]\n",
    "    return rv\n",
    "\n",
    "# test the function\n",
    "ans = cumsum(np.array([1, 2, 3]))\n",
    "correct = np.array([1, 3, 6])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = cumsum(np.array([1, -2, 3]))\n",
    "correct = np.array([1, -1, 2])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = cumsum(np.array([]))\n",
    "correct = np.array([])\n",
    "assert np.allclose(ans, correct), f'expected {correct}, but got {ans}'\n",
    "\n",
    "print('all tests pass')\n",
    "\n",
    "# time your code with and without the jit decorator\n",
    "n = 10000\n",
    "data = np.random.random(n)\n",
    "sum_max_3(data) # warm up the jit\n",
    "%timeit cumsum(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise: sum-max-3\n",
    "Write a function `sum_max_3` that takes an array as its single argument and returns the sum of the three largest values in the array. Time it with an without the `@njit` decorator. How fast can you make it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "@njit\n",
    "def sum_max_3(a):\n",
    "    '''return the sum of the 3 largest values in a'''\n",
    "    assert len(a) >= 3, 'input array must have at least 3 elements'\n",
    "    max4 = np.zeros(4)\n",
    "    max4[0] = -math.inf\n",
    "    max4[1:] = a[:3]\n",
    "    max4.sort()\n",
    "    for v in a[3:]:\n",
    "        max4[0] = v\n",
    "        if max4[0] > max4[1]:\n",
    "            max4[0], max4[1] = max4[1], max4[0]\n",
    "            if max4[1] > max4[2]:\n",
    "                max4[1], max4[2] = max4[2], max4[1]\n",
    "                if max4[2] > max4[3]:\n",
    "                    max4[2], max4[3] = max4[3], max4[2]\n",
    "    return max4[1:].sum()\n",
    "\n",
    "# test the function\n",
    "ans = sum_max_3(np.array([1, 2, 3]))\n",
    "correct = 6\n",
    "assert ans == correct, f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = sum_max_3(np.array([1, 1, 2, 3]))\n",
    "correct = 6\n",
    "assert ans == correct, f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = sum_max_3(np.array([1, 2, 3, 10]))\n",
    "correct = 15\n",
    "assert ans == correct, f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = sum_max_3(np.array([-1, 11, 30, -40, 10]))\n",
    "correct = 51\n",
    "assert ans == correct, f'expected {correct}, but got {ans}'\n",
    "\n",
    "ans = sum_max_3(np.array([-1, -11, -30, -40, 10]))\n",
    "correct = -2\n",
    "assert ans == correct, f'expected {correct}, but got {ans}'\n",
    "\n",
    "print('all tests pass')\n",
    "\n",
    "# time your code with and without the jit decorator\n",
    "n = 10000\n",
    "data = np.random.random(n)\n",
    "sum_max_3(data) # warm up the jit\n",
    "%timeit sum_max_3(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba and Objects\n",
    "In many cases its good practice to write object oriented code. For example, if you have a program dealing with 2D vectors its often easier to create a class representing the vector and then write functions that take the object as its argument rather than passing around the vector coordinates explicitly. Let's look at an example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom dtypes\n",
    "It's often a good idea to use containers, e.g., dicts, tuples, objects, to pass around collections of variables since it makes the code easier to read and manage. Unfortunately numba doesn't support these well. However, we can use numpy custom data types (dtypes) to the same effect! For the C programmers, custom dtypes are essentially C structs, i.e., not full-fledged objects. The point of custom dtypes is to organize your code without losing the performance benefit of numba.\n",
    "\n",
    "#### 2D Vector Datatype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vector2d_dtype = np.dtype([\n",
    "    ('x', np.float64),\n",
    "    ('y', np.float64),\n",
    "])\n",
    "\n",
    "def create_vectors(xs, ys):\n",
    "    '''create an array of vectors (xs[i], ys[i]) for all i'''\n",
    "    return np.array([(x, y) for x, y in zip(xs, ys)], dtype=vector2d_dtype)\n",
    "\n",
    "xs = [1, 2]\n",
    "ys = [3, 4]\n",
    "vs = create_vectors(xs, ys)\n",
    "\n",
    "print(vs)\n",
    "\n",
    "v = vs[0] # get the first point\n",
    "print(v)\n",
    "print(v['x'], v['y']) # access elements of the \"object\"\n",
    "\n",
    "# you can also access an array of all x values\n",
    "print('xs:', vs['x'])\n",
    "print('xs:', vs['y'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Shortest distance problem for 2D vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "@njit # show that this is needed\n",
    "def distance(v1, v2):\n",
    "    '''return the distance between vectors v1 and v2'''\n",
    "    return math.sqrt(math.pow(v1['x'] - v2['x'], 2) + math.pow(v1['y'] - v2['y'], 2))\n",
    "\n",
    "@njit\n",
    "def shortest_distance(vs):\n",
    "    '''return the shortest distance between any pair of vectors'''\n",
    "    n = len(vs)\n",
    "    mind = math.inf\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            v1, v2 = vs[i], vs[j]\n",
    "            # d = math.sqrt(math.pow(v1['x']-v2['x'], 2) + math.pow(v1['y']-v2['y'], 2))\n",
    "            d = distance(v1, v2)\n",
    "            mind = min(d, mind)\n",
    "    return mind\n",
    "\n",
    "vectors = create_vectors(np.random.random(n), np.random.random(n))\n",
    "shortest_distance(vectors) # warm up the jit\n",
    "%timeit shortest_distance(vectors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implementation using dicts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_vectors_dct(xs, ys):\n",
    "    '''create an array of vectors {'x': x, 'y': y} for all i'''\n",
    "    return [{'x': x, 'y': y} for x, y in zip(xs, ys)]\n",
    "\n",
    "def distance_dct(v1, v2):\n",
    "    '''return the distance between vectors v1 and v2'''\n",
    "    return math.sqrt(math.pow(v1['x']-v2['x'], 2) + math.pow(v1['y']-v2['y'], 2))\n",
    "\n",
    "def shortest_distance_dct(vs):\n",
    "    '''return the shortest distance between any pair of vectors.'''\n",
    "    n = len(vs)\n",
    "    mind = math.inf\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            v1, v2 = vs[i], vs[j]\n",
    "            d = distance_dct(v1, v2)\n",
    "            mind = min(d, mind)\n",
    "    return mind\n",
    "\n",
    "# benchmark performance\n",
    "n = 100 # number of vectors\n",
    "vectors = create_vectors_dct(np.random.random(n), np.random.random(n))\n",
    "%timeit shortest_distance_dct(vectors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implementation using objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "class Vector2D:\n",
    "    def __init__(self, x, y):\n",
    "        self.x = x\n",
    "        self.y = y\n",
    "    def __repr__(self):\n",
    "        return f'Vector2D(x: {self.x}, y={self.y})'\n",
    "\n",
    "def distance(v1, v2):\n",
    "    '''return the distance between vectors v1 and v2'''\n",
    "    return math.sqrt(math.pow(v1.x-v2.x, 2) + math.pow(v1.y-v2.y, 2))\n",
    "\n",
    "def shortest_distance(vs):\n",
    "    '''return the shortest distance between any pair of vectors.'''\n",
    "    n = len(vs)\n",
    "    mind = math.inf\n",
    "    for i in range(n):\n",
    "        for j in range(i+1, n):\n",
    "            v1, v2 = vs[i], vs[j]\n",
    "            # d = math.sqrt(math.pow(v1.x-v2.x, 2) + math.pow(v1.y-v2.y, 2))\n",
    "            d = distance(v1, v2)\n",
    "            mind = min(d, mind)\n",
    "    return mind\n",
    "\n",
    "# benchmark performance\n",
    "n = 100 # number of vectors\n",
    "vectors = [Vector2D(random.random(), random.random()) for _ in range(n)]\n",
    "%timeit shortest_distance(vectors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom ufuncs\n",
    "In numpy there are lots of functions that can take either scalar or vector arguments. These are referred to as ufuncs and they used to be quite complicated to make."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import vectorize\n",
    "\n",
    "@vectorize\n",
    "def divide_by_larger(a, b):\n",
    "    '''return the smaller number of a and b divided by the larger number'''\n",
    "    if a > b:\n",
    "        return b / a\n",
    "    return a / b\n",
    "\n",
    "a = 2\n",
    "b = 3\n",
    "rv = divide_by_larger(a, b)\n",
    "print(rv)\n",
    "\n",
    "n = 100000\n",
    "a = 2 * np.ones(n)\n",
    "b = 3 * np.ones(n)\n",
    "rv = divide_by_larger(a, b)\n",
    "print(rv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automagic Parallelization\n",
    "Numba also allows us to write code that automatically executes in parallel.\n",
    "* Automagically parallelize parallelize your code to exploit multi-core CPUs and even GPUs.\n",
    "* https://numba.pydata.org/numba-doc/dev/user/vectorize.html#dynamic-universal-functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import float64\n",
    "\n",
    "@vectorize([float64(float64, float64)], target='parallel')\n",
    "def divide_by_larger(a, b):\n",
    "    '''return the smaller number of a and b divided by the larger number'''\n",
    "    if a > b:\n",
    "        return b / a\n",
    "    return a / b\n",
    "\n",
    "n = 100000\n",
    "a = 2 * np.ones(n)\n",
    "b = 3 * np.ones(n)\n",
    "%timeit divide_by_larger(a, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercises\n",
    "* Give them a function written using loops that they should rewrite using array operations.\n",
    "* Consider giving them a problem and a naive implementation and have them make it go fast. Should be something a bit complicated. Should be some speedup from using njit (or njit doens't work without changes) but there should be additional things to do as well. What's the speedup factor compared to the naive implementation?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# naive implementation\n",
    "\n",
    "import math\n",
    "\n",
    "vector2d_dtype = np.dtype([\n",
    "    ('x', np.float64),\n",
    "    ('y', np.float64),\n",
    "])\n",
    "\n",
    "def create_vectors(xs, ys):\n",
    "    '''create an array of vectors (xs[i], ys[i]) for all i'''\n",
    "    return np.array([(x, y) for x, y in zip(xs, ys)], dtype=vector2d_dtype)\n",
    "\n",
    "def magnitude(v):\n",
    "    '''return the magntiude (length) of the vector v.'''\n",
    "    return math.sqrt(math.pow(v['x'], 2) + math.pow(v['y'], 2))\n",
    "\n",
    "def longest_in_quadrant(vs, xsign, ysign):\n",
    "    '''return the longest vector in vs with x (y) coordinate sign equal to xsign (ysign).'''\n",
    "    rv = None\n",
    "    for v in vs:\n",
    "        if v['x'] * xsign < 0: # differing sign on x\n",
    "            continue\n",
    "        if v['y'] * ysign < 0: # differing sign on y\n",
    "            continue\n",
    "        if rv is None: # no previously found candidate\n",
    "            rv = v\n",
    "            continue\n",
    "        if magnitude(v) > magnitude(rv):\n",
    "            rv = v\n",
    "    return rv\n",
    "\n",
    "def longest_in_quadrants(vs):\n",
    "    '''return an array composed of the longest vector in v in the \n",
    "    top right, bottom right, bottom left, top left quadrant, in that order.\n",
    "    \n",
    "    '''\n",
    "    return [\n",
    "        longest_in_quadrant(vs, 1, 1),\n",
    "        longest_in_quadrant(vs, 1, -1),\n",
    "        longest_in_quadrant(vs, -1, -1),\n",
    "        longest_in_quadrant(vs, -1, 1),\n",
    "    ]\n",
    "\n",
    "n = 1000\n",
    "vectors = create_vectors(np.random.random(n)*2-1, np.random.random(n)*2-1)\n",
    "ans = longest_in_quadrants(vectors)\n",
    "print(ans)\n",
    "%timeit longest_in_quadrants(vectors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# our optimized implementation\n",
    "from numba import vectorize, int16, float64\n",
    "\n",
    "@vectorize([float64(float64, float64)], target='parallel')\n",
    "def magnitudes(x, y):\n",
    "    '''return the magntiude (length) of the vector with coordinates (x, y).'''\n",
    "    return math.sqrt(math.pow(x, 2) + math.pow(y, 2))\n",
    "\n",
    "@vectorize([int16(float64, float64)], target='parallel')\n",
    "def quadrants(x, y):\n",
    "    '''return the quadrant index of the vector with coordinates (x, y)'''\n",
    "    if x > 0:\n",
    "        if y > 0:\n",
    "            return 1\n",
    "        return 2\n",
    "    if y < 0:\n",
    "        return 3\n",
    "    return 4\n",
    "\n",
    "@jit\n",
    "def longest_in_quadrants(vs):\n",
    "    '''return an array composed of the longest vector in v in the \n",
    "    top right, bottom right, bottom left, top left quadrant, in that order.\n",
    "    \n",
    "    '''\n",
    "    n = len(vs)\n",
    "    ms = magnitudes(vs['x'], vs['y']) # compute magnitudes\n",
    "    qs = quadrants(vs['x'], vs['y']) # compute magnitude in parallel\n",
    "    rv = np.zeros(4, dtype=vector2d_dtype) # store longest vector found so far\n",
    "    rvm = np.zeros(4) # store magnitude of rv vectors\n",
    "    for i in range(n):\n",
    "        m = ms[i]\n",
    "        q = qs[i]\n",
    "        if m > rvm[q-1]: # compare vector length\n",
    "            rv[q-1] = vs[i]\n",
    "            rvm[q-1] = m\n",
    "    return rv\n",
    "\n",
    "n = 1000\n",
    "vectors = create_vectors(np.random.random(n)*2-1, np.random.random(n)*2-1)\n",
    "longest_in_quadrants(vectors) # warm up the jit\n",
    "%timeit longest_in_quadrants(vectors)"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": "9a7f0b82bed342f090fa8f7809d84653",
   "lastKernelId": "8e7b52e9-f9a4-4f96-8aaa-2c550f7fa28a"
  },
  "kernelspec": {
   "display_name": "dynpie3-2020a",
   "language": "python",
   "name": "dynpie3-2020a"
  },
  "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}