scipy.ndimage.fourier_filter似乎只适用于索引为0,0，在数组中的向量

Question

我正在尝试在3D阵列的第3轴上实现移动平均值。由于移动窗口可能非常大，我想在傅立叶域中进行以获得不错的速度。但它似乎只适用于arr[0,0,:]中的向量。

这是一个包含测试用例的最小示例。第一个断言通过，第二个断言没有通过。

import numpy as np
import scipy.ndimage as sni

def moving_average_in_ft(arr, size, axis=-1):
    """Perform moving average of given size along given axis of arr
    Performed in Fourier domain.
    Returns a complex result.
    """
    arr_ft = np.fft.fft(arr)
    nd_size = [1,] * arr.ndim # all other axes do not average
    nd_size[axis] = size
    return np.fft.ifft(sni.fourier_uniform(arr_ft, nd_size))


def test_moving_average_in_ft_and_time():
    a = np.random.rand(2, 2, 20)
    size = 9
    a_filt_ft = moving_average_in_ft(a, size).real
    a_filt_time = sni.uniform_filter(a, (1, 1, size), mode='wrap')
    # for the first ensemble it works
    ensamble_1 = (0, 0, slice(None))
    np.testing.assert_allclose(a_filt_ft[ensamble_1], a_filt_time[ensamble_1], atol=0.01)
    # but all the other ensembles don't
    np.testing.assert_allclose(a_filt_ft, a_filt_time, atol=0.1)