将 scipy.ndimage.convolve 应用于三维 xarray DataArray

Question

我有一个xarray尺寸为 x、y、z的 3D DataArray，我试图scipy.ndimage.convolve在每个 xy 平面上应用，同时将输出保持为 DataArray。当然，我正试图用它xr.apply_ufunc来做到这一点。如果我只为一架飞机这样做，它就可以完美运行：

da=xr.DataArray(np.random.rand(5,5,5), dims=("x", "y", "z"))

kernel=np.ones((3,3))

from scipy.ndimage import convolve

conv1 = lambda x: convolve(x, kernel, mode="wrap")

print(xr.apply_ufunc(conv1, da[:,:,0])) # works successfully

我现在试图想出一种方法来为每个 xy 平面做同样的事情。我认为会起作用的是使用np.apply_along_axisor np.apply_over_axes，但它们都不起作用。

我可以遍历轴，将所有内容都放在一个列表中，然后进行连接，但我试图使用它xr.apply_ufunc来避免属性出现问题。有没有办法做到这一点？

这是我认为应该工作的示例，但事实并非如此：

np.apply_over_axes(conv1, c, axes=(0,1))

但这失败了

TypeError: <lambda>() takes 1 positional argument but 2 were given

Answer 1

使用形状为 (3, 3, 1) 的内核代替 (3, 3) 怎么样？kernel2d = np.ones((3, 3))conv2d = lambda x: convolve(x, kernel2d, mode="wrap")result2d = xr.apply_ufunc(conv2d, da[:, :, 0])kernel3d = np.ones((3, 3, 1))conv3d = lambda x: convolve(x, kernel3d, mode="wrap")result3d = xr.apply_ufunc(conv3d, da)(result2d == result3d[:, :, 0]).all()&nbsp; # -> True另一种选择是在 中使用矢量化逻辑xr.apply_ufunc，这可能更接近您尝试执行的操作kernel = np.ones((3, 3))conv = lambda x: convolve(x, kernel, mode="wrap")result = xr.apply_ufunc(conv, da, input_core_dims=[['x', 'y']],&nbsp;&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; output_core_dims=[['x', 'y']],&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; vectorize=True)(result2d == result.transpose('x', 'y', 'z')).all()&nbsp; # --> True此选项只是为了方便而准备的，因此它可能比计算矢量化的第一个要慢得多。

Answer 2

我想出的一个可能的答案是手动执行此操作：def conv_rx(da, axis="z"):&nbsp; &nbsp; planes = [ xr.apply_ufunc(conv1, da.sel(z=z)) for z in da.z ]&nbsp; &nbsp; new = xr.concat(planes, dim=axis)&nbsp; &nbsp; return new.transpose(*da.dims)这会产生正确的结果。但是，我对此不太满意，因为它不优雅而且速度很慢。