docs/source/examples/dask/efficient_dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using `dask` with `climpred`\n",
    "\n",
    "This demo demonstrates `climpred`'s capabilities with `dask` https://docs.dask.org/en/latest/array.html. This enables enables out-of-memory and parallel computation for large datasets with `climpred`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import dask\n",
    "import climpred\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Load a `Client` to use `dask.distributed`: https://stackoverflow.com/questions/51099685/best-practices-in-setting-number-of-dask-workers\n",
    "- (Optionally) Use the `dask` dashboard to visualize performance: https://github.com/dask/dask-labextension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of CPUs: 4, number of threads: 2, number of workers: 2, processes: False\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Client</h3>\n",
       "<ul style=\"text-align: left; list-style: none; margin: 0; padding: 0;\">\n",
       "  <li><b>Scheduler: </b>inproc://136.172.147.245/47102/1</li>\n",
       "  <li><b>Dashboard: </b><a href='http://136.172.147.245:8787/status' target='_blank'>http://136.172.147.245:8787/status</a></li>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Cluster</h3>\n",
       "<ul style=\"text-align: left; list-style:none; margin: 0; padding: 0;\">\n",
       "  <li><b>Workers: </b>2</li>\n",
       "  <li><b>Cores: </b>4</li>\n",
       "  <li><b>Memory: </b>34.36 GB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'inproc://136.172.147.245/47102/1' processes=2 threads=4, memory=34.36 GB>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dask.distributed import Client\n",
    "import multiprocessing\n",
    "ncpu = multiprocessing.cpu_count()\n",
    "processes = False\n",
    "nworker = 2\n",
    "threads = ncpu // nworker\n",
    "print(\n",
    "    f\"Number of CPUs: {ncpu}, number of threads: {threads}, number of workers: {nworker}, processes: {processes}\",\n",
    ")\n",
    "client = Client(\n",
    "    processes=processes,\n",
    "    threads_per_worker=threads,\n",
    "    n_workers=nworker,\n",
    "    memory_limit=\"64GB\",\n",
    ")\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generic\n",
    "ny, nx = 256, 220\n",
    "nl, ni, nm = 20, 12, 10\n",
    "init = xr.DataArray(np.random.random((nl, ni, nm, ny, nx)), dims=('lead', 'init', 'member', 'y', 'x'))\n",
    "init.name='var'\n",
    "init['init'] = np.arange(3000, 3300, 300 // ni)\n",
    "init['lead'] = np.arange(1,1+init.lead.size)\n",
    "control = xr.DataArray(np.random.random((300, ny, nx)),dims=('time', 'y', 'x'))\n",
    "control.name='var'\n",
    "control['time'] = np.arange(3000, 3300)\n",
    "\n",
    "pm = climpred.PerfectModelEnsemble(init).add_control(control)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## compute skill with `compute_perfect_model`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "kw = {'comparison':'m2e', 'metric':'rmse', 'dim':['init', 'member']}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### compute skill without `dask`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 12 s, sys: 7.08 s, total: 19.1 s\n",
      "Wall time: 19.1 s\n"
     ]
    }
   ],
   "source": [
    "%time s = pm.verify(**kw)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 2 core Mac Book Pro 2018: CPU times: user 11.5 s, sys: 6.88 s, total: 18.4 s Wall time: 19.6 s\n",
    "- 24 core mistral node: CPU times: user 9.22 s, sys: 10.3 s, total: 19.6 s Wall time: 19.5 s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `verify()` skill with `dask`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to use `dask` efficient, we need to chunk the data appropriately. Processing chunks of data lazily with `dask` creates a tiny overhead per dask, therefore chunking mostly makes sense when applying it to large data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td>\n",
       "<table>\n",
       "  <thead>\n",
       "    <tr><td> </td><th> Array </th><th> Chunk </th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th> Bytes </th><td> 1.08 GB </td> <td> 540.67 MB </td></tr>\n",
       "    <tr><th> Shape </th><td> (20, 12, 10, 256, 220) </td> <td> (20, 12, 10, 128, 220) </td></tr>\n",
       "    <tr><th> Count </th><td> 2 Tasks </td><td> 2 Chunks </td></tr>\n",
       "    <tr><th> Type </th><td> float64 </td><td> numpy.ndarray </td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</td>\n",
       "<td>\n",
       "<svg width=\"380\" height=\"189\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"34\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"37\" x2=\"34\" y2=\"37\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"0\" y2=\"37\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"34\" y1=\"0\" x2=\"34\" y2=\"37\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.000000,0.000000 34.105574,0.000000 34.105574,37.188039 0.000000,37.188039\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"17.052787\" y=\"57.188039\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >12</text>\n",
       "  <text x=\"54.105574\" y=\"18.594019\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(0,54.105574,18.594019)\">20</text>\n",
       "\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"104\" y1=\"0\" x2=\"123\" y2=\"19\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"104\" y1=\"60\" x2=\"123\" y2=\"79\" />\n",
       "  <line x1=\"104\" y1=\"120\" x2=\"123\" y2=\"139\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"104\" y1=\"0\" x2=\"104\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"123\" y1=\"19\" x2=\"123\" y2=\"139\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"104.000000,0.000000 123.405241,19.405241 123.405241,139.405241 104.000000,120.000000\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"104\" y1=\"0\" x2=\"207\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"123\" y1=\"19\" x2=\"226\" y2=\"19\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"104\" y1=\"0\" x2=\"123\" y2=\"19\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"207\" y1=\"0\" x2=\"226\" y2=\"19\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"104.000000,0.000000 207.125000,0.000000 226.530241,19.405241 123.405241,19.405241\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"123\" y1=\"19\" x2=\"226\" y2=\"19\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"123\" y1=\"79\" x2=\"226\" y2=\"79\" />\n",
       "  <line x1=\"123\" y1=\"139\" x2=\"226\" y2=\"139\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"123\" y1=\"19\" x2=\"123\" y2=\"139\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"226\" y1=\"19\" x2=\"226\" y2=\"139\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"123.405241,19.405241 226.530241,19.405241 226.530241,139.405241 123.405241,139.405241\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"174.967741\" y=\"159.405241\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >220</text>\n",
       "  <text x=\"246.530241\" y=\"79.405241\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,246.530241,79.405241)\">256</text>\n",
       "  <text x=\"103.702620\" y=\"149.702620\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,103.702620,149.702620)\">10</text>\n",
       "</svg>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<xarray-var, shape=(20, 12, 10, 256, 220), dtype=float64, chunksize=(20, 12, 10, 128, 220), chunktype=numpy.ndarray>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chunked_dim = 'y'\n",
    "chunks = {chunked_dim:init[chunked_dim].size // nworker}\n",
    "pm_chunked = pm.chunk(chunks)\n",
    "# if memory allows\n",
    "# pm_chunked = pm_chunked.persist()\n",
    "pm_chunked.get_initialized()['var'].data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 35.8 s, sys: 18.5 s, total: 54.3 s\n",
      "Wall time: 19.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "s_chunked = pm_chunked.verify(**kw)\n",
    "assert dask.is_dask_collection(s_chunked)\n",
    "s_chunked = s_chunked.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 2 core Mac Book Pro 2018: CPU times: user 2min 35s, sys: 1min 4s, total: 3min 40s Wall time: 2min 10s\n",
    "- 24 core mistral node: CPU times: user 26.2 s, sys: 1min 37s, total: 2min 3s Wall time: 5.38 s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    xr.testing.assert_allclose(s,s_chunked,atol=1e-6)\n",
    "except AssertionError:\n",
    "    for v in s.data_vars:\n",
    "        (s-s_chunked)[v].plot(robust=True, col='lead')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- The results `s` and `s_chunked` are identical as requested.\n",
    "- Chunking reduces Wall time from 20s to 5s on supercomputer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `bootstrap()` skill\n",
    "\n",
    "This speedup translates into `bootstrap_perfect_model`, where bootstrapped resamplings of intializialized, uninitialized and persistence skill are computed and then translated into p values and confidence intervals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "kwp = kw.copy()\n",
    "kwp['iterations'] = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### without `dask`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time s_p = pm.bootstrap(**kwp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 2 core Mac Book Pro 2018: CPU times: user 2min 3s, sys: 1min 22s, total: 3min 26s Wall time: 3min 43s\n",
    "- 24 core mistral node: CPU times: user 1min 51s, sys: 1min 54s, total: 3min 45s Wall time: 3min 25s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### bootstrap skill with `dask`\n",
    "\n",
    "When `ds` is chunked, `bootstrap_perfect_model` performs all skill calculations on resampled inputs in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time s_p_chunked = pm_chunked.bootstrap(**kwp).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 2 core Mac Book Pro 2018: CPU times: user 2min 35s, sys: 1min 4s, total: 3min 40s Wall time: 2min 10s\n",
    "- 24 core mistral node: CPU times: user 2min 55s, sys: 8min 8s, total: 11min 3s Wall time: 1min 53s"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:climpred-dev] *",
   "language": "python",
   "name": "conda-env-climpred-dev-py"
  },
  "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.11"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": false,
   "autoclose": true,
   "autocomplete": true,
   "bibliofile": "large.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": false,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}