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

【问题标题】：Applying scipy.ndimage.convolve to tridimensional xarray DataArray将 scipy.ndimage.convolve 应用于三维 xarray DataArray
【发布时间】：2018-09-28 14:38:00
【问题描述】：

我有一个尺寸为 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

【问题讨论】：

标签： python scipy python-xarray

【解决方案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

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

【讨论】：

我觉得很愚蠢，我没有想到添加长度一维。谢谢！

【解决方案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)

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

【讨论】：