将 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 convolveconv1 = 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()  # -> 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

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

Answer 2

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

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)

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