numba_on_dask.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "from scipy.sparse import coo_matrix\n",
    "import dask.array as da\n",
    "\n",
    "from numba import jit, prange, guvectorize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# sparse matrix\n",
    "ds = xr.open_dataset(\"weights.nc\")\n",
    "n_s = ds.dims['n_s']\n",
    "col = ds['col'].values - 1\n",
    "row = ds['row'].values - 1\n",
    "S = ds['S'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# input data\n",
    "data = np.random.rand(500, 240000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True, nogil=True)\n",
    "def sparse_dot(data_out, data, col, row, S):\n",
    "    for j in range(data.shape[0]):\n",
    "        for i in range(S.size):\n",
    "            data_out[j, row[i]] += data[j, col[i]]*S[i]\n",
    "\n",
    "@jit(nopython=True, nogil=True, parallel=True)\n",
    "def sparse_dot_pa(data_out, data, col, row, S):\n",
    "    for j in prange(data.shape[0]):\n",
    "        for i in range(S.size):\n",
    "            data_out[j, row[i]] += data[j, col[i]]*S[i]\n",
    "\n",
    "# cannot use nopython mode to create array\n",
    "@jit(nogil=True)\n",
    "def apply_A(data, parallel=False):\n",
    "    if parallel:\n",
    "        func = sparse_dot_pa\n",
    "    else:\n",
    "        func = sparse_dot\n",
    "    data_out = np.zeros([data.shape[0], 120000])\n",
    "    func(data_out, data, col, row, S)    # use global col, row, S here\n",
    "    return data_out\n",
    "\n",
    "out_numba = apply_A(data) # reference result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "697 ms ± 22 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "315 ms ± 3.86 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "# reference performance\n",
    "%timeit apply_A(data)\n",
    "%timeit apply_A(data, parallel=True) # obvious speed-up with parallelization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# guvectorize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((500, 240000), (480000,), (480000,), (480000,), (500, 120000))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# shapes of input arguments that guvectorize would want to know\n",
    "data.shape, col.shape, row.shape, S.shape, out_numba.shape "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the output grid size 120000 is not contained in any input arguments,\n",
    "# so I just do something silly: create a dummy array to pass the shape info\n",
    "shape_arr = np.arange(120000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "@guvectorize([\"void(float64[:], int32[:], int32[:], float64[:], int64[:], float64[:])\"],\n",
    "             \"(n),(k),(k),(k),(m)->(m)\",\n",
    "             target='parallel')\n",
    "def sparse_dot_guvec(data, col, row, S, shape_arr, data_out):\n",
    "    # only one loop for the grid dimension\n",
    "    # let numba vectorize over extra dimension\n",
    "    for i in range(S.size):\n",
    "        data_out[row[i]] += data[col[i]]*S[i]\n",
    "        \n",
    "out_guvec = sparse_dot_guvec(data, col, row, S, shape_arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# result is correct\n",
    "np.array_equal(out_numba, out_guvec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Doesn't seem to be faster than mannually parallelizing with `prange`. So we will still use the previous simple way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "352 ms ± 4.99 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit sparse_dot_guvec(data, col, row, S, shape_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Apply numba function on dask array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dask.array<array, shape=(500, 240000), dtype=float64, chunksize=(50, 240000)>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_dask = da.from_array(data, chunks=(50, 240000))\n",
    "data_dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray 'data' (extra_dims: 500, grid_dims: 240000)>\n",
       "dask.array<shape=(500, 240000), dtype=float64, chunksize=(50, 240000)>\n",
       "Dimensions without coordinates: extra_dims, grid_dims"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dr_dask = xr.DataArray(data_dask, \n",
    "                       dims=['extra_dims', 'grid_dims'],\n",
    "                       name='data')\n",
    "dr_dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dr_out_pa = xr.apply_ufunc(apply_A, dr_dask, \n",
    "                           input_core_dims=[['grid_dims']],\n",
    "                           output_core_dims=[['out_grid']],\n",
    "                           output_sizes={'out_grid': 120000},\n",
    "                           dask='parallelized', \n",
    "                           output_dtypes=[float])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# result is correct\n",
    "np.array_equal(dr_out_pa, out_numba)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The serial case is slower than the pure numpy version, and the parallel efficiency is not too great."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.12 s ± 14.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "809 ms ± 17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "772 ms ± 20.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "for n in [1, 2, 4]:\n",
    "    %timeit dr_out_pa.compute(num_workers=n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Much faster without concatenating. But still not comparable with Numba."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "872 ms ± 18 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "550 ms ± 18.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "511 ms ± 14.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "for n in [1, 2, 4]:\n",
    "    %timeit dr_out_pa.persist(num_workers=n)"
   ]
  }
 ],
 "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.2"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}