def ds_concat()

in bayesmark/xr_util.py [0:0]

• 11 lines of code
• 6 McCabe index (conditional complexity)

def ds_concat(ds_dict, dims):
    """Concatenate a dictionary of :class:`xarray:xarray.Dataset` similar to :func:`pandas:pandas.concat`, and a
    generalization of :func:`.da_concat`.

    Parameters
    ----------
    ds_dict : dict(tuple(str), :class:`xarray:xarray.DataArray`)
        Dictionary of :class:`xarray:xarray.Dataset` to combine. The keys are tuples of index values. The
        :class:`xarray:xarray.Dataset` must have compatible coordinates, and all have the same variables.
    dims : list(str)
        The names of the new dimensions we create for the dictionary keys. This must be of the same length as the
        key tuples in `ds_dict`.

    Returns
    -------
    ds : :class:`xarray:xarray.Dataset`
        Combined dataset. For each variable `var`, the new dimensions will be ``input_ds[var].dims + dims``.
    """
    assert len(ds_dict) > 0
    assert len(dims) > 0
    assert all(len(kk) == len(dims) for kk in ds_dict)

    # Get an arbitrary element as the reference
    k0 = list(ds_dict.keys())[0]

    # Check all vars the same
    vars_, = set([tuple(ds) for ds in ds_dict.values()])

    # Now combine da for each variable, one at a time
    ds = xr.Dataset(coords=ds_dict[k0].coords)
    for vv in vars_:
        da_dict = OrderedDict([(kk, da[vv]) for kk, da in ds_dict.items()])
        ds[vv] = da_concat(da_dict, dims)

    return ds