def da_concat()

in bayesmark/xr_util.py [0:0]

• 20 lines of code
• 5 McCabe index (conditional complexity)

def da_concat(da_dict, dims):
    """Concatenate a dictionary of :class:`xarray:xarray.DataArray` similar to :func:`pandas:pandas.concat`.

    Parameters
    ----------
    da_dict : dict(tuple(str), :class:`xarray:xarray.DataArray`)
        Dictionary of :class:`xarray:xarray.DataArray` to combine. The keys are tuples of index values. The
        :class:`xarray:xarray.DataArray` must have compatible coordinates.
    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 `da_dict`.

    Returns
    -------
    da : :class:`xarray:xarray.DataArray`
        Combined data array. The new dimensions will be ``input_da.dims + dims``.
    """
    assert len(da_dict) > 0
    assert all(len(da.dims) > 0 for da in da_dict.values()), "0-dimensional DataArray not supported"
    assert all_unique(dims)

    cur_dims = list(da_dict.values())[0].dims
    assert all(da.dims == cur_dims for da in da_dict.values())
    assert len(set(cur_dims) & set(dims)) == 0

    def squeeze(tt):
        if len(tt) == 1:
            return tt[0]
        return tt

    D = OrderedDict([(squeeze(kk), da.to_series()) for kk, da in da_dict.items()])
    df = pd.concat(D, axis=1)

    assert df.columns.nlevels == len(dims)
    df.columns.names = dims

    df = df.stack(level=list(range(df.columns.nlevels)))
    assert isinstance(df, pd.Series)
    da = df.to_xarray()
    assert isinstance(da, xr.DataArray)
    return da