-
|
Hi, 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') |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 2 replies
-
|
if you have a function that consumes and produces |
Beta Was this translation helpful? Give feedback.
-
|
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. |
Beta Was this translation helpful? Give feedback.
-
|
you can pass additional arguments to the function applied by |
Beta Was this translation helpful? Give feedback.
you can pass additional arguments to the function applied by
apply_ufuncusing thekwargskeyword argument. So what I'd recommend is to write a function that works onnumpyarrays and is vectorized (possibly usingnumba.guvectorizeornumpy.vectorize), and you pass in the georeference data askwargs={"spatial_ref": crs}or similar.