python - 使用 numpy 和 fft 在 for 循环中进行多核处理

Question

我使用 numpy 和 fft 计算向量。 我使用了 numpy 广播方法和 for 循环。 两种方法的速度相似。 如何使用多核、numpy 和 fft 计算向量？

import numpy as np
from numpy.fft import fft, ifft

num_row, num_col = 6000, 13572

ss = np.ones((num_row, num_col), dtype=np.complex128)
sig = np.random.standard_normal(num_col) * 1j * np.random.standard_normal(num_col)

# for loop    
for idx in range(num_row):
    ss[idx, :] = ifft(fft(ss[idx, :]) * sig)

# broadcast
ss = ifft(fft(ss, axis=1) * sig, axis=1)

结果

loop : 10.798867464065552 sec
broadcast : 11.298897981643677 sec

Answer 1

您可以将 axis 参数用于 fft 和 ifft 以及广播:

ss = np.ones((num_row, num_col), dtype=np.complex128)
sig = np.random.standard_normal(num_col) * 1j * np.random.standard_normal(num_col)

ss = ifft(fft(ss, axis=1) * sig, axis=1)