Python numpy.fft 改变步幅

Question

亲爱的 stackoverflow 社区!

今天我发现在高端集群架构上，尺寸为 1921 x 512 x 512 的 2 个立方体的元素乘法大约需要 27 秒。这太长了，因为在当前实现中，我必须执行至少 256 次这样的计算才能对功率谱进行方位平均。我发现性能缓慢主要是由于不同的步幅结构(一种情况是 C，另一种情况是 FORTRAN)。两个数组之一是新生成的 bool 网格(C 顺序)，另一个(FORTRAN 顺序)来自 3D numpy.fft.fftn()输入网格的傅立叶变换(C 阶)。除了反转轴(这只是一个解决方法)之外，numpy.fft.fftn() 改变步幅和如何防止这种情况的想法有什么原因吗？通过类似的步幅(FT 网格的 ndarray.copy())，可以实现约 4 秒，这是一个巨大的进步。

因此问题如下:

考虑数组:

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(ran,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