tutorial hierarchical zarr stores

intermediate/hierarchical_zarr_store.ipynb

@@ -0,0 +1,251 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# Zarr Stores with Xarray\n",
+    "\n",
+    "## Learning Objectives:\n",
+    "\n",
+    "- Learn how to read a Zarr store with a group hierarchical structure with `xr.DataTree` and `xarray`'s `\"zarr\"` engine\n",
+    "- Learn how to select `dask` arrays from a Zarr store\n",
+    "- Explore how to use Zarr stores with `xarray` for computations and visualizations\n",
+    "\n",
+    "## What is Zarr?\n",
+    "\n",
+    "The Zarr data format is an open, community-maintained format designed for efficient, scalable storage of large N-dimensional arrays. It stores data as compressed and chunked arrays in a format well-suited to parallel processing and cloud-native workflows. \n",
+    "\n",
+    "### Zarr Data Organization:\n",
+    "- **Arrays**: N-dimensional arrays that can be chunked and compressed.\n",
+    "- **Groups**: A container for organizing multiple arrays and other groups with a hierarchical structure.\n",
+    "- **Metadata**: JSON-like metadata describing the arrays and groups, including information about data types, dimensions, chunking, compression, and user-defined key-value fields. \n",
+    "- **Dimensions and Shape**: Arrays can have any number of dimensions, and their shape is defined by the number of elements in each dimension.\n",
+    "- **Coordinates & Indexing**: Zarr supports coordinate arrays for each dimension, allowing for efficient indexing and slicing.\n",
+    "\n",
+    "The diagram below from [the Zarr v3 specification](https://wiki.earthdata.nasa.gov/display/ESO/Zarr+Format) showing the structure of a Zarr store:\n",
+    "\n",
+    "![ZarrSpec](https://zarr-specs.readthedocs.io/en/latest/_images/terminology-hierarchy.excalidraw.png)\n",
+    "\n",
+    "\n",
+    "NetCDF and Zarr share similar terminology and functionality, but the key difference is that NetCDF is a single file, while Zarr is a directory-based “store” composed of many chunked files, making it better suited for distributed and cloud-based workflows."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
+   "source": [
+    "## Opening a dataset\n",
+    "\n",
+    "Let's open up a precipitation zarr store, with `xr.open_datatree(engine='zarr')`. This dataset was derived from \"GPM_3IMERGHH_07\" and \"M2T1NXFLX_5.12.4\" products and has a group hierarchical structure.\n",
+    "**NOTE** We selected `consolidated=True`, Zarr supports consolidated metadata, which allows you to store all metadata in a single file and can improve performance when reading metadata."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import xarray as xr\n",
+    "\n",
+    "precipitation_store = \"https://pub-45a1d62ac8d94c4c89f4dc63681a98ed.r2.dev/precipitation.zarr\"\n",
+    "\n",
+    "precip_dt = xr.open_datatree(precipitation_store, engine=\"zarr\", chunks={}, consolidated=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "## Variable selection\n",
+    "As with on disk storage methods, data from zarr stores in the cloud can be accessed with either dict-like syntax or method based syntax with the `\"zarr\"` engine."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "precip_dt.observed[\"precipitation\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "This returns an `xr.DataArray` object with the underly data as a chunked `dask.Array`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6",
+   "metadata": {},
+   "source": [
+    "## Time slicing\n",
+    "\n",
+    "Like a dataset stored on disk we can do time slicing on our zarr store. Each time slice is one hour data with a total of 10 hours of data. \n",
+    "\n",
+    "Let's try getting the first 5 hours of data with `.sel(time=)`. \n",
+    "\n",
+    "Since the time slices are ordered we can get a subset of the array of our time coordinate and pass it to the `.sel` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "time_index = precip_dt.time[0:5]\n",
+    "precip_dt.sel(time=time_index)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8",
+   "metadata": {},
+   "source": [
+    "We can also subset by time with a `datetime` string."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "precip_dt.sel(time=slice(\"2021-08-29T07:30:00\", \"2021-08-29T16:30:00\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "10",
+   "metadata": {},
+   "source": [
+    "Or by the index of our time dimension `.isel(time=slice())`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "precip_dt.isel(time=slice(0, 5))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "12",
+   "metadata": {
+    "vscode": {
+     "languageId": "plaintext"
+    }
+   },
+   "source": [
+    "## Chunking\n",
+    "Chunking is the process of dividing arrays into smaller pieces, which allows for parallel processing and efficient storage.\n",
+    "\n",
+    "To examine the chunks in our Zarr store, with `xarray` you can use the `chunks` attribute on the `xr.DataArray` object."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "precip_dt.observed[\"precipitation\"].data.chunks"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "14",
+   "metadata": {},
+   "source": [
+    "### Selecting by chunks\n",
+    "\n",
+    "Since our underlying data arrays are `dask.Array`, we can access data from each chunked array in our Zarr store. \n",
+    "\n",
+    "Let's get the first chunk of the \"observed/precipitation\" variable in our zarr store.\n",
+    "\n",
+    "We add `.data` to our `xr.DataArray` to get access the `dask.Array`. The `.blocks[]` method allows you to index by chunk and `.compute()` returns the `np.ndarray`. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "precip_dt.observed[\"precipitation\"].data.blocks[0, 0, 0].compute()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "16",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "precip_dt.observed[\"precipitation\"].mean(dim=\"time\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "17",
+   "metadata": {},
+   "source": [
+    "## Exercise"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "18",
+   "metadata": {},
+   "source": [
+    "::::{admonition} Exercise\n",
+    ":class: tip\n",
+    "\n",
+    "Can you calculate and plot the mean precipitation starting at 09:55 for the reanalysis group in this zarr store?\n",
+    "\n",
+    ":::{admonition} Solution\n",
+    ":class: dropdown\n",
+    "\n",
+    "```python\n",
+    "precip_dt.reanalysis['precipitation'].sel(time=slice('2021-08-29T09:55:00', '2021-08-29T16:30:00')).mean(dim='time').plot()\n",
+    "```\n",
+    ":::\n",
+    "::::"
+   ]
+  }
+ ],
+ "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
+}