diff --git a/_toc.yml b/_toc.yml index cfe7f887..4a4f1177 100644 --- a/_toc.yml +++ b/_toc.yml @@ -53,6 +53,7 @@ parts: - file: intermediate/hvplot - file: intermediate/datastructures-intermediate.ipynb - file: intermediate/BiologyDataset.ipynb + - file: intermediate/hierarchical_zarr_store.ipynb - file: intermediate/remote_data/index sections: - file: intermediate/remote_data/cmip6-cloud.ipynb diff --git a/intermediate/hierarchical_zarr_store.ipynb b/intermediate/hierarchical_zarr_store.ipynb new file mode 100644 index 00000000..44001fa4 --- /dev/null +++ b/intermediate/hierarchical_zarr_store.ipynb @@ -0,0 +1,331 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Zarr Stores with Xarray\n", + "\n", + "## Learning Objectives:\n", + "- Learn about the Zarr data format and xarray's \"zarr\" backend engine\n", + "- Learn how to read a Zarr store with a hierarchical structure with `xr.DataTree`\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. Xarray’s Zarr backend allows xarray to leverage these capabilities, including the ability to store and analyze datasets far too large fit onto disk (particularly in combination with [dask](https://docs.xarray.dev/en/latest/user-guide/dask.html#dask)).\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": [ + "## Reading a zarr store\n", + "\n", + "With xarray's \"zarr\" backend, we can read data from cloud storage buckets. The Zarr store we will be using for this tutorial has groups: \"observed\" and \"reanalysis\", each containing a \"precipitation\" data variable derived from *[GPM_3IMERGHH_07](https://disc.gsfc.nasa.gov/datasets/GPM_3IMERGHH_07/summary)* and *[M2T1NXFLX_5.12.4](https://disc.gsfc.nasa.gov/datasets/M2T1NXFLX_5.12.4/summary)* products, respectively.\n", + "\n", + "Let's read in our zarr store as an `xr.DataTree`. " + ] + }, + { + "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": [ + ":::{note} We selected `\"zarr\"` backend engine, which tells xarray to load and decode a dataset from a Zarr store. The `chunks={}` parameter is used to load the data into a dask array. And `consolidated=True` enables zarr’s consolidated metadata capability. This lets us read all of the metadata from a single file which can improve performance.\n", + ":::" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "## Variable selection\n", + "With our `DataTree` object, we can select variables from our Zarr store with either dictionary and or attribute like syntax." + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "### Dictionary-like interface" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "precip_dt[\"observed/precipitation\"]" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "### Attribute-like access" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "precip_dt.observed.precipitation" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "### Dictionary and Attribute like access" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "precip_dt.observed[\"precipitation\"]" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + ":::{note} All of these variable selection options return the same \"precipitation\" `xr.DataArray` object, as a chunked `dask.Array`, from the \"observed\" group of our Zarr store.\n", + ":::" + ] + }, + { + "cell_type": "markdown", + "id": "12", + "metadata": {}, + "source": [ + "## Time slicing\n", + "\n", + "We can index and subset by time on our `xr.DataTree` object. Each time slice in our Zarr store represents one hour of data with a total of 10 hours of data.\n", + "\n", + "Let's explore the different ways we can get the first 5 hours of data. " + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "### Label-based indexing" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "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": "15", + "metadata": {}, + "outputs": [], + "source": [ + "time_index = precip_dt.time[0:5]\n", + "precip_dt.sel(time=time_index)" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "### Datetime indexing\n", + "We can also subset by time with a `datetime` string." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "precip_dt.sel(time=slice(\"2021-08-29T07:30:00\", \"2021-08-29T16:30:00\"))" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "### Positional Indexing\n", + "Or by the index of our time dimension `.isel(time=slice())`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "precip_dt.isel(time=slice(0, 5))" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "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": "21", + "metadata": {}, + "outputs": [], + "source": [ + "precip_dt.observed[\"precipitation\"].data.chunks" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "### Selecting by chunks\n", + "\n", + "Since we loaded our data as a `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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "precip_dt.observed[\"precipitation\"].data.blocks[0, 0, 0].compute()" + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + ":::{note}\n", + "We added `.data` to our `xr.DataArray` to access the `dask.Array`. The `.blocks[]` method allows you to index by chunk and `.compute()` returns the `np.ndarray`. \n", + ":::" + ] + }, + { + "cell_type": "markdown", + "id": "25", + "metadata": {}, + "source": [ + "## Exercise" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "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} Hint\n", + ":class: dropdown\n", + "This is how you could calculate mean \"precipitation\"\n", + "\n", + "```python\n", + "precip_dt.reanalysis['precipitation'].mean(dim='time')\n", + "```\n", + ":::\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 +}