notebooks/08.Make.generalized.ufuncs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generalized ufuncs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've just seen how to make our own ufuncs using `vectorize`, but what if we need something that can operate on an input array in any way that is not element-wise?\n",
    "\n",
    "Enter `guvectorize`.  \n",
    "\n",
    "There are several important differences between `vectorize` and `guvectorize` that bear close examination.  Let's take a look at a few simple examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy\n",
    "from numba import guvectorize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "@guvectorize('int64[:], int64, int64[:]', '(n),()->(n)')\n",
    "def g(x, y, result):\n",
    "    for i in range(x.shape[0]):\n",
    "        result[i] = x[i] + y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Declaration of input/output layouts\n",
    "* No return statements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = numpy.arange(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the cell below we call the function `g` with a preallocated array for the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "result = numpy.zeros_like(x)\n",
    "result = g(x, 5, result)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But wait!  We can still call `g` as if it were defined as `def g(x, y)`\n",
    "\n",
    "```python\n",
    "res = g(x, 5)\n",
    "print(res)\n",
    "```\n",
    "\n",
    "We don't recommend this as it can have unintended consequences if some of the elements of the `results` array are not operated on by the function `g`.  (The advantage is that you can preserve existing interfaces to previously written functions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "@guvectorize('float64[:,:], float64[:,:], float64[:,:]', \n",
    "            '(m,n),(n,p)->(m,p)')\n",
    "def matmul(A, B, C):\n",
    "    m, n = A.shape\n",
    "    n, p = B.shape\n",
    "    for i in range(m):\n",
    "        for j in range(p):\n",
    "            C[i, j] = 0\n",
    "            for k in range(n):\n",
    "                C[i, j] += A[i, k] * B[k, j]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "a = numpy.random.random((500, 500))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "out = matmul(a, a, numpy.zeros_like(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%timeit matmul(a, a, numpy.zeros_like(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%timeit a @ a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And it also supports the `target` keyword argument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def g(x, y, res):\n",
    "    for i in range(x.shape[0]):\n",
    "        res[i] = x[i] + numpy.exp(y)\n",
    "        \n",
    "g_serial = guvectorize('float64[:], float64, float64[:]', \n",
    "                       '(n),()->(n)')(g)\n",
    "g_par = guvectorize('float64[:], float64, float64[:]', \n",
    "                    '(n),()->(n)', target='parallel')(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%timeit res = g_serial(numpy.arange(1000000).reshape(1000, 1000), 3)\n",
    "%timeit res = g_par(numpy.arange(1000000).reshape(1000, 1000), 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [Exercise: Writing signatures](./exercises/08.GUVectorize.Exercises.ipynb#Exercise:-2D-Heat-Transfer-signature)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's up with these boundary conditions?\n",
    "\n",
    "```python\n",
    "for i in range(I):\n",
    "        Tn[i, 0] = T[i, 0]\n",
    "        Tn[i, J - 1] = Tn[i, J - 2]\n",
    "\n",
    "    for j in range(J):\n",
    "        Tn[0, j] = T[0, j]\n",
    "        Tn[I - 1, j] = Tn[I - 2, j]\n",
    "```\n",
    "\n",
    "We don't pass in `Tn` explicitly, which means Numba allocates it for us (thanks!) but it's allocated using `numpy.empty_like` so if we don't touch every value in `Tn` in the function, those empty values will stick around and cause trouble.  \n",
    "\n",
    "Solutions?  The one above, or pass it in explicitly after doing something like `Tn = Ti.copy()`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## [Exercise: Remove the vanilla loops](./exercises/08.GUVectorize.Exercises.ipynb#Exercise:-2D-Heat-Transfer-Time-loop)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The example above loops in time outside of the `vectorize`d function.  That means it's looping in vanilla Python which is not the fastest thing in the world.  \n",
    "\n",
    "Move the time loop inside the function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demo: Why not `jit` the `run_ftcs` function?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because, at the moment, it won't work.  (bummer)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@guvectorize('float64[:,:], float64[:,:]', '(n,n)->(n,n)')\n",
    "def gucopy(a, b):\n",
    "    I, J = a.shape\n",
    "    for i in range(I):\n",
    "        for j in range(J):\n",
    "            b[i, j] = a[i, j]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from numba import jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def make_a_copy():\n",
    "    a = numpy.random.random((25,25))\n",
    "    b = gucopy(a)\n",
    "    \n",
    "    return a, b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "a, b = make_a_copy()\n",
    "assert numpy.allclose(a, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "make_a_copy.inspect_types()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}