Python numpy.fft 改变步伐

Question

亲爱的 Whosebug 社区！

今天我发现在高端集群架构上，2 个尺寸为 1921 x 512 x 512 的立方体的元素乘法需要大约 27 秒。这太长了，因为在当前的实现中，我必须为功率谱的方位角平均执行至少 256 次这样的计算。我发现性能缓慢主要是由于步幅结构不同（一种情况下是 C，另一种情况下是 FORTRAN）。两个数组之一是新生成的布尔网格（C阶），另一个（FORTRAN阶）来自FT网格的3Dnumpy.fft.fftn() Fourier transform of an input grid (C order). Any reasons why numpy.fft.fftn() changes the strides and ideas on how to prevent that except for reversing the axes (which would be just a workaround)? With similar strides (ndarray.copy()），~4s是可以实现的，一个巨大的进步。

因此问题如下：

考虑数组：

ran = np.random.rand(1921, 512, 512)
ran.strides
(2097152, 4096, 8)

a = np.fft.fftn(ran)
a.strides
(16, 30736, 15736832)

正如我们所见，步幅结构不同。如何避免这种情况（不使用 a = np.fft.fftn(运行, axes = (1,0))）？是否还有其他可能影响步幅结构的 numpy 数组例程？在这些情况下可以做什么？

像往常一样非常感谢有用的建议！

Answer 1

您可以使用 scipy.fftpack.fftn（同样由 hpaulj 建议）而不是 numpy.fft.fftn，看起来它正在做您想要的。然而，它的表现略差：

import numpy as np
import scipy.fftpack

ran = np.random.rand(192, 51, 51)  # not much memory on my laptop
a = np.fft.fftn(ran)
b = scipy.fftpack.fftn(ran)

ran.strides
(20808, 408, 8)
a.strides
(16, 3072, 156672)
b.strides
(41616, 816, 16)

timeit -n 100 np.fft.fftn(ran)
100 loops, best of 3: 37.3 ms per loop
timeit -n 100 scipy.fftpack.fftn(ran)
100 loops, best of 3: 41.3 ms per loop

Answer 2

Any reasons why numpy.fft.fftn() changes the strides and ideas on how to prevent that except for reversing the axes (which would be just a workaround)?

计算数组的多维 DFT 包括连续计算每个维度上的一维 DTF。有两种策略：

将一维 DTF 计算限制为连续的一维数组。由于数组是连续的，与 latency/cache 未命中相关的问题将会减少。这种策略有一个主要缺点：数组要在每个维度之间进行转置。这很可能是 numpy.fft 采用的策略。在计算结束时，数组已被转置。为了避免不必要的计算，返回转置数组并修改步幅。
为跨步阵列启用一维 DDFT 计算。这可能会引发一些与延迟相关的问题。它是 fftw 的策略，可通过接口 pyfftw 使用。因此，输出数组具有与输入数组相同的步幅。

计时numpy.fftn和pyfftw.numpy.fftn执行 and there or there会告诉你FFTW是否真的是西方最快的傅立叶变换...

要检查 numpy 是否使用第一种策略，请查看 numpy/fft/fftpack.py。在第 81-85 行，对 work_function(a, wsave) 的调用（即来自 FFTPACK, arguments documented there 的 fftpack.cfftf）包含在对 numpy.swapaxes() 执行换位的调用之间。
scipy.fftpack.fftn似乎没有改变步幅...不过，它似乎使用了第一种策略。 scipy.fftpack.fftn() calls scipy.fftpack.zfftnd() which calls zfft(), based on zfftf1 which does not seem to handle strided DFTs. Moreover, zfftnd() calls many times the function flatten() 执行转置。
再举个例子：对于并行分布式内存多维DFT，FFTW-MPI uses the first strategy to avoid any MPI communications between processes during 1D DTFs. Of course, functions to transpose the array已经不远了，过程中涉及到很多MPI通信

Are there any other numpy array routines that could affect stride structure? What can be done in those cases?

可以search the github repository of numpy for swapaxes: 这个功能只用过几次。因此，在我看来，这个 "change of strides" 是 fft.fftn() 特有的，并且大多数 numpy 函数保持步幅不变。

最后，"change of strides" 是第一种策略的一个特征，没有办法阻止它。唯一的解决方法是在计算结束时交换轴，这很昂贵。但是您可以依赖 pyfftw，因为 fftw 以非常有效的方式实施了第二个策略。 DFT计算会更快，如果不同数组的步幅一致，后续计算也会更快。

Python numpy.fft 改变步伐

Python numpy.fft changes strides

python

arrays

memory-management

numpy

fft