initial draft of the datatree computation tutorial

Author: keewis

intermediate/hierarchical_computation.ipynb

@@ -0,0 +1,219 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# Hierarchical computations\n",
+    "\n",
+    "In this lesson, we extend what we learned about [basic computation](#basic-computation) to hierarchical datasets. By the end of the lesson, we will be able to:\n",
+    "\n",
+    "- Apply basic arithmetic and label-aware reductions to xarray DataTree objects\n",
+    "- Apply arbitrary functions across all nodes across a tree"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import xarray as xr\n",
+    "import numpy as np\n",
+    "\n",
+    "xr.set_options(keep_attrs=True, display_expand_attrs=False, display_expand_data=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2",
+   "metadata": {},
+   "source": [
+    "## Example dataset\n",
+    "\n",
+    "First we load the NMC reanalysis air temperature dataset and arrange it to form a hierarchy of temporal resolutions:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds = xr.tutorial.open_dataset(\"air_temperature\")\n",
+    "\n",
+    "ds_daily = (\n",
+    "    ds.resample(time=\"D\").mean(\"time\").assign(mask=lambda ds: ds[\"air\"].isel(time=0) >= 293.15)\n",
+    ")\n",
+    "ds_weekly = (\n",
+    "    ds.resample(time=\"W\").mean(\"time\").assign(mask=lambda ds: ds[\"air\"].isel(time=0) >= 293.15)\n",
+    ")\n",
+    "ds_monthly = (\n",
+    "    ds.resample(time=\"ME\").mean(\"time\").assign(mask=lambda ds: ds[\"air\"].isel(time=0) >= 293.15)\n",
+    ")\n",
+    "\n",
+    "tree = xr.DataTree.from_dict({\"daily\": ds_daily, \"weekly\": ds_weekly, \"monthly\": ds_monthly})\n",
+    "tree"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4",
+   "metadata": {},
+   "source": [
+    "## Arithmetic\n",
+    "\n",
+    "As an extension to `Dataset`, `DataTree` objects automatically apply arithmetic to all variables within all nodes:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tree - 273.15"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6",
+   "metadata": {},
+   "source": [
+    "## Reductions\n",
+    "\n",
+    "In a similar way, we can reduce all nodes in the datatree at once:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tree.mean(dim=[\"lat\", \"lon\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8",
+   "metadata": {},
+   "source": [
+    "## Applying functions designed for `Dataset` with `map_over_datasets`\n",
+    "\n",
+    "What if we wanted to convert the data to log-space? For a `Dataset` or `DataArray`, we could just use {py:func}`xarray.ufuncs.log`, but that does not support `DataTree` objects, yet:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "xr.ufuncs.log(tree)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "10",
+   "metadata": {},
+   "source": [
+    "Note how the result is a empty `Dataset`?\n",
+    "\n",
+    "To map a function to all nodes, we can use {py:func}`xarray.map_over_datasets` and {py:meth}`xarray.DataTree.map_over_datasets`: "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tree.map_over_datasets(xr.ufuncs.log)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "12",
+   "metadata": {},
+   "source": [
+    "We can also use a custom function to perform more complex operations, like subtracting a group mean:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def demean(ds):\n",
+    "    return ds.groupby(\"time.day\") - ds.groupby(\"time.day\").mean()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "14",
+   "metadata": {},
+   "source": [
+    "Applying that to the dataset raises an error, though:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tree.map_over_datasets(demean)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "16",
+   "metadata": {},
+   "source": [
+    "The reason for this error is that the root node does not have any variables, and thus in particular no `\"time\"` coordinate. To avoid the error, we have to skip computing the function for that node:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def demean(ds):\n",
+    "    if \"time\" not in ds.coords:\n",
+    "        return ds\n",
+    "    return ds.groupby(\"time.day\") - ds.groupby(\"time.day\").mean()\n",
+    "\n",
+    "\n",
+    "tree.map_over_datasets(demean)"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}