Apply a function along the dimension

[mhungen]

Hi,
How can I apply a function (which maps a DataArray to a DataArray) along a dimension and combine the result back to a DataArray (same dims and coords as input)?

I can't use xarray.apply_ufunc because I need the labeled data inside the function that I want to apply.

I found a "hack" that utilizes groupBy, but groupBy seems to sort the coordinates of my dimension (which makes sense for grouping by time).

Is there a better solution for that?

import numpy as np
import pytest
import xarray as xr

c, y, x = 4, 20, 20
rgba_img = xr.DataArray(
    data=np.zeros(y * x * c).reshape(y, x, c),
    dims=['y', 'x', 'c'],
    coords={
        'c': ['r', 'g', 'b', 'a']
    }
)

def func_2d(da_2d: xr.DataArray) -> xr.DataArray:
    # do important stuff on the 2D DataArray
    
    if len(da_2d.shape) != 2:
        raise ValueError
    return da_2d + 1

# what I expected to work
with pytest.raises(ValueError):
    rgba_img.groupby('c', squeeze=True).map(func_2d)

# works but the index will be sorted
mapped_v1 = rgba_img.groupby('c').map(lambda f: func_2d(f.squeeze()))
assert mapped_v1.dims == ('y', 'x', 'c')
assert np.all(mapped_v1.coords['c'] == ['a', 'b', 'g', 'r']) # not what I want

# that works but is a bit complex
mapped_v2 = (rgba_img
             .sortby('c')
             .groupby('c', squeeze=True)
             .map(func_2d)
             .reindex_like(rgba_img)
             )
assert set(mapped_v2.sizes.items()) == {
    ('c', c),
    ('y', y),
    ('x', x)
}
assert mapped_v2.dims == ('y', 'x', 'c')

[keewis]

if you have a function that consumes and produces DataArray objects, you can just call this directly on rgba_array (no need to do fancy things like apply_ufunc). Though if you need to vectorize all other dimensions, you might need to use map_blocks (be aware, though, that this can be pretty slow).

[mhungen]

In the given example, I want to apply a correction to each channel of an RGBA image. It would be possible to adjust the correction function to handle 3 dimensions, but that adds a lot of complexity to this function (should also handle RGB and grayscale). What if someone wanted to do the same analysis over a time series of RGBA images (e.g.: adding a fourth dimension)?

My approach was to break the data down in the simplest way possible (2D) and use xarray to split the data into easily managable tasks. apply_ufunc would be a perfect fit but i need the metadata (georeferencing).
Is there a solution for this as well?

[keewis]

you can pass additional arguments to the function applied by apply_ufunc using the kwargs keyword argument. So what I'd recommend is to write a function that works on numpy arrays and is vectorized (possibly using numba.guvectorize or numpy.vectorize), and you pass in the georeference data as kwargs={"spatial_ref": crs} or similar.