sparse_dot_benchmark.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",
    "from numba import jit, prange"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Performance of different methods:\n",
    "\n",
    "| Method | Time (ms)  |\n",
    "|--------|------------|\n",
    "| scipy.sparse.coo_matrix.dot      | 650 |\n",
    "| scipy.sparse.csr_matrix.dot      | 1000|\n",
    "| scipy.sparse.csc_matrix.dot      | 1000|\n",
    "| Numba loop for COO algorithm     |  650|\n",
    "| **Numba parallel loop (2 cores)**|**300**|\n",
    "| Cython loop for COO algorithm    | 650 |\n",
    "| ESMPy native method              | 3000 |\n",
    "\n",
    "Data used for the benchmark:\n",
    "\n",
    "| Data   | grid size | extra dim | flattened shape | size  | I/O time (ms)  |\n",
    "|--------|-----------|-----------|-----------------|-------| -------------- |\n",
    "| Input  | 600x400   | 500       | 500x240000      | 1 GB  | 600  | \n",
    "| Output | 400x300   | 500       | 500x120000      | 0.5 GB| 250  |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Key points:**\n",
    "\n",
    "- All Python methods (Scipy/Numba/Cython) are much faster than ESMPy's original Fortran regridding engine ([ESMPy timing](http://xesmf.readthedocs.io/en/latest/Backend.html)). Thus xESMF must have its own regridding backend, instead of completely relying on ESMPy.\n",
    "- Given Numba's excellent serial and parallel performance, we plan to use it as the default backend for xESMF in the next version.\n",
    "- The regridding efficiency is getting close to I/O efficiency, so further optimizing regridding cannot significantly speed-up the entire workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sparse matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read a typical weight file generated by ESMPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:  (n_s: 480000)\n",
       "Dimensions without coordinates: n_s\n",
       "Data variables:\n",
       "    S        (n_s) float64 ...\n",
       "    col      (n_s) int32 ...\n",
       "    row      (n_s) int32 ..."
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = xr.open_dataset(\"weights.nc\")\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(120000, 240000)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = coo_matrix((S, (row, col))) \n",
    "A.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.text.Text at 0x112846be0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZsAAADUCAYAAABK4Y71AAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAG/xJREFUeJzt3X+0ldV95/H3R4ioVRCUUgI6YKTtQms0sAixTse1nBGi\naTAzTEImGUlkZLq01qQr04E6q3Glf4w2raZOq1MajWiNoiZZMq1oCSaTrEmEgL8QLeUqWGFQiBgw\nbePkku/88ewTHw73Xs6559c+53xea511nrPP83tdng97P/vZRxGBmZlZKx3X6R0wM7Pe57AxM7OW\nc9iYmVnLOWzMzKzlHDZmZtZyDhszM2s5hw0gaaGk7ZIGJK3o9P60gqRdkrZKekbS5lQ2SdJ6STvS\n+8TS/CvT+dguaUGpfE5az4Ck2yQplY+TtCaVb5Q0o93HWAtJd0naJ+n5UllbzoOkpWkbOyQtbc8R\n12aY83KjpD3pb+YZSZeVvuv58yLpDEnfkvSCpG2Srk/lff/3MioR0dcvYAzwEnAWcDzwLDC70/vV\nguPcBZxeVfZHwIo0vQK4OU3PTudhHDAznZ8x6btNwHxAwDrgg6n8GuB/puklwJpOH/Mw5+E3gPcB\nz7fzPACTgJfT+8Q0PbHT5+MY5+VG4HNDzNsX5wWYCrwvTZ8C/H069r7/exnNyzUbmAcMRMTLEfH/\ngAeARR3ep3ZZBKxO06uBK0rlD0TE2xGxExgA5kmaCoyPiCej+BdxT9UylXU9DFxS+d9bTiLiO8CB\nquJ2nIcFwPqIOBARbwLrgYXNP8LRGea8DKcvzktE7I2Ip9L0W8CLwDT89zIqDpvij+fV0ufdqazX\nBPBNSVskLU9lUyJib5p+DZiSpoc7J9PSdHX5EctExCBwEDit2QfRIu04D936d3adpOdSM1uluajv\nzktq3roA2Ij/XkbFYdM/LoqI84EPAtdK+o3yl+l/XH0/dpHPwxHuoGhePh/YC/xJZ3enMySdDHwN\n+ExEHCp/57+X2jlsYA9wRunz9FTWUyJiT3rfB3yDovnw9VTFJ73vS7MPd072pOnq8iOWkTQWmAC8\n0YpjaYF2nIeu+zuLiNcj4nBE/Az4S4q/Geij8yLpXRRBc19EfD0V++9lFBw28ANglqSZko6nuEm3\ntsP71FSSfkHSKZVp4FLgeYrjrPRyWQo8kqbXAktST5mZwCxgU2o6OCRpfmpXvrJqmcq6FgNPpP/1\ndYN2nIfHgUslTUzNUZemsmxVLqjJRyj+ZqBPzks6hjuBFyPiltJX/nsZjU73UMjhBVxG0dPkJeCG\nTu9PC47vLIpeMs8C2yrHSNE2vAHYAXwTmFRa5oZ0PraTes6k8rkUF52XgD8DlMpPAB6iuCm6CTir\n08c9zLm4n6JJ6KcU7eDL2nUegKtS+QDw6U6fixrOy73AVuA5iovi1H46L8BFFE1kzwHPpNdl/nsZ\n3atywGZmZi3jZjQzM2s5h42ZmbWcw8bMzFrOYWNmZi3X02GjOgbYLD1VbyU+L0PzeRmaz8vRfE4K\nPRs2ksYAf07xxPxs4OOSZo+wiP8ghubzMjSfl6H5vBzN54QeDhv6e4BNM7Os9OxzNpIWAwsj4j+l\nz/8ReH9E/HZpnuWk/3Ucd+L4OWMn/GJD2/y1aRMaWj5H+/fvZ/LkyZ3ejez4vAzN5+VovX5OtmzZ\n8sOIOOYBjm3HzuQqIlYBqwDGTZ0VU5d+qaH1/TC977rp8gb3zMysO0h6pZb5erkZrWMD2c1Y8TfM\nWPE37diUmVlX6OWw6fgAmw4dM7NCzzajRcSgpN+mGCl1DHBXRGzrxL5UAsfNa2bWr3o2bAAi4lHg\n0U7vR4VDx8z6VS83o2XLzWtm1m8cNh3k0DGzfuGwyYBDx8x6ncMmyeGBTIeOmfWqnh1BoF5z586N\nzZs3A2RzwXdHAjPLnaQtETH3mPM5bArlsCnLIXgcOmaWq1rDxs1ox5DDhd7Na2bW7VyzSYar2ZTl\ncsHPIQDNzMDNaHWrJWwqHDpmZgWHTZ3qCZsKh46Z9TuHTZ1GEzZlOQSPQ8fM2s0dBNoshwu9OxKY\nWa5cs0kardmU5XLBzyEAzay3uRmtTs0MmwqHjpn1OodNnVoRNhUOHTPrVQ6bOrUybMpyCB6Hjpk1\nizsIZCqHC707EphZu7lmk7SrZlOWywU/hwA0s+7kZrQ6dSJsKhw6Ztatsm1Gk3SGpG9JekHSNknX\np/JJktZL2pHeJ5aWWSlpQNJ2SQtK5XMkbU3f3SZJqXycpDWpfKOkGe0+znrsuunyLC70bl4zs1Zp\ne81G0lRgakQ8JekUYAtwBfAp4EBE3CRpBTAxIv6rpNnA/cA84N3AN4FfjojDkjYBvwNsBB4FbouI\ndZKuAc6LiN+StAT4SER8bKT96mTNZig5XPRzCEAzy1u2NZuI2BsRT6Xpt4AXgWnAImB1mm01RQCR\nyh+IiLcjYicwAMxLoTU+Ip6MIjHvqVqmsq6HgUsqtZ5ukcOF3jUdM2uWjvZGS81bF1DUTKZExN70\n1WvAlDQ9DXi1tNjuVDYtTVeXH7FMRAwCB4HThtj+ckmbJW3ev39/E46oudy8Zma9omNhI+lk4GvA\nZyLiUPm7VFNpefteRKyKiLkRMXfy5Mmt3tyoOXTMrNt1JGwkvYsiaO6LiK+n4tdT01jlvs6+VL4H\nOKO0+PRUtidNV5cfsYykscAE4I3mH0l7OXTMrFt1ojeagDuBFyPiltJXa4GlaXop8EipfEnqYTYT\nmAVsSk1uhyTNT+u8smqZyroWA09ED/XxduiYWbfpRG+0i4DvAluBn6Xi36e4b/MgcCbwCvDRiDiQ\nlrkBuAoYpGh2W5fK5wJ3AycC64DrIiIknQDcS3E/6ACwJCJeHmm/cuuNVo9cLvg5BKCZtZcf6qxT\nN4dNhUPHzNrNYVOnXgibCoeOmbWLw6ZOvRQ2FQ4dM2s1h02dejFsKhw6ZtYqDps69XLYVDh0zKzZ\nHDZ16oewKcsheBw6Zt0v27HRLA85XOj9nI5Z/3DNJum3mk1ZLhf8HALQzOrjZrQ69XPYVDh0zKxe\nDps6OWze4dAxs1o5bOrksBlaDsHj0DHLlzsIWFPkcKF3RwKz7ueaTeKazbHlcsHPIQDNrOBmtDo5\nbGrn0DGzCodNnRw29XPomJnDpk4Om8bkEDwOHbP2cwcBa6scLvTuSGCWL9dsEtdsmieXC34OAWjW\n69yMVieHTfM5dMx6n8OmTg6b1nHomPWu7O/ZSBoj6WlJf50+T5K0XtKO9D6xNO9KSQOStktaUCqf\nI2lr+u42SUrl4yStSeUbJc1o9/HZO3bddPnPX53kezpmndPJDgLXAy+WPq8ANkTELGBD+oyk2cAS\n4BxgIXC7pDFpmTuAq4FZ6bUwlS8D3oyIs4FbgZtbeyhWq04HDjh0zDqhI2EjaTpwOfDlUvEiYHWa\nXg1cUSp/ICLejoidwAAwT9JUYHxEPBlFW+A9VctU1vUwcEml1mOdl0MtBxw6Zu3UqZrNl4DfA35W\nKpsSEXvT9GvAlDQ9DXi1NN/uVDYtTVeXH7FMRAwCB4HTqndC0nJJmyVt3r9/f0MHZPVz6Jj1j7aH\njaQPAfsiYstw86SaSst7LkTEqoiYGxFzJ0+e3OrN2TAcOma9b+xwX0jayggX/Ig4b5Tb/HXgw5Iu\nA04Axkv6K+B1SVMjYm9qItuX5t8DnFFafnoq25Omq8vLy+yWNBaYALwxyv21NikHTicv+pVt5xCA\nZr1ipJrNh4DfBB5Lr0+k16PpNSoRsTIipkfEDIob/09ExCeBtcDSNNtS4JE0vRZYknqYzaToCLAp\nNbkdkjQ/3Y+5smqZyroWp224j3cXyeFC75qOWfMc8zkbSU9HxAVVZU9FxPsa3rh0MfC5iPiQpNOA\nB4EzgVeAj0bEgTTfDcBVwCDwmYhYl8rnAncDJwLrgOsiIiSdANwLXAAcAJZExMsj7Yufs8lXLhf8\nHALQLDdNe6hT0jPAtRHxf9LnC4HbI+L8puxpJhw2+XPomOWnmWEzB7iL4r6HgDeBqyLiqWbsaC4c\nNt3DoWOWj6YPVyNpAkBEHGxw37LksOlOOQSPQ8f6WcNhI+l3R1owIm4Z5b5lyWHTvXIIHHDoWH+q\nNWyG7foMnNLE/TFrmcpFvtOh4y7TZsPzqM+Jaza9o9OhU+HQsX7QtFGfJf2ypA2Snk+fz5P035qx\nk2at4BEJzPJTS2+0/w38F+AvKs/bSHo+Is5tw/61jWs2vS2Hi34OAWjWbM38PZuTImJTVdng6HbL\nrDNyuNC7pmP9rJaw+aGk95DGSZO0GNg78iJm+XHzmlnn1NKMdhawCriQ4oHOncAnI2JXy/eujdyM\n1n9yueDnEIBmo9WKhzp/ATguIt5qdOdy5LDpXw4ds9FrxkOdn4yIvxru4U4/1Gm9KIfgcehYN2lG\nB4GT0vspw7zMek4OF3rf07FeNNIIAu9J7y9ExEPt2BmzHHhEArPmG6kZbStwHrClGb9dkzs3o9lw\nOh06FQ4dy1EzxkZ7jKL32cmSDpXXDUREjG9wH826gms6Zo2rpevzIxGxqE370zGu2Vg9Oh084NCx\nPDRtBIF+CBqzeuVwoXdHAusmHvU5cc3GRiuXC34OAWj9p5ljozWdpFMlPSzp7yS9KOkDkiZJWi9p\nR3qfWJp/paQBSdslLSiVz5G0NX13mySl8nGS1qTyjZJmtP8orV94GByzY6vlJwaur6WsTn8KPBYR\nvwq8F3gRWAFsiIhZwIb0GUmzgSXAOcBC4HZJY9J67gCuBmal18JUvgx4MyLOBm4Fbm5wf82OyaFj\nNrxaOgg8Vd31WdLTlZ8bqHuD0gTgGeCsKG1c0nbg4ojYK2kq8O2I+BVJKwEi4r+n+R4HbgR2Ad9K\ngYWkj6fl/3Nlnoj4vqSxwGvA5BjhYN2MZq2Qw0U/hwC03tVwM5qkj0v6X8BMSWtLr28BBxrYt5nA\nfuArkp6W9OU07tqUiKiMJv0aMCVNTwNeLS2/O5VNS9PV5UcsExGDwEHgtCGOcbmkzZI279+/v4FD\nMhtaDhd613QsByM9Z/M9ip8SOB34k1L5W8BzDW7zfcB1EbFR0p+SmswqIiIktbznQkSsohjRmrlz\n57qnhLWEn9MxGyFsIuIV4BXgA03e5m5gd0RsTJ8fpgib1yVNLTWj7Uvf7wHOKC0/PZXtSdPV5eVl\ndqdmtAnAG00+DrO6OHSsn9XSQeAtSYfS6yeSDleNKFCXiHgNeFXSr6SiS4AXgLXA0lS2FHgkTa8F\nlqQeZjMpOgJsSk1uhyTNT73QrqxaprKuxcATI92vMWsndySwfjRSMxoAEfHzEZ7TRX0RML/B7V4H\n3CfpeOBl4NMUwfegpGUUNaqPpu1vk/QgRSANAtdGxOG0nmuAu4ETgXXpBXAncK+kAYr7S0sa3F+z\npisHTicv+q7pWDuM6qHORnqj5cq90azTcqllOHSsHs0YiLOyon9b+ngcMBf4SQP7ZmZD8D0d62XH\nDBvgN0vTgxTPt3i8NLMWcehYL/LYaImb0SxXnQ6dCoeODaXWZrRaRhA4i2J4mflAAN8HPhsRLzdj\nR3PhsLFukEPwOHSsrJkDcX4VeBCYCrwbeAi4v7HdM7PRyOFC7y7TNhq11Gyei4jzqsqejYj3tnTP\n2sw1G+s2uVzwcwhA65xmNqPdTPHz0A9QNKN9DJgIfBEgIhoZJy0bDhvrVg4d66Rmhs3OEb6OiDir\n3p3LkcPGup1DxzqhmWFzQkT85Fhl3c5hY70kh+Bx6PSHZnYQ+F6NZWaWiRwu9O5IYGXDPtQp6Zco\nfhfmREkXAEpfjQdOasO+mVkD/HCo5WSkEQQWAJ+iGLr/llL5W8Dvt3CfzKyJHDqWg1ru2fy7iPha\nm/anY3zPxvpFp0OnwqHTG5o2ECdwrqRzqgsj4guj2jMz6yj/tIF1Qi0dBH4M/GN6HQY+CMxo4T6Z\nWZvkcKF3R4L+UPdAnJLGAY9HxMUt2aMOcTOa9btcLvg5BKDVrpldn6udRNFpwMx6iH+u2lqplh9P\n20oxTA3AGGAy4Ps1Zj3KvdesFWrpjfYvSh8HgdcjYrCle9UBbkYzG1qnQ6fCoZOnpjWjRcQrwKkU\nv9j5EWB2E3bus5K2SXpe0v2STpA0SdJ6STvS+8TS/CslDUjaLmlBqXyOpK3pu9skKZWPk7QmlW+U\nNKPRfTbrV25es2Y4ZthIuh64D/jF9LpP0nWj3aCkacDvAHMj4lyKprklwApgQ0TMAjakz0ianb4/\nB1gI3C5pTFrdHcDVwKz0WpjKlwFvRsTZwK3AzaPdXzMrOHSsEbV0EFgGvD8i/iAi/oDiFzuvbnC7\nYymGwRlL0eHg/wKLgNXp+9XAFWl6EfBARLwdETuBAWCepKnA+Ih4Moq2wHuqlqms62Hgkkqtx8wa\n49Cx0aglbETxfE3FYd4ZJ61uEbEH+GPgH4C9wMGI+FtgSkTsTbO9BkxJ09OAV0ur2J3KpqXp6vIj\nlkn3lw4Cpx11YNJySZslbd6/f/9oD8msLzl0rB61jCDwFWCjpG+kz1cAd452g+lezCJgJvAj4CFJ\nnyzPExEhqb4HgEYhIlYBq6DoINDq7Zn1Ivdes1ocM2wi4hZJ3wYuSkWfjoinG9jmvwZ2RsR+AElf\nBy4EXpc0NSL2piayfWn+PcAZpeWnp7I9HPm8T6W8vMzu1FQ3AXijgX02s2Nw6NhIaqnZEBFPAU81\naZv/AMyXdBLwz8AlwGaK4XCWAjel90fS/GuBr0q6BXg3RUeATRFxWNIhSfOBjcCVwP8oLbMU+D6w\nGHgi6h0qwcxGxWOv2VBGM4JAQyJiI8VN+6eArWkfVlGEzL+RtIOi9nNTmn8b8CDwAvAYcG1EVO4h\nXQN8maLTwEvAulR+J3CapAHgd0k928ysvXK40PueTh7qHhutV/mhTrPWyuWCn0MA9pJaH+p02CQO\nG7P2cOj0FodNnRw2Zu3l0OkNDps6OWzMOieH4HHojE4rf2LAzKypcrjQuyNBa7lmk7hmY5aHXC74\nOQRgN3AzWp0cNmZ5ceh0B4dNnRw2Znly6OTNYVMnh41Z/nIIHofOkdxBwMx6Tg4XenckGB3XbBLX\nbMy6Sy4X/BwCsJPcjFYnh41Zd3LodJbDpk4OG7Pu5tDpDIdNnRw2Zr0jh+Dpl9BxBwEz61s5XOjd\nkeBIrtkkrtmY9aZcLvg5BGAruBmtTg4bs97m0GkNh02dHDZm/cGh01wOmzo5bMz6Tw7B0+2h4w4C\nZmbHkMOFvl86ErQsbCTdJWmfpOdLZZMkrZe0I71PLH23UtKApO2SFpTK50jamr67TZJS+ThJa1L5\nRkkzSsssTdvYIWlpq47RzLrfrpsud+i0QStrNncDC6vKVgAbImIWsCF9RtJsYAlwTlrmdklj0jJ3\nAFcDs9Krss5lwJsRcTZwK3BzWtck4PPA+4F5wOfLoWZmNhSHTmu1LGwi4jvAgariRcDqNL0auKJU\n/kBEvB0RO4EBYJ6kqcD4iHgyiptL91QtU1nXw8AlqdazAFgfEQci4k1gPUeHnpnZkBw6rTG2zdub\nEhF70/RrwJQ0PQ14sjTf7lT20zRdXV5Z5lWAiBiUdBA4rVw+xDJHkLQcWA5w5plnju6IzKwnlQOn\nkxf9yrZzCMBGdKyDQKqpdLQrXESsioi5ETF38uTJndwVM8tYDhf6bq/ptDtsXk9NY6T3fal8D3BG\nab7pqWxPmq4uP2IZSWOBCcAbI6zLzGzU3LzWmHaHzVqg0jtsKfBIqXxJ6mE2k6IjwKbU5HZI0vx0\nP+bKqmUq61oMPJFqS48Dl0qamDoGXJrKzMwa5tAZnZbds5F0P3AxcLqk3RQ9xG4CHpS0DHgF+ChA\nRGyT9CDwAjAIXBsRh9OqrqHo2XYisC69AO4E7pU0QNERYUla1wFJfwj8IM33hYio7qhgZtaQSuB0\n+oLfLfd0PIJA4hEEzKxRnQ4eaH/oeAQBM7M2y6F2kWvzmms2iWs2ZtZMuVzwWx2AHoizTg4bM2uF\nXg8dh02dHDZm1kq9GjoOmzo5bMysXXIInmaFjjsImJllqh87Erhmk7hmY2adkEMtB0YfgG5Gq5PD\nxsw6qVtDx2FTJ4eNmeWg20LHYVMnh42Z5SaH4DlW6LiDgJlZl+uljgSu2SSu2ZhZznKo5cDRAehm\ntDo5bMysG+QWOrWGTbt/FtrMzBqQ208b1MphY2bWhcrNWZ0MnuN/6ew5tcznDgJmZl0uh44Ex+Kw\nMTPrAbn8XPVwHDZmZj0k19Bx2JiZ9aDcQqdlYSPpLkn7JD1fKvuipL+T9Jykb0g6tfTdSkkDkrZL\nWlAqnyNpa/ruNklK5eMkrUnlGyXNKC2zVNKO9FraqmM0M8tdJXQ6HTytrNncDSysKlsPnBsR5wF/\nD6wEkDQbWAKck5a5XdKYtMwdwNXArPSqrHMZ8GZEnA3cCtyc1jUJ+DzwfmAe8HlJE1twfGZmXaWT\ngdOysImI7wAHqsr+NiIG08cngelpehHwQES8HRE7gQFgnqSpwPiIeDKKp0/vAa4oLbM6TT8MXJJq\nPQuA9RFxICLepAi46tAzM+tLnarldPI5m6uANWl6GkX4VOxOZT9N09XllWVeBYiIQUkHgdPK5UMs\ncwRJy4Hl6ePb5SY/+7nTgR92eicy5PMyNJ+Xo2V9Tmp9TmY4gwf31TRfR8JG0g3AIHBfJ7ZfERGr\ngFVpnzbXMuRCv/F5GZrPy9B8Xo7mc1Joe280SZ8CPgR8It4ZmG0PcEZptumpbA/vNLWVy49YRtJY\nYALwxgjrMjOzDmlr2EhaCPwe8OGI+KfSV2uBJamH2UyKjgCbImIvcEjS/HQ/5krgkdIylZ5mi4En\nUng9DlwqaWLqGHBpKjMzsw5pWTOapPuBi4HTJe2m6CG2EhgHrE89mJ+MiN+KiG2SHgReoGheuzYi\nDqdVXUPRs+1EYF16AdwJ3CtpgKIjwhKAiDgg6Q+BH6T5vhARR3RUGMaqBg63l/m8DM3nZWg+L0fz\nOcE/MWDWMEnfi4gLm7zOGcCFEfHVYb7/InAZ8CjwEvBPEXHPCOu7EfhxRPzxENv564g4txn7bTYc\nj/ps1qBmB00yA/gPwJBhQ9GLclKpBcAsax6uxqxBkn6c3i+W9G1JD6eRMu4rjXixS9IfpdEwNkk6\nO5XfLWlx9bqAm4B/KekZSZ+t2t5a4GRgi6SPSbpR0ufSd++R9JikLZK+K+lXh9jfOZKelfQscG0L\nTonZURw2Zs11AfAZYDZwFvDrpe8ORsSvAX8GfOkY61kBfDcizo+IW8tfRMSHgX9O362pWm4VcF1E\nzAE+B9w+xLq/kuZ5b60HZdYoh41Zc22KiN0R8TPgGYrmsIr7S+8faPaGJZ0MXAg8JOkZ4C+AqVXz\nnAqcmkb4ALi32fthNhTfszFrrrdL04c58t9YDDE9SPpPn6TjgOMb2PZxwI8i4vwG1mHWEq7ZmLXP\nx0rv30/Tu4DKcCEfBt6Vpt8CTqln5RFxCNgp6d8DqPDeqnl+BPxI0kWp6BP1bMNstBw2Zu0zUdJz\nwPVA5ab/XwL/Kt2s/wDwj6n8OeBwupH/2aNXNaxPAMvS+rZRDFhb7dPAn6emNo3iOMzq5udszNpA\n0i5gbkRkOyCjWSu5ZmNmZi3nmo2ZmbWcazZmZtZyDhszM2s5h42ZmbWcw8bMzFrOYWNmZi33/wEG\nGkWXLUxPggAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x10ed58eb8>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.spy(A)\n",
    "plt.xlabel('input field')\n",
    "plt.ylabel('output field')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input grid size is `Nlat*Nlon=400*600=240000`. The size of additional dimension is 500 (think of `Ntime*Nlev`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(500, 240000)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = np.random.rand(500, A.shape[1])\n",
    "data.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output grid size is `Nlat*Nlon=300*400=120000`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(500, 120000)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# placeholder for numba/cython\n",
    "data_out = np.zeros([data.shape[0], A.shape[0]])\n",
    "data_out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.96, 0.48)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# data size in GB\n",
    "data.nbytes/1e9, data_out.nbytes/1e9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test I/O efficiency"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In real workflow, users typically read input data from disk and then dump regridding results to disk. Making the regridding calculation too much faster than I/O will have little benefit on the overall efficiency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "xr.DataArray(data).to_netcdf('data_in.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "676 ms ± 72.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit xr.open_dataarray('data_in.nc').load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "263 ms ± 14.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit xr.DataArray(data_out).to_netcdf('data_out.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# scipy.sparse dot\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Common memory order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# convert COO format to other formats\n",
    "A_csc = A.tocsc()\n",
    "A_csr = A.tocsr()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To perform `A.dot()` on data, we need to pass the transpose so their shapes agree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(500, 120000)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_scipy = A.dot(data.T).T  # record results for later reference\n",
    "out_scipy.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "COO is faster than CSC/CSR, contradicts the common belief that CSC/CSR are more efficient for matrix multiplication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "692 ms ± 45.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "1.13 s ± 67.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "1.11 s ± 92 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit A.dot(data.T).T\n",
    "%timeit A_csc.dot(data.T).T\n",
    "%timeit A_csr.dot(data.T).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Memory order matters.** The input argument `data.T` is Fortran-ordered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "  C_CONTIGUOUS : False\n",
       "  F_CONTIGUOUS : True\n",
       "  OWNDATA : False\n",
       "  WRITEABLE : True\n",
       "  ALIGNED : True\n",
       "  UPDATEIFCOPY : False"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.T.flags"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reversed memory order"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If `data.T` is C-ordered, then CSC/CSR will be much faster than COO."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data_T = data.T.copy()  # make a C-ordered copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "  C_CONTIGUOUS : True\n",
       "  F_CONTIGUOUS : False\n",
       "  OWNDATA : True\n",
       "  WRITEABLE : True\n",
       "  ALIGNED : True\n",
       "  UPDATEIFCOPY : False"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_T.flags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.5 s ± 129 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "458 ms ± 54.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "400 ms ± 17.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit A.dot(data_T).T\n",
    "%timeit A_csc.dot(data_T).T\n",
    "%timeit A_csr.dot(data_T).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scipy summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For most NetCDF data, the grid dimensions (lat and lon) are the fastest-changing dimensions. For example, the shape could be `(Ntime, Nlev, Nlat, Nlon)` (C-ordered). We flatten it to `(Ntime*Nlev, Nlat*Nlon)` (still C-ordered), then pass the tranpose `(Nlat*Nlon, Ntime*Nlev)` (Fortran-ordered) for `A.dot()` operation. Thus COO format should be used for most NetCDF data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The internal code for [`scipy.sparse.coo_matrix.dot`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.coo_matrix.dot.html) is a **just a plain algorithm written in C++**, without any specialize library\n",
    ": [scipy/coo.h](https://github.com/scipy/scipy/blob/master/scipy/sparse/sparsetools/coo.h#L144)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus we replicate the same algorithm with Numba/Cython, and try to parallelize it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Serial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@jit(nopython=True, nogil=True)\n",
    "def sparse_dot(data, col, row, S):\n",
    "    # A.shape[0] = 120000 is hard-coded here\n",
    "    data_out = np.zeros((data.shape[0], 120000), dtype=np.float64)\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",
    "    return data_out\n",
    "\n",
    "data_out = sparse_dot(data, col, row, S); # compile jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# exactly the same as Scipy result\n",
    "np.array_equal(data_out, out_scipy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba have almost the same performance as Scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "630 ms ± 11.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit sparse_dot(data, col, row, S)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@jit(nopython=True, parallel=True, nogil=True)\n",
    "def sparse_dot_parallel(data, col, row, S):\n",
    "    # A.shape[0] = 120000 is hard-coded here\n",
    "    data_out = np.zeros((data.shape[0], 120000), dtype=np.float64)\n",
    "    for j in prange(data.shape[0]):\n",
    "        data_out[j, :] = 0.0\n",
    "        for i in range(S.size):\n",
    "            data_out[j, row[i]] += data[j, col[i]]*S[i]\n",
    "            \n",
    "    return data_out\n",
    "\n",
    "sparse_dot_parallel(data, col, row, S); # compile jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array_equal(data_out, out_scipy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "310 ms ± 16 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit sparse_dot_parallel(data, col, row, S)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parallelization on my dual-core Mac gives 2x speed-up. Regridding now takes almost the same time as I/O."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%load_ext Cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "# adding --compile-args=-O3 doesn't seem to help\n",
    "cimport cython\n",
    "import numpy as np # need to import numpy again in cython context\n",
    "\n",
    "@cython.wraparound(False) # turning off wraparound gives 20% speed-up\n",
    "@cython.boundscheck(False)\n",
    "cpdef cython_dot(double[:, :] data, int[:] col, int[:] row, double[:] S):\n",
    "    cdef double[:, :] data_out = np.zeros((data.shape[0], 120000))\n",
    "    cdef int i, j \n",
    "    for j in range(data.shape[0]):\n",
    "        for i in range(S.shape[0]):\n",
    "            data_out[j, row[i]] += data[j, col[i]]*S[i]\n",
    "            \n",
    "    return data_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# exactly the same as Scipy result\n",
    "data_out = cython_dot(data, col, row, S)\n",
    "np.array_equal(data_out, out_scipy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cython's performance is also similar to Scipy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "648 ms ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit cython_dot(data, col, row, S)"
   ]
  }
 ],
 "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
}