<jupyter_start><jupyter_text>Dask and XarrayThis notebook demonstrates one of xarray's most powerful features: the ability to wrap dask arrays and allow users to seamlessly execute analysis code in parallel. Learning Objectives- Learn that xarray DataArrays and Datasets are "dask collections" i.e. you can execute top-level dask functions such as dask.visualize(xarray_object)- Learn that all xarray built-in operations can transparently use dask- Learn that xarray provides tools to easily parallelize custom functions across blocks of dask-backed xarray objects. Prerequisites| Concepts | Importance | Notes || --- | --- | --- || Familiarity with Dask Array | Necessary | || Familiarity with xarray | Necessary | |- **Time to learn**: *15-20 minutes*  SetupFirst let's set up a `LocalCluster` using `dask.distributed`.<jupyter_code>import dask
import dask.array as da
import xarray as xr
from dask.distributed import Client, LocalCluster
cluster = LocalCluster()
client = Client(cluster)
client<jupyter_output><empty_output><jupyter_text>Reading data with Dask and XarrayRecall that a dask's array consists of many chunked arrays:<jupyter_code>darr = da.ones((2000, 300), chunks=(200, 50))
darr
darr.compute()<jupyter_output><empty_output><jupyter_text>To read data as dask arrays with xarray, we need to specify the `chunks` argument to `open_dataset()` function.<jupyter_code>ds = xr.open_dataset(
    "data/tos_Omon_CESM2_historical_r11i1p1f1_gr_200001-201412.nc", engine="netcdf4", chunks={}
)
ds<jupyter_output><empty_output><jupyter_text>Passing `chunks={}` to `open_dataset()` works, but since we didn't tell dask how to split up (or chunk) the array, Dask will create a single chunk for our array.<jupyter_code>ds = xr.open_dataset(
    "data/tos_Omon_CESM2_historical_r11i1p1f1_gr_200001-201412.nc",
    engine="netcdf4",
    chunks={"time": 90, "lat": 180, "lon": 360},
)
ds
ds.tos
ds.tos.chunks<jupyter_output><empty_output><jupyter_text>Xarray data structures are first-class dask collectionsThis means you can call the following functions - `dask.visualize(...)`- `dask.compute(...)`- `dask.persist(...)`on both xarray DataArrays and Datasets backed by dask-arrays.<jupyter_code>dask.visualize(ds)<jupyter_output><empty_output><jupyter_text>Parallel and Lazy computation using `dask.array` with xarrayXarray seamlessly wraps dask so all computation is deferred until explicitly requested.<jupyter_code>z = ds.tos.mean(['lat', 'lon']).dot(ds.tos.T)
z<jupyter_output><empty_output><jupyter_text>As you can see, `z` contains a dask array. This is true for all xarray built-in operations including subsetting<jupyter_code>z.isel(lat=0)
dask.visualize(z)
%%time
z.compute()<jupyter_output><empty_output><jupyter_text>Reading multiple datasets with `open_mfdataset`Xarray provides a built-in function `xr.open_mfdataset()` for opening multiple files as a single dataset. This makes it easy to work with data from multiple files as one logical dataset. For demonstration purposes, let's revisit our example in [Dask Delayed Notebook](./08-dask-delayed.ipynb). In this example, we loop over a list of files (for four ensemble members), and we compute the anomaly for each ensemble member as follows:<jupyter_code>import pathlib

data_dir = pathlib.Path("data/")
files = sorted(data_dir.glob("tos_Omon_CESM2*"))

results = {}
for file in files:

    # Read in file
    ds = dask.delayed(xr.open_dataset)(file, engine='netcdf4')

    # Compute anomaly
    gb = ds.tos.groupby('time.month')
    tos_anom = gb - gb.mean(dim='time')

    # Save the computed anomaly and record the name of the ensemble member
    results[file.stem.split('_')[-3]] = tos_anom


