Looping over and mapping slices of a DataArray in xarray-CodePudding

Say I have a multidimensional DataArray and I would like to loop over slices according to some dimension and change them. For example, the first dimension is time and I would like for each time to receive a DataArray that represents a slice of that time and map it to a different DataArray of the same size.

I can use apply_ufunc but then I lose the ability to use labelled dimensions within the function that operates on them. I thought maybe I could use map_blocks but I couldn't understand how to specify a dimension to loop over.

Edit:

I've implemented what I wanted to do like this:

def xarray_map_over_dimension(
    data_array: xr.DataArray, func: Callable, dim: str, *args, **kwargs
) -> None:
    """
    For an n-dimensional DataArray this function will map n-1 dimensional slices by iterating over a given dimension.
    This is similar to xarray's apply_ufunc except that it calls :func with DataArray objects rather than ndarrays

    WARNING: this function will modify data_array inplace
    """
    for i, data_slice in data_array.groupby(dim):
        data_array[{dim: i}] = func(data_slice, *args, **kwargs)

Basically I wanted something like apply_ufunc that would give me data_array slices instead of ndarray slices. I thought this was quite a common use case so I figured there would be a standard way to do this.

CodePudding user response：

Ok turns out what I wanted was data_array.groupby(dim).map(func) https://docs.xarray.dev/en/stable/generated/xarray.core.groupby.DatasetGroupBy.map.html