Xarray and Dask Example

binder/environment.yml

@@ -17,6 +17,7 @@ dependencies:
   - nomkl
   - h5py
   - py-xgboost
+  - xarray
   - pip:
     - graphviz
     - dask_xgboost

xarray.ipynb

@@ -0,0 +1,361 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Xarray with Dask Arrays\n",
+    "\n",
+    "<img src=\"images/dataset-diagram-logo.png\" \n",
+    "     align=\"right\"\n",
+    "     width=\"66%\"\n",
+    "     alt=\"Xarray Dataset\">\n",
+    "     \n",
+    "**Xarray** is an open source project and Python package that extends the labeled data functionality of pandas N-dimensional array-like datasets. It shares a similar API to Numpy and Pandas and supports Dask arrays under the hood."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Start Dask Client for Dashboard\n",
+    "\n",
+    "Starting the Dask Client is optional.  It will provide a dashboard which \n",
+    "is useful to gain insight on the computation.  \n",
+    "\n",
+    "The link to the dashboard will become visible when you create the client below.  We recommend having it open on one side of your screen while using your notebook on the other side.  This can take some effort to arrange your windows, but seeing them both at the same is very useful when learning."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dask.distributed import Client, progress\n",
+    "client = Client(n_workers=2, threads_per_worker=2, memory_limit='1GB')\n",
+    "client"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Open a sample dataset\n",
+    "\n",
+    "We will use some of xarray's tutorial data for this example. All we need to do is specify the chunk shape and xarray will automatically create Dask arrays for the variables in the dataset."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import xarray as xr\n",
+    "\n",
+    "ds = xr.tutorial.load_dataset('air_temperature',\n",
+    "                              chunks={'lat': 25, 'lon': 25, 'time': 1000})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "By specifying the chunk sizes by dimension name, we have configured the data variables in our `Dataset` to be returned backed by Dask arrays. In xarray, `Datasets` are dict-like container of labeled arrays akin to the `pandas.DataFrame`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Quickly inspecting the Dataset above, we'll note that this Dataset has three dimensions akin to axis in NumPy (lat, lon, and time), three coordinate variables akin to Pandas.Index objects (also named lat, lon, and time), and one data variable (air). Xarray also holds Dataset specific metadata in as Attributes."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da = ds['air']\n",
+    "da"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Each data variable in xarray is called a `DataArray`. These are the fundemental labeled array object in xarray. Much like the Dataset, DataArray's still have dimensions and coordinate that support many of its array opperations."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da.data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Accessing the underlying array of data is done via the `data` property. Here we can see that we have a Dask array."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use Standard Xarray Operations\n",
+    "\n",
+    "Most common xarray operations operate identically on xarray objects backed by dask. The "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da2 = da.groupby('time.month').mean('time')\n",
+    "da3 = da - da2\n",
+    "da3"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Call `.compute()` when you want your result as a xarray.DataArray with data stored as NumPy arrays.\n",
+    "\n",
+    "If you started `Client()` above then you may want to watch the status page during computation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "computed_da = da3.compute()\n",
+    "type(computed_da.data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "computed_da"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Persist data in memory\n",
+    "\n",
+    "If you have the available RAM for your dataset then you can persist data in memory.  \n",
+    "\n",
+    "This allows future computations to be much faster."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da = da.persist()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Time Series Operations\n",
+    "\n",
+    "Because we have a datetime index time-series operations work efficiently"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da.resample(time='1w').mean('time').std('time')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da.resample(time='1w').mean('time').std('time').compute().plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da_smooth = da.rolling(time=30).mean()\n",
+    "da_smooth"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Since xarray stores each of its coordinate variables in memory, slicing by label is trivial and entirely lazy."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%time da.sel(time='2013-01-01T18:00:00')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%time da.sel(time='2013-01-01T18:00:00').compute()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Custom workflows and automatic parallelization"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Almost all of xarray’s built-in operations work on dask arrays. If you want to use a function that isn’t wrapped by xarray, one option is to extract dask arrays from xarray objects (.data) and use dask directly.\n",
+    "\n",
+    "Another option is to use xarray’s apply_ufunc(), which can automate embarrassingly parallel “map” type operations where a functions written for processing NumPy arrays should be repeatedly applied to xarray objects containing dask arrays. It works similarly to dask.array.map_blocks() and dask.array.atop(), but without requiring an intermediate layer of abstraction.\n",
+    "\n",
+    "Here we show an example using NumPy operations and a fast function from bottleneck, which we use to calculate Spearman’s rank-correlation coefficient:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import xarray as xr\n",
+    "import bottleneck\n",
+    "\n",
+    "def covariance_gufunc(x, y):\n",
+    "    return ((x - x.mean(axis=-1, keepdims=True))\n",
+    "            * (y - y.mean(axis=-1, keepdims=True))).mean(axis=-1)\n",
+    "\n",
+    "def pearson_correlation_gufunc(x, y):\n",
+    "    return covariance_gufunc(x, y) / (x.std(axis=-1) * y.std(axis=-1))\n",
+    "\n",
+    "def spearman_correlation_gufunc(x, y):\n",
+    "    x_ranks = bottleneck.rankdata(x, axis=-1)\n",
+    "    y_ranks = bottleneck.rankdata(y, axis=-1)\n",
+    "    return pearson_correlation_gufunc(x_ranks, y_ranks)\n",
+    "\n",
+    "def spearman_correlation(x, y, dim):\n",
+    "    return xr.apply_ufunc(\n",
+    "        spearman_correlation_gufunc, x, y,\n",
+    "        input_core_dims=[[dim], [dim]],\n",
+    "        dask='parallelized',\n",
+    "        output_dtypes=[float])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In the examples above, we were working with an some air temperature data. For this example, we'll calculate the spearman correlation using the raw air temperature data with the smoothed version that we also created (`da_smooth`). For this, we'll also have to rechunk the data ahead of time."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "corr = spearman_correlation(da.chunk({'time': -1}),\n",
+    "                            da_smooth.chunk({'time': -1}),\n",
+    "                            'time')\n",
+    "corr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "corr.plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "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.6.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}