时间戳

首页 > 解决方案 > 将 scipy.ndimage.convolve 应用于三维 xarray DataArray

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

问题描述

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

标签： pythonscipypython-xarray

解决方案

使用形状为 (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

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

推荐阅读