Numpy/Scipy中的卷积计算

Question

我在做一些计算工作时，发现我的程序有一个瓶颈，主要是一个函数的表现不太好（np 是 numpy，sp 是 scipy）：

def mix1(signal1, signal2):
    spec1 = np.fft.fft(signal1, axis=1)
    spec2 = np.fft.fft(signal2, axis=1)
    return np.fft.ifft(spec1*spec2, axis=1)

这两个信号的形状是 (C, N)，其中 C 是数据集的数量（通常少于20），而 N 是每个数据集中的样本数量（大约5000）。每个数据集（行）的计算完全独立于其他数据集。

我认为这只是一个简单的卷积运算，所以我尝试用以下代码替代它：

def mix2(signal1, signal2):
    outputs = np.empty_like(signal1)

    for idx, row in enumerate(outputs):
        outputs[idx] = sp.signal.convolve(signal1[idx], signal2[idx], mode='same')

    return outputs

...只是想看看结果是否相同。但结果并没有，我的问题是：

1. 为什么会这样？
2. 有没有更好的方法来计算 mix1() 的等效结果？

（我意识到 mix2 可能不会更快，但它可能是并行处理的一个不错起点。）

这是我用来快速检查的完整脚本：

import numpy as np
import scipy as sp
import scipy.signal

N = 4680
C = 6

def mix1(signal1, signal2):
    spec1 = np.fft.fft(signal1, axis=1)
    spec2 = np.fft.fft(signal2, axis=1)
    return np.fft.ifft(spec1*spec2, axis=1)

def mix2(signal1, signal2):
    outputs = np.empty_like(signal1)

    for idx, row in enumerate(outputs):
        outputs[idx] = sp.signal.convolve(signal1[idx], signal2[idx], mode='same')

    return outputs

def test(num, chans):
    sig1 = np.random.randn(chans, num)
    sig2 = np.random.randn(chans, num)
    res1 = mix1(sig1, sig2)
    res2 = mix2(sig1, sig2)

    np.testing.assert_almost_equal(res1, res2)

if __name__ == "__main__":
    np.random.seed(0x1234ABCD)
    test(N, C)

Answer 1

之前提到过，scipy.signal.convolve这个函数并不是用来做循环卷积的。如果你想在实际空间中进行循环卷积（而不是用fft的方法），我建议你使用scipy.ndimage.convolve这个函数。它有一个模式参数，可以设置为'wrap'，这样就可以实现循环卷积了。

for idx, row in enumerate(outputs):
    outputs[idx] = sp.ndimage.convolve(signal1[idx], signal2[idx], mode='wrap')

Answer 2

scipy.signal.fftconvolve 是一个用快速傅里叶变换（FFT）来进行卷积运算的工具，它是用 Python 编写的。你可以查看它的源代码，了解它是怎么工作的，然后可以修正你的 mix1 函数。

Answer 3

我测试了一下，现在可以确认几点：

1) numpy.convolve 不是循环的，这和 FFT 代码给出的结果不同：

2) FFT 在内部并不会自动填充到 2 的幂次方。看看下面这些操作的速度差别就知道了：

x1 = np.random.uniform(size=2**17-1)
x2 = np.random.uniform(size=2**17)

np.fft.fft(x1)
np.fft.fft(x2)

3) 归一化并不是一个区别——如果你简单地做一个循环卷积，通过计算 a(k)*b(i-k) 的和，你会得到 FFT 代码的结果。

问题在于，填充到 2 的幂次方会改变结果。我听说有一些聪明的方法可以通过使用长度的质因数来解决这个问题（在《数值食谱》中提到过，但没有具体代码），但我从来没有见过有人真的这样做过。