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

Question

我有一个尺寸为 x、y、z 的 3D xarray DataArray，我正在尝试在每个 x-y 平面上应用 scipy.ndimage.convolve，同时将输出保持为 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

我现在正在尝试想出一种方法来对每个 x-y 平面执行相同的操作。我认为可以使用 np.apply_along_axis 或 np.apply_over_axes，但其中 none 有效。

我可以遍历轴，将所有内容放在列表中，然后连接，但我正在尝试使用 xr.apply_ufunc 来避免属性问题。有办法吗？

这是一个我认为应该起作用但不起作用的例子：

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

但这失败了

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

Answer 1

我想到的一个可能的答案是手动执行此操作：

def conv_rx(da, axis="z"):
    planes = [ xr.apply_ufunc(conv1, da.sel(z=z)) for z in da.z ]
    new = xr.concat(planes, dim=axis)
    return new.transpose(*da.dims)

这会产生正确的结果。但是，我对此不是很满意，因为它不优雅而且速度很慢。

Answer 2

使用形状为 (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()  # -> 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']], 
                        output_core_dims=[['x', 'y']],
                        vectorize=True)
(result2d == result.transpose('x', 'y', 'z')).all()  # --> True

这个选项只是为了方便而准备的，因此它可能比第一个计算向量化的选项慢得多。

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

Applying scipy.ndimage.convolve to tridimensional xarray DataArray

python

scipy

python-xarray