apply_ufunc_with_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",
    "from dask.diagnostics import ProgressBar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Update: This kind of works now. But the performance is worse than Numba.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just copy the code from [sparse_dot_benchmark.ipynb](./sparse_dot_benchmark.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "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\n",
    "A = coo_matrix((S, (row, col))) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(500, 240000)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = np.random.rand(500, A.shape[1])\n",
    "data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 475 ms, sys: 167 ms, total: 642 ms\n",
      "Wall time: 650 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(500, 120000)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# reference result and baseline performance\n",
    "%time out_scipy = A.dot(data.T).T\n",
    "out_scipy.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# apply_ufunc on numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def apply_A(data):\n",
    "    # use global A here!\n",
    "    return A.dot(data.T).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Has the same performance as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 466 ms, sys: 165 ms, total: 631 ms\n",
      "Wall time: 643 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0.11250547,  0.62152447,  0.32295978, ...,  0.27827256,\n",
       "         0.34809867,  0.36379499],\n",
       "       [ 0.64895496,  0.26802886,  0.41272402, ...,  0.46974352,\n",
       "         0.64499186,  0.64985559],\n",
       "       [ 0.82787185,  0.33176759,  0.56978941, ...,  0.71662387,\n",
       "         0.43066475,  0.18261389],\n",
       "       ..., \n",
       "       [ 0.35410965,  0.36236262,  0.7849376 , ...,  0.50895257,\n",
       "         0.3384691 ,  0.73448701],\n",
       "       [ 0.87502444,  0.58880469,  0.35482145, ...,  0.69282072,\n",
       "         0.40735966,  0.61793078],\n",
       "       [ 0.24804293,  0.48838714,  0.46165374, ...,  0.6340203 ,\n",
       "         0.87557678,  0.82797878]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time xr.apply_ufunc(apply_A, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# apply_ufunc on xarray DataArray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray 'data' (extra_dims: 500, grid_dims: 240000)>\n",
       "array([[ 0.021726,  0.311896,  0.695838, ...,  0.178822,  0.947601,  0.153578],\n",
       "       [ 0.906841,  0.161121,  0.273547, ...,  0.657841,  0.965106,  0.675433],\n",
       "       [ 0.920618,  0.71062 ,  0.13484 , ...,  0.397721,  0.204345,  0.13308 ],\n",
       "       ..., \n",
       "       [ 0.513338,  0.047198,  0.542131, ...,  0.11914 ,  0.883951,  0.778235],\n",
       "       [ 0.996765,  0.825025,  0.589761, ...,  0.281366,  0.581218,  0.696507],\n",
       "       [ 0.034942,  0.963049,  0.411179, ...,  0.926111,  0.872794,  0.905661]])\n",
       "Coordinates:\n",
       "  * extra_dims  (extra_dims) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 ...\n",
       "  * grid_dims   (grid_dims) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 ..."
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dr = xr.DataArray(data, \n",
    "                  dims=['extra_dims', 'grid_dims'],\n",
    "                  coords=[np.arange(500), np.arange(240000)],\n",
    "                  name='data'\n",
    "                 )\n",
    "dr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Has the same performance as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 516 ms, sys: 193 ms, total: 709 ms\n",
      "Wall time: 730 ms\n"
     ]
    }
   ],
   "source": [
    "%time dr_out = xr.apply_ufunc(apply_A, dr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# result is correct\n",
    "np.array_equal(dr_out.data, out_scipy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use dask array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dask.array<array, shape=(500, 240000), dtype=float64, chunksize=(50, 240000)>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_dask = da.from_array(data, chunks=(50, 240000))\n",
    "data_dask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Horribly slow...**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.93 s, sys: 1.05 s, total: 4.98 s\n",
      "Wall time: 5.09 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(500, 120000)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time out_dask = apply_A(data_dask)\n",
    "out_dask.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# result is correct\n",
    "np.array_equal(out_dask, out_scipy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# apply_ufunc on dask array, serial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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": 13,
     "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": "markdown",
   "metadata": {},
   "source": [
    "**As slow as the previous case...**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.98 s, sys: 1.05 s, total: 5.03 s\n",
      "Wall time: 5.13 s\n"
     ]
    }
   ],
   "source": [
    "%time dr_out_dask = xr.apply_ufunc(apply_A, dr_dask, dask='allowed')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# result is correct\n",
    "np.array_equal(dr_out_dask, out_scipy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# apply_ufunc on dask array, parallelized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 140 ms, sys: 39.4 ms, total: 179 ms\n",
      "Wall time: 178 ms\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "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": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray 'data' (extra_dims: 500, out_grid: 120000)>\n",
       "dask.array<shape=(500, 120000), dtype=float64, chunksize=(50, 120000)>\n",
       "Dimensions without coordinates: extra_dims, out_grid"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dr_out_pa # not computed yet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Much faster than the previous case, but still slower than the pure numpy case.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "731 ms ± 10.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit dr_out_pa.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# result is correct\n",
    "np.array_equal(dr_out_pa.data, out_scipy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Low-level implementation dask delayed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from dask import delayed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Manully break data into 10 chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50, 240000)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i = 0 \n",
    "data[50*i:50*(i+1), :].shape # input chunk shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50, 120000)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "apply_A(data[50*i:50*(i+1), :]).shape # output chunk shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A list of output results from 10 chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Delayed('apply_A-8dbdf735-6195-4d95-b157-50000f70eb0b'),\n",
       " Delayed('apply_A-e1adddc9-9ae7-4956-9c31-a439e9ab8b72'),\n",
       " Delayed('apply_A-a330b035-3b1b-41a9-a020-f64f69eba30a'),\n",
       " Delayed('apply_A-0ba32025-8e20-4ba9-81ec-806883ae0a11'),\n",
       " Delayed('apply_A-e2163887-b057-494c-b27d-a6864fd1e786'),\n",
       " Delayed('apply_A-03707b33-db0b-4679-8c23-3c1a4d609131'),\n",
       " Delayed('apply_A-99771b3c-6b08-4688-abba-a5a66d5468e8'),\n",
       " Delayed('apply_A-2e96076a-10d1-43bf-a824-590f8e1d339c'),\n",
       " Delayed('apply_A-2d3ed76b-4d70-4ea9-88d7-9da3b77c5d9d'),\n",
       " Delayed('apply_A-268087fd-d16e-4d5a-90b1-46bdfd400940')]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_delayed = [delayed(apply_A)(data[50*i:50*(i+1), :]) for i in range(10)]\n",
    "out_delayed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert to dask arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>,\n",
       " dask.array<from-value, shape=(50, 120000), dtype=float64, chunksize=(50, 120000)>]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dask_subarrays = [da.from_delayed(d, (50, 120000), float) for d in out_delayed]\n",
    "dask_subarrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Merge into a single dask array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dask.array<concatenate, shape=(500, 120000), dtype=float64, chunksize=(50, 120000)>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dask_merge = da.concatenate(dask_subarrays, axis=0)\n",
    "dask_merge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array_equal(dask_merge, out_scipy) # result is correct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parallelization does give some speed-up. But the serial version is slower than pure numpy version, and the parallel efficiency is worse than Numba."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "725 ms ± 16.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "529 ms ± 5.83 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "505 ms ± 6.37 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit dask_merge.compute(num_workers=1)\n",
    "%timeit dask_merge.compute(num_workers=2)\n",
    "%timeit dask_merge.compute(num_workers=4)"
   ]
  }
 ],
 "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
}