如何使用xarray.apply_ufunc（）将1D中值滤镜应用于3D DataArray

Question

我有3维DataArray（使用xarray）。 我想沿特定尺寸对其应用一维尺寸。 具体来说，我想应用scipy.signal.medfilt()函数，但同样，它应该是一维的。

到目前为止，我已经通过以下方式成功实现了这一点：

for sample in data_raw.coords["sample"]:
    for experiment in data_raw.coords["experiment"]:
        data_filtered.loc[sample,experiment,:] = signal.medfilt(data_raw.loc[sample,experiment,:], 15)

（我的数据数组的维度为“样本”，“实验”和“ wave_number。此代码沿“ wave_number”维度应用过滤器）

问题在于计算需要花费很长时间，而且我的直觉告诉我，尽管像这样的坐标循环是一种低效的方法。 所以我正在考虑使用xarray.apply_ufunc()函数，尤其是因为我在相同的代码中以类似的方式使用了它：

xr.apply_ufunc(np.linalg.norm, data, kwargs={"axis": 2}, input_core_dims=[["wave_number"]])

（这将计算沿“ wave_number”维的矢量长度。）

我最初也像在这里的第一个代码一样在坐标中循环。

问题是当我尝试

xr.apply_ufunc(signal.medfilt, data_smooth, kwargs={"kernel_size": 15})

它返回一个全为零的数据数组，大概是因为它应用了3D中值滤波器并且该数据数组包含NaN条目。 我意识到这里的问题是我需要将scipy.signal.medfilt()函数提供给一维数组，但是不幸的是，无法指定沿其应用滤镜的轴（与numpy.linalg.norm()不同） 。

那么，如何在不循环通过坐标的情况下应用一维中值滤波器？

Answer 1

如果我理解正确，则应该这样使用它：

xr.apply_ufunc(signal.medfilt, data_smooth, kwargs={"kernel_size": 15}, input_core_dims = [['wave_number']], vectorize=True)

使用vectorize = True可以对输入函数进行矢量化，以应用于定义为保留核心尺寸的数组切片。

尽管如此，如文档中所述：

存在此选项是为了方便起见，但它总是比提供预向量化功能慢

因为实现本质上是一个for循环。 但是，我仍然比通过自己制作循环来获得更快的结果。