# Compute the results
# dask.compute() returns a tuple here with a single item. So, ensure to grab this one item by using the 0 index
computed_results = dask.compute(results)[0]
# Combine the results in our dataarray by concatenating the results across a new dimension `ensemble_member`
dset_anom = xr.concat(list(computed_results.values()), dim='ensemble_member')
dset_anom['ensemble_member'] = list(computed_results.keys())
dset_anom<jupyter_output><empty_output><jupyter_text>Instead of explicitly looping over the list of files to construct xarray datasets, we can pass the list of files to [`xr.open_mfdataset()`](https://xarray.pydata.org/en/stable/generated/xarray.open_mfdataset.htmlxarray.open_mfdataset) and xarray will construct a single dataset for us:<jupyter_code>dset = xr.open_mfdataset(
    sorted(files),
    concat_dim='ensemble_member',
    combine="nested",
    parallel=True,
    data_vars=['tos'],
    engine="netcdf4",
    chunks={'time': 90},
)
# Add coordinate labels for the newly created `ensemble_member` dimension
dset["ensemble_member"] = ['r11i1p1f1', 'r7i1p1f1', 'r8i1p1f1', 'r9i1p1f1']
dset<jupyter_output><empty_output><jupyter_text>By default, open_mfdataset() will chunk each netCDF file into a single Dask array; supply the chunks argument to control the size of the resulting Dask arrays.    In more complex cases, you can open each file individually using open_dataset(..., chunks={...}) and merge the results into a single dataset.        Passing the keyword argument parallel=True to open_mfdataset() will speed up the reading of large multi-file datasets by executing those read tasks in parallel using dask.delayed.<jupyter_code># Compute anomaly
gb = dset.tos.groupby('time.month')
tos_anom = gb - gb.mean(dim='time')
tos_anom
tos_anom.sel(lon=310, lat=50, method='nearest').plot(col='ensemble_member', col_wrap=2, size=4);<jupyter_output><empty_output><jupyter_text>Note that using plotting functionality will automatically trigger computation of required results. So, with xarray's `open_mfdataset()`, the following code```pythonresults = {}for file in files:     Read in file    ds = dask.delayed(xr.open_dataset)(file, engine='netcdf4')     Compute anomaly    gb = ds.tos.groupby('time.month')    tos_anom = gb - gb.mean(dim='time')     Save the computed anomaly and record the name of the ensemble member    results[file.stem.split('_')[-3]] = tos_anom Compute the results dask.compute() returns a tuple here with a single item. So, ensure to grab this one item by using the 0 indexcomputed_results = dask.compute(results)[0] Combine the results in our dataarray by concatenating the results across a new dimension `ensemble_member`dset_anom = xr.concat(list(computed_results.values()), dim='ensemble_member')dset_anom['ensemble_member'] = list(computed_results.keys())```becomes ```pythondset = xr.open_mfdataset(sorted(files), concat_dim='ensemble_member',                          combine="nested", parallel=True, data_vars=['tos'],                         engine="netcdf4", chunks={'time': 90}) Add coordinate labels for the newly created `ensemble_member` dimensiondset["ensemble_member"] = ['r11i1p1f1', 'r7i1p1f1', 'r8i1p1f1', 'r9i1p1f1']  Compute anomalygb = dset.tos.groupby('time.month')tos_anom = gb - gb.mean(dim='time')```This latter version is cleaner and easier to maintain than the version with loops.<jupyter_code>cluster.close()
client.close()
%load_ext watermark
%watermark --time --python --updated --iversion<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10750139431121027, 'fraction_numerical': 0.01687116564417178, 'mean_word_length': 5.586776859504132, 'pattern_counts': {'":': 3, '<': 45, '<?xml version=': 0, '>': 45, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '430206', 'license': 'CC0-1.0', 'n_tokens_mistral': 2380, 'n_tokens_neox': 2291, 'n_words': 696, 'path': 'notebooks/10-dask-and-xarray.ipynb', 'repo_name': 'mgrover1/xarray-tutorial', 'size': '7172'}