{ "cells": [ { "cell_type": "markdown", "id": "990cfa4c-2117-4435-9806-ff9048890398", "metadata": { "tags": [] }, "source": [ "\"Dask\n", "\"Xarray\"\n",\n", "\n", "# Parallelizing Xarray with Dask\n", "\n", "### In this tutorial, you learn:\n", "\n", "* Using Dask with Xarray\n", "* Read/write netCDF files with Dask\n", "* Dask backed Xarray objects and operations\n", "* Extract Dask arrays from Xarray objects and use Dask array directly.\n", "* Xarray built-in operations can transparently use dask\n", "\n", "### Prerequisites\n", "| Concepts | Importance | Notes |\n", "| --- | --- | --- |\n", "| [Intro to Xarray](https://foundations.projectpythia.org/core/xarray/xarray-intro) | Necessary | |\n", "| Dask Arrays | Necessary | |\n", "| Dask DataFrames | Necessary | |\n", "\n", "- **Time to learn**: 40 minutes\n", "---------\n", "\n", "## Introduction\n", "\n", "### Xarray Quick Overview\n", "\n", "\n", " \n", "Xarray is an open-source Python library designed for working with *labelled multi-dimensional* data. By *multi-dimensional* data (also often called *N-dimensional*), we mean data that has many independent dimensions or axes (e.g. latitude, longitude, time). By labelled we mean that these axes or dimensions are associated with coordinate names (like \"latitude\") and coordinate labels like \"30 degrees North\".\n", "\n", "Xarray provides pandas-level convenience for working with this type of data.\n", "\n", "\n", "\"Xarray\"\n",\n", "\n", "*Image credit: Xarray Contributors*\n", "\n", "The dataset illustrated has two variables (`temperature` and `precipitation`) that have three dimensions. Coordinate vectors (e.g., latitude, longitude, time) that describe the data are also included.\n", "\n", " \n", "#### Xarray Data Structures\n", "\n", "Xarray has two fundamental data structures:\n", "\n", "* `DataArray` : holds a single multi-dimensional variable and its coordinates\n", "* `Dataset` : holds multiple DataArrays that potentially share the same coordinates\n", "\n", "\n", "**Xarray DataArray**\n", "\n", "A `DataArray` has four essential attributes:\n", "* `data`: a `numpy.ndarray` holding the values.\n", "* `dims`: dimension names for each axis (e.g., latitude, longitude, time).\n", "* `coords`: a dict-like container of arrays (coordinates) that label each point (e.g., 1-dimensional arrays of numbers, datetime objects or strings).\n", "* `attrs`: a dictionary to hold arbitrary metadata (attributes).\n", "\n", "**Xarray DataSet**\n", "\n", "A dataset is simply an object containing multiple Xarray DataArrays indexed by variable name." ] }, { "cell_type": "markdown", "id": "598055e0-3bac-491b-8b7f-d7313a306bc8", "metadata": {}, "source": [ "### Xarray can wrap many array types like Numpy and Dask.\n", "\n", "Let's start with a random 2D NumPy array, for example this can be SST (sea-surface temperature) values of a domain with dimension of 300x450 grid:" ] }, { "cell_type": "code", "execution_count": null, "id": "0b6d9a1d-6520-4374-a178-ad91af454628", "metadata": { "tags": [] }, "outputs": [], "source": [ "import numpy as np \n", "import dask.array as da\n", "import xarray as xr\n", "\n", "xr.set_options(display_expand_data=False);" ] }, { "cell_type": "code", "execution_count": null, "id": "c5185246-289d-4bc3-a355-4c5101bd6ddd", "metadata": { "tags": [] }, "outputs": [], "source": [ "# -- numpy array \n", "sst_np = np.random.rand(300,450)\n", "type(sst_np)" ] }, { "cell_type": "markdown", "id": "c17adcd3-672d-4a46-8b16-1f579aa29e8b", "metadata": {}, "source": [ "As we saw in the previous tutorial, we can convert them to a Dask Array:" ] }, { "cell_type": "code", "execution_count": null, "id": "744e09ca-7a23-428b-9032-1808610c19b0", "metadata": { "tags": [] }, "outputs": [], "source": [ "sst_da = da.from_array( sst_np)\n", "sst_da" ] }, { "cell_type": "markdown", "id": "62d098b3-a7fc-4562-bc7f-41c19b3c9280", "metadata": {}, "source": [ "This is great and fast! BUT\n", "* What if we want to attach coordinate values to this array?\n", "* What if we want to add metadata (e.g. units) to this array?" ] }, { "cell_type": "code", "execution_count": null, "id": "2006822d-c6db-4e67-995e-732d92ff10b6", "metadata": { "tags": [] }, "outputs": [], "source": [ "# similarly we can convert them to xarray datarray\n", "sst_xr = xr.DataArray(sst_da)\n", "sst_xr" ] }, { "cell_type": "markdown", "id": "d1eb3ef7-8413-4578-bd1f-93f488d6b344", "metadata": {}, "source": [ "A simple DataArray without dimensions or coordinates isn't much use." ] }, { "cell_type": "code", "execution_count": null, "id": "81aad904-f133-4a28-b1fc-aecc3fab7a11", "metadata": { "tags": [] }, "outputs": [], "source": [ "# we can add dimension names to this:\n", "sst_xr = xr.DataArray(sst_da,dims=['lat','lon'])\n", "\n", "sst_xr.dims" ] }, { "cell_type": "markdown", "id": "5dd6bea8-35ff-4f7f-91aa-cb808d30621c", "metadata": {}, "source": [ "We can add our coordinates with values to it :\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d9525a4b-e99b-46e5-9154-f0badf205ff9", "metadata": { "tags": [] }, "outputs": [], "source": [ "# -- create some dummy values for lat and lon dimensions\n", "lat = np.random.uniform(low=-90, high=90, size=300)\n", "lon = np.random.uniform(low=-180, high=180, size=450)\n", "\n", "sst_xr = xr.DataArray(sst_da,\n", " dims=['lat','lon'],\n", " coords={'lat': lat, 'lon':lon},\n", " attrs=dict(\n", " description=\"Sea Surface Temperature.\",\n", " units=\"degC\")\n", " )\n", "sst_xr" ] }, { "cell_type": "markdown", "id": "c58938ac-372a-42dc-8168-a59e4b45294b", "metadata": {}, "source": [ "Xarray data structures are a very powerful tool that allows us to use metadata to express different analysis patterns (slicing, selecting, groupby, averaging, and many other things). " ] }, { "cell_type": "markdown", "id": "a9be8aad-9135-45f5-a0e7-3d1d26e77869", "metadata": {}, "source": [ "
\n", "\n", "Take Away \n", "\n", "Xarray DataArray provides a wrapper around arrays, and uses labeled dimensions and coordinates to support metadata-aware operations (e.g. `da.sum(dim=\"time\")` instead of `array.sum(axis=-1)`)\n", "\n", "**Xarray can wrap dask arrays instead of numpy arrays.** \n", "\n", "This capability turns Xarray into an extremely useful tool for Big Data earth science.\n", "
" ] }, { "cell_type": "markdown", "id": "3e848a64-407d-4492-a680-56a292e9eec3", "metadata": {}, "source": [ "With this introduction, let's start our tutorial on features of Xarray and Dask:" ] }, { "cell_type": "markdown", "id": "82b9060b-7aab-48a0-ab22-7f07300b3fa9", "metadata": {}, "source": [ "### Setup: Spinning up a cluster" ] }, { "cell_type": "code", "execution_count": null, "id": "778c8357-fcca-4b06-bce6-9c7531c0181b", "metadata": { "tags": [] }, "outputs": [], "source": [ "from dask.distributed import LocalCluster, Client\n", "cluster = LocalCluster()\n", "client = Client(cluster)\n", "client\n" ] }, { "cell_type": "markdown", "id": "f4cc4275-c0ce-42fa-a6b4-033c597b5c9c", "metadata": {}, "source": [ "-----------\n", "## Reading data with Dask and Xarray\n", "\n", "### Reading multiple netCDF files with `open_mfdataset`\n", "\n", "Xarray provides a function called `open_dataset` function that allows us to load a netCDF dataset into a Python data structure. To read more about this function, please see [xarray `open_dataset` API documentation](https://docs.xarray.dev/en/stable/generated/xarray.open_dataset.html). \n", "\n", "Xarray also provides `open_mfdataset`, which open multiple files as a single xarray dataset. Passing the argument `parallel=True` will speed up reading multiple datasets by executing these tasks in parallel using Dask Delayed under the hood. \n", "\n", "In this example, we are going to examine a subset of CESM2 Large Ensemble Data Sets (LENS). We will use 2m temperature (TREFHT) for this analysis. \n", "\n", "To learn more about LENS dataset, please visit:\n", "* [LENS official website](https://www.cesm.ucar.edu/community-projects/lens/data-sets)\n", "* [LENS paper](https://www.cesm.ucar.edu/community-projects/lens/data-sets)\n", "\n", "For this tutorial, we only look at a small subset of data. If you don't have the data, running the following code enables you to download, prepare, and stage the required datasets (`../data/` folder) for this cookbook." ] }, { "cell_type": "code", "execution_count": null, "id": "a53169fd-46dd-41bd-94ca-aed13e0f6357", "metadata": { "tags": [] }, "outputs": [], "source": [ "!./get_data.sh notebook3" ] }, { "cell_type": "markdown", "id": "172706d1-73e3-41b8-b0af-3ec6168d6242", "metadata": {}, "source": [ "We can open up multiple files using `open_mfdataset` function. " ] }, { "cell_type": "markdown", "id": "8c23d960-93f4-44c0-9f3d-b92ff8083848", "metadata": {}, "source": [ "### Constructing Xarray Datasets from files" ] }, { "cell_type": "code", "execution_count": null, "id": "d1b5337d-28c3-4f88-b9c3-ce3d92445c37", "metadata": { "tags": [] }, "outputs": [], "source": [ "import os\n", "import glob\n", "\n", "var = 'TREFHT'\n", "\n", "# find all LENS files for 1 ensemble\n", "data_dir = '../data/data_for_cesm'\n", "files = glob.glob(os.path.join(data_dir, 'b.e21.BSSP370smbb.f09_g17.LE2-1301.013*.nc'))\n", "\n", "print(\"All files: [\", len(files), \"files]\")" ] }, { "cell_type": "code", "execution_count": null, "id": "920f5f93-db35-4ae8-b9bd-d9374efabe40", "metadata": { "tags": [] }, "outputs": [], "source": [ "%%time\n", "ds = xr.open_mfdataset(\n", " sorted(files),\n", " # concatenate along this dimension\n", " concat_dim=\"time\",\n", " # concatenate files in the order provided\n", " combine=\"nested\",\n", " # parallelize the reading of individual files using dask\n", " # This means the returned arrays will be dask arrays\n", " parallel=True,\n", " # these are netCDF4 files, use the h5netcdf package to read them\n", " engine=\"h5netcdf\",\n", " # hold off on decoding time\n", " decode_cf=False,\n", " # specify that data should be automatically chunked\n", " chunks=\"auto\",\n", ")\n", "ds = xr.decode_cf(ds)\n", "ds" ] }, { "cell_type": "markdown", "id": "2d757c05-8258-44ff-a232-1a278edc2e28", "metadata": {}, "source": [ "For complex scenarios, you can access each file individually by utilizing the `open_dataset` function with the specified `chunk`s and then combine the outputs into a single dataset later." ] }, { "cell_type": "markdown", "id": "56497135-6ad3-4323-bb8a-4f073eaf8939", "metadata": {}, "source": [ "Note that the \"real\" values are not displayed, since that would trigger actual computation. \n", "
\n", "\n", "Xarray automatically wraps Dask Arrays and Dask is lazy, meaning that operations are not computed until we explicitly request them, for example by calling `.compute()`.\n", "\n", "
\n", "\n", "\n", "Please see previous notebooks for more information on \"lazy evaluation\". \n", "\n", "The represntation of `TREFHT` DataArray shows details of chunks and chunk-sizes of Xarray DataArray:" ] }, { "cell_type": "code", "execution_count": null, "id": "ac192d70-f708-4db9-822c-4f2848cb0e02", "metadata": { "tags": [] }, "outputs": [], "source": [ "tref = ds.TREFHT\n", "tref" ] }, { "cell_type": "code", "execution_count": null, "id": "d511a271-fe12-4552-b59c-672605543cb2", "metadata": {}, "outputs": [], "source": [ "tref.chunks" ] }, { "cell_type": "markdown", "id": "0fbf4db0-f4b8-4b48-8337-12ace3882861", "metadata": {}, "source": [ "* How many chunks do we have? \n", "* What is the size of each chunk size?\n", "\n", "Here we can see that we have a total of 9 chunks - equal to the number of our netCDF files. In general `open_mfdataset` will return one chunk per netCDF file.\n", "\n", "
\n", "\n", "**WARNING:** The chunk structure within the file is important. When re-chunking the dataset after creation with `ds.chunk()` it is recommended to only use multiples of the on-file chunk shape.\n", "\n", "
\n", "\n", "We can check what that shape is by looking at the encoding:" ] }, { "cell_type": "code", "execution_count": null, "id": "c54779d9-92d5-4272-b300-94391b34bb52", "metadata": { "tags": [] }, "outputs": [], "source": [ "tref.encoding" ] }, { "cell_type": "markdown", "id": "d9592706-78ac-4ae1-b0b7-c626ab95ec16", "metadata": {}, "source": [ "
\n", "\n", "**TIP:** The `chunks` parameter can significantly affect total performance when using Dask Arrays. `chunks` should be small enough that each chunk fit in the memory, but large enough to avoid that the communication overhead. \n", "\n", "
" ] }, { "cell_type": "markdown", "id": "c094dcba-fdc8-4457-bf49-5c6ddf9b28a3", "metadata": {}, "source": [ "A good rule of thumb is to create arrays with a minimum chunksize of at least one million elements. Here we have 120x192x288 elements in each chunk (except for the last chunk). \n", "With large arrays (10+ GB), the cost of queuing up Dask operations can be noticeable, and you may need even larger chunksizes. \n", "\n", "**Additional Reading**\n", "\n", " - [dask.array best practices](https://docs.dask.org/en/stable/array-best-practices.html#select-a-good-chunk-size)\n", " - [NCAR chunking tutorial](https://ncar.github.io/dask-tutorial/notebooks/06-dask-chunking.html)\n", " - [Dask blog post on chunking](https://blog.dask.org/2021/11/02/choosing-dask-chunk-sizes)" ] }, { "cell_type": "markdown", "id": "ad1c1948-992e-463d-9631-35779ea204df", "metadata": {}, "source": [ "### Xarray data structures are Dask collections.\n", "\n", "This means you can call the following Dask-related functions on Xarray Data Arrays and Datasets:\n", "\n", "* `.visualize()`\n", "* `.compute()`\n", "* `.persist()`\n", "\n", "For more information about Dask Arrays, please see [Dask Array chapter](https://projectpythia.org/dask-cookbook/notebooks/01-dask-array)." ] }, { "cell_type": "code", "execution_count": null, "id": "d5468f03-be6e-44b9-92ef-2de287f3d990", "metadata": {}, "outputs": [], "source": [ "tref_mean = tref.mean('time')\n", "tref_mean.data.dask" ] }, { "cell_type": "markdown", "id": "f1628d55-7bda-42f4-b615-b86274571ec9", "metadata": {}, "source": [ "If we check Dask Task Graph for `tref_mean`, we can see all the steps required for calculating it (from opening the netcdf file to calculating mean and aggreagting it). " ] }, { "cell_type": "markdown", "id": "889e9971-43eb-4837-8d13-e560de21f908", "metadata": {}, "source": [ "### Getting concrete values\n", "At some point, you will want to actually do the calculations and receive concrete values from Dask.\n", "\n", "There are two ways to compute values on dask arrays.\n", "\n", " 1. `compute()` returns a new xarray object with the data now represented as a numpy array.\n", " 2. `load()` replaces the dask array in the xarray object with a numpy array. Equivalent to `ds = ds.compute()`.\n", "\n", "`.load()` operates *inplace* and `.compute()` returns a new xarray object.\n", "\n", "#### Distributed non-blocking concrete values\n", "\n", "There is another option available third option : “persisting”. `.persist()` loads the values into distributed RAM. The values are computed but remain distributed across workers. So essentially `persist` turns a lazy Dask collection into a Dask collection where the results are either fully computed or actively computing in the background.\n", "\n", "So `ds.air.persist()` is still backed by a Dask array. This is useful if you will be repeatedly using a dataset for computation but it is too large to load into local memory. \n", "\n", "Read more: [Dask user guide](https://docs.dask.org/en/stable/generated/dask.dataframe.Series.persist.html)\n" ] }, { "cell_type": "markdown", "id": "8ce5b8a3-cea9-4d04-bac0-dd574da5877d", "metadata": { "tags": [] }, "source": [ "## How to access underlying data in an Xarray object?" ] }, { "cell_type": "markdown", "id": "cc67499d-dbba-4844-9fed-cb125cc5048a", "metadata": {}, "source": [ "There are two basic ways to extract values from an Xarray object:\n", "\n", "1. Using `.data` will return a Dask array. For example:" ] }, { "cell_type": "code", "execution_count": null, "id": "de86a448-d04f-4a39-b1c6-85366734d590", "metadata": {}, "outputs": [], "source": [ "tref.data" ] }, { "cell_type": "markdown", "id": "3c3f5541-c879-40be-a0dc-85b54d498969", "metadata": {}, "source": [ "**This means that for Dask-backed Xarray object, we can access the values using `.compute`**" ] }, { "cell_type": "code", "execution_count": null, "id": "68961208-6a78-49a4-8787-b7937d66a5b6", "metadata": {}, "outputs": [], "source": [ "%%time\n", "\n", "tref.data.compute()" ] }, { "cell_type": "markdown", "id": "56b9429c-7a4f-4c86-982f-90656cc65662", "metadata": {}, "source": [ "2. We can also use `.values` to see the \"real\" values of Xarray object. Another option is using `.to_numpy`. Both of these option return the values of underlying Dask object in a numpy array. " ] }, { "cell_type": "code", "execution_count": null, "id": "e56f49b0-7865-4be4-8e36-011377cd20b4", "metadata": {}, "outputs": [], "source": [ "%%time\n", "tref.to_numpy()" ] }, { "cell_type": "markdown", "id": "7f2ce398-ff13-4251-945b-7b871bdfcb81", "metadata": {}, "source": [ "## Computation" ] }, { "cell_type": "markdown", "id": "b692aefa-001a-4dcc-9247-039f1790863f", "metadata": {}, "source": [ "All built-in Xarray methods (`.mean`, `.max`, `.rolling`, `.groupby` etc.) support dask arrays.\n", "\n", "Now, let's do some computations on this Xarray dataset. " ] }, { "cell_type": "markdown", "id": "0f4d163a-bc44-41f9-b190-ea5cdcd91cd1", "metadata": {}, "source": [ "### Single Point Calculations\n", "\n", "To start out, let's do the calculations on a single point first. First, we extract the time series data at a grid point and save it to a variable. Here we select the closest point using `.sel` and load the data. " ] }, { "cell_type": "code", "execution_count": null, "id": "c28fb0e2-6de1-45ca-ba79-7d50b3aa1842", "metadata": {}, "outputs": [], "source": [ "tref_boulder = tref.sel(lat=40.0150, lon=-105.2705, method='nearest').load()" ] }, { "cell_type": "markdown", "id": "dea9cfcd-ba5e-4c45-bfd1-c079c7a34e73", "metadata": {}, "source": [ "
\n", "\n", "**WARNING:** Remember as soon as you call `.load()` you are telling Dask to trigger computation.\n", "\n", "
" ] }, { "cell_type": "code", "execution_count": null, "id": "afa4f016-4a0a-4220-ad1f-bc4b7d1ecc32", "metadata": {}, "outputs": [], "source": [ "# -- take annual average\n", "tb = tref_boulder.resample(time='AS').mean()\n", "tb" ] }, { "cell_type": "markdown", "id": "06872e46-54ec-4df1-a0cb-4e25653f5f7d", "metadata": {}, "source": [ "We can either see the values of our DataArray in the text representation above or by plotting it:" ] }, { "cell_type": "code", "execution_count": null, "id": "2e829f32-11ff-4e7a-bf7c-ceb3d606793f", "metadata": {}, "outputs": [], "source": [ "tb.plot()" ] }, { "cell_type": "markdown", "id": "9c868e1f-6d5a-42a9-94c0-3c8bf70e8020", "metadata": {}, "source": [ "### Calculations over all grids" ] }, { "cell_type": "code", "execution_count": null, "id": "4f8ecc4c-f99a-4a6f-b46c-608e7ec2e60e", "metadata": {}, "outputs": [], "source": [ "# change the unit from Kelvin to degree Celsius \n", "tref_c = tref - 273.15\n", "tref_c" ] }, { "cell_type": "code", "execution_count": null, "id": "225daee2-b0f5-4cca-8105-c053b6f14106", "metadata": {}, "outputs": [], "source": [ "%%time \n", "\n", "tref_c = tref_c.load()" ] }, { "cell_type": "code", "execution_count": null, "id": "be6fa1c8-a550-4fd1-8418-cae2844fe657", "metadata": {}, "outputs": [], "source": [ "%%time\n", "\n", "# Compute monthly anomaly\n", "\n", "# -- 1. calculate monthly average\n", "tref_grouped = tref.groupby('time.month')\n", "tmean = tref_grouped.mean(dim='time')\n", "\n", "#-- 2. calculate monthly anomaly\n", "tos_anom = tref_grouped - tmean\n", "tos_anom" ] }, { "cell_type": "code", "execution_count": null, "id": "e9df3af5-4089-4014-aecc-cb2fe62a0f0d", "metadata": {}, "outputs": [], "source": [ "%%time\n", "computed_anom = tos_anom.load()\n", "type(computed_anom)" ] }, { "cell_type": "code", "execution_count": null, "id": "7a161d2b-9321-459f-9b36-c4ae73115eab", "metadata": {}, "outputs": [], "source": [ "tos_anom.sel(lon=310, lat=50, method='nearest').plot( size=4)" ] }, { "cell_type": "code", "execution_count": null, "id": "40d5ee76-a2ae-41d7-9c45-46cf9c19bc7a", "metadata": {}, "outputs": [], "source": [ "tos_anom.sel(time='2030-01-01').plot()" ] }, { "cell_type": "markdown", "id": "b6d15e18-788b-4763-b4ae-f4ee0f644cf3", "metadata": {}, "source": [ "
\n", "\n", "TIP: Using Xarray plotting functionality automatically triggers computations on the Dask Array, similar to `.compute()`.\n", "\n", "
" ] }, { "cell_type": "markdown", "id": "3da23b0c-d3f7-4bfb-aa80-19eeb640823e", "metadata": {}, "source": [ "We can do more complex calculations too:" ] }, { "cell_type": "code", "execution_count": null, "id": "fb13b6b0-b881-44b4-82af-60514515a8eb", "metadata": {}, "outputs": [], "source": [ "rolling_mean = tref.rolling(time=5).mean()\n", "rolling_mean # contains dask array" ] }, { "cell_type": "code", "execution_count": null, "id": "edacaaa8-61b7-4f68-a751-a9f71d5c4dde", "metadata": {}, "outputs": [], "source": [ "timeseries = rolling_mean.isel(lon=1, lat=20) # no activity on dashboard\n", "timeseries # contains dask array" ] }, { "cell_type": "code", "execution_count": null, "id": "20a6caff-aa35-4968-8c71-bd083c3cd136", "metadata": {}, "outputs": [], "source": [ "computed = rolling_mean.compute() # activity on dashboard\n", "computed # has real numpy values" ] }, { "cell_type": "markdown", "id": "dfe03506-897e-4ae8-bf5c-b990f43882d8", "metadata": {}, "source": [ "### Supplementary Material: Advanced workflows and automatic parallelization using `apply_ufunc`\n", "Most of xarray’s built-in operations work on Dask arrays. If you want to use a function that isn’t wrapped by Xarray to work with Dask, 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()` function. `xr.apply_ufunc()` can automate embarrassingly parallel “map” type operations where a function written for processing NumPy arrays, but we want to apply it on our Xarray DataArray. \n", "\n", "`xr.apply_ufunc()` give users capability to run custom-written functions such as parameter calculations in a parallel way. See the [Xarray tutorial material on apply_ufunc](https://tutorial.xarray.dev/advanced/apply_ufunc/apply_ufunc.html) for more.\n", "\n", "In the example below, we calculate the saturation vapor pressure by using `apply_ufunc()` to apply this function to our Dask Array chunk by chunk. " ] }, { "cell_type": "code", "execution_count": null, "id": "c67a70aa-93f2-40d8-b8f3-896602603648", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "\n", "def sat_p(t):\n", " \"\"\"Calculate saturation vapor pressure using Clausius-Clapeyron equation\"\"\"\n", " return 0.611 * np.exp(17.67 * (t-273.15)*((t-29.65)**(-1)))" ] }, { "cell_type": "code", "execution_count": null, "id": "3e3aa9ad-67e6-4b95-b1ed-dd238dc12d32", "metadata": {}, "outputs": [], "source": [ "es = xr.apply_ufunc(sat_p, tref, dask=\"parallelized\", output_dtypes=[float])\n", "es" ] }, { "cell_type": "code", "execution_count": null, "id": "2dea08bd-fb26-49ab-ab13-f7c0ebd5c25c", "metadata": {}, "outputs": [], "source": [ "es.compute()" ] }, { "cell_type": "markdown", "id": "43e90ad3-3af4-43bc-8f36-d4aff2960263", "metadata": {}, "source": [ "The data used for this tutorial is from one ensemble member. What if we want to use multiple ensemble members? So far, we only run on one machine, what if we run an HPC cluster? We will go over this in the next tutorial." ] }, { "cell_type": "markdown", "id": "5169db51-92ba-4257-a4ab-d2bd6b524909", "metadata": {}, "source": [ "### Dask + Xarray Good Practices\n", "\n", "
\n", "\n", "Summary of Dask + Xarray Good Practices \n", "\n", "\n", "The good practices regarding Dask + Xarray is the same as the good practices for Dask only.\n", "\n", "* Similar to Dask DataFrames, it is more efficient to first do spatial and temporal indexing (e.g. .sel() or .isel()) and filter the dataset early in the pipeline, especially before calling resample() or groupby(). \n", "\n", "* Chunk sizes should be small enough to fit into the memory at once but large enough to avoid the additional communication overhead. Good chunk size ~100 MB. \n", "\n", "* It is always better to chunk along the `time` dimension. \n", "\n", "* Avoid too many tasks since each task will introduce 1ms of overhead. \n", "\n", "* When possible, use `xr.apply_ufunc` to apply an unvectorized function to the Xarray object. \n", "\n", "
\n", "\n", "\n", "\n", "### Close you local Dask Cluster\n", "It is always a good practice to close the Dask cluster you created." ] }, { "cell_type": "code", "execution_count": null, "id": "9bd19bbe-d887-4fbb-b4eb-5dc0881e9bf6", "metadata": {}, "outputs": [], "source": [ "client.shutdown()" ] }, { "cell_type": "markdown", "id": "e0d12769-eaba-47fc-97e0-6d5a4298fbc4", "metadata": {}, "source": [ "## Summary\n", "\n", "In this notebook, we have learned about:\n", "\n", "* Using Dask with Xarray\n", "* Read/write netCDF files with Dask\n", "* Dask backed Xarray objects and operations\n", "* Extract Dask arrays from Xarray objects and use Dask array directly..\n", "* Customized workflows using `apply_ufunc`\n", "\n", "## Resources and references\n", "\n", "### Reference\n", "* [Dask Array Docs](https://docs.dask.org/en/stable/array.html)\n", "* [Dask Examples](https://examples.dask.org/)\n", "* [Dask Code](https://github.com/dask/dask/)\n", "* [Dask Blog](https://blog.dask.org/)\n", "* [Xarray Docs](https://xarray.pydata.org/)\n", "* [Xarray + Dask docs](https://docs.xarray.dev/en/stable/user-guide/dask.html), particularly the [best practices](https://docs.xarray.dev/en/stable/user-guide/dask.html#best-practices)\n", "* [Xarray Tutorial material](https://tutorial.xarray.dev/intro.html)\n", " \n", "\n", "### Ask for help\n", "* [`dask`](http://stackoverflow.com/questions/tagged/dask) tag on Stack Overflow, for usage questions\n", "* [github discussions: dask](https://github.com/dask/dask/discussions) for general, non-bug, discussion, and usage questions\n", "* [github issues: dask](https://github.com/dask/dask/issues/new) for bug reports and feature requests\n", "* [github discussions: xarray](https://github.com/pydata/xarray/discussions) for general, non-bug, discussion, and usage questions\n", "* [github issues: xarray](https://github.com/pydata/xarray/issues/new) for bug reports and feature requests" ] }, { "cell_type": "code", "execution_count": null, "id": "57f2ad23-bd82-4fd2-89ef-8ca07feba117", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.10.12" }, "widgets": { "application/vnd.jupyter.widget-state+json": { "state": {}, "version_major": 2, "version_minor": 0 } } }, "nbformat": 4, "nbformat_minor": 5 }