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.
