<jupyter_start><jupyter_text>Writing multiple netCDF files in parallel with xarray and daskA typical computation workflow with xarray consists of:1. reading one or more netCDF files into an xarray dataset backed by dask using `xr.open_mfdataset()` or `xr.open_dataset(chunks=...)`,2. applying some transformation to the input dataset, and3. saving the resulting output to disk in a netCDF file using `xr.to_netcdf()`.The last step (3) can easily lead to a large netCDF file (>=10GB in size). As a result, this step can take a very long time to complete (since it is run in serial), and sometimes may hang. So, to avoid these issues one can use one of the lesser-used but helpful xarray capabilities: the [`xr.save_mfdataset()`](https://xarray.pydata.org/en/latest/generated/xarray.save_mfdataset.html) function. This function allows users to write multiple datasets to disk as netCDF files simultaneously. The `xr.save_mfdataset()` function signature looks like this:```pythonxr.save_mfdataset(    datasets,    paths,    mode='w',    format=None,    groups=None,    engine=None,    compute=True,)Docstring:Write multiple datasets to disk as netCDF files simultaneously.```  Please show me the code  Package imports<jupyter_code>import numpy as np
import xarray as xr
from distributed import Client, performance_report

xr.__version__
client = Client()
client<jupyter_output><empty_output><jupyter_text>Load some toy dataset<jupyter_code>ds = xr.tutorial.open_dataset('rasm', chunks={'time': 12})
ds
ds.Tair.data.visualize()<jupyter_output><empty_output><jupyter_text>Perform a computation on input dataset As you can see, our input dataset is chunked along the `time` dimension. This produces an xarray dataset with 3 chunks. For illustrative purposes, let's apply some arbitray computation on our dataset:<jupyter_code>result = np.sqrt(np.sin(ds) ** 2 + np.cos(ds) ** 2)
result.Tair.data.visualize()<jupyter_output><empty_output><jupyter_text>Our resulting xarray dataset has the same dimensions and the same number of chunks as our input dataset.  Create a helper function to split a dataset into sub-datasetsOur main objective is to save this resulting dataset into multiple netCDF files by saving each chunk in its own netCDF file. To accomplish this, we need to create a function that allows us to split a dataset into sub-datasets for each chunk. The following code snippet was adapted from this [comment](https://github.com/pydata/xarray/issues/1093issuecomment-259213382) by [<NAME>](https://github.com/shoyer).<jupyter_code>import itertools


def split_by_chunks(dataset):
    chunk_slices = {}
    for dim, chunks in dataset.chunks.items():
        slices = []
        start = 0
        for chunk in chunks:
            if start >= dataset.sizes[dim]:
                break
            stop = start + chunk
            slices.append(slice(start, stop))
            start = stop
        chunk_slices[dim] = slices
    for slices in itertools.product(*chunk_slices.values()):
        selection = dict(zip(chunk_slices.keys(), slices))
        yield dataset[selection]
datasets = list(split_by_chunks(result))<jupyter_output><empty_output><jupyter_text>Let's confirm that we have three items corresponding to our three chunks:<jupyter_code>print(len(datasets))<jupyter_output>3<jupyter_text>Each item in the returned list consists of an xarray dataset corresponding to each sliced chunk. Now that we have a list of datasets, we need another helper function that will generate a filepath for us when given an xarray dataset as input:  Create a helper function for generating a filepath<jupyter_code>def create_filepath(ds, prefix='filename', root_path="."):
    """
    Generate a filepath when given an xarray dataset
    """
    start = ds.time.data[0].strftime("%Y-%m-%d")
    end = ds.time.data[-1].strftime("%Y-%m-%d")
    filepath = f'{root_path}/{prefix}_{start}_{end}.nc'
    return filepath<jupyter_output><empty_output><jupyter_text>Let's run a small test to make sure this function is working as expected:<jupyter_code>create_filepath(datasets[1])<jupyter_output><empty_output><jupyter_text>Let's create a list of paths to which to save each corresponding dataset.<jupyter_code>paths = [create_filepath(ds) for ds in datasets]
paths<jupyter_output><empty_output><jupyter_text>Invoke `xr.save_mfdataset()` At this point we have two lists:    1. `datasets`: List of datasets to save.2. `paths`: List of paths to which to save each corresponding dataset.We are ready to invoke the `xr.save_mfdataset()`:<jupyter_code>xr.save_mfdataset(datasets=datasets, paths=paths)<jupyter_output><empty_output><jupyter_text>Confirm that the output files were properly written<jupyter_code>new_ds = xr.open_mfdataset(paths, combine='by_coords')
new_ds
try:
    xr.testing.assert_identical(result, new_ds)
except AssertionError:
    print('The datasets are not identical!')
else:
    print('The datasets are identical!')<jupyter_output>The datasets are identical!
{'fraction_non_alphanumeric': 0.09131223543640395, 'fraction_numerical': 0.006248740173352147, 'mean_word_length': 5.014545454545455, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 42, 'https://': 3, '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': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '657455', 'license': 'Apache-2.0', 'n_tokens_mistral': 1503, 'n_tokens_neox': 1461, 'n_words': 511, 'path': 'posts/2020/writing-multiple-netcdf-files-in-parallel-with-xarray-and-dask.ipynb', 'repo_name': 'jukent/esds', 'size': '4968'}