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，但是它们都不起作用。

我可以遍历整个轴，将所有内容放入列表中，然后进行串联，但是我试图使用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)

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

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

2 个答案: