如何加快嵌入了numpy函数的python代码的速度？

Question

这是我代码中的速率限制功能

def timepropagate(wv1, ham11,
                  ham12, ham22, scalararray, nt):
    wv2 = np.zeros((nx, ny), 'c16')
    fw1 = np.zeros((nx, ny), 'c16')
    fw2 = np.zeros((nx, ny), 'c16')
    for t in range(0, nt, 1):
        wv1, wv2 = scalararray*wv1, scalararray*wv2
        fw1, fw2 = (np.fft.fft2(wv1), np.fft.fft2(wv2))
        fw1 = ham11*fw1+ham12*fw2
        fw2 = ham12*fw1+ham22*fw2
        wv1, wv2 = (np.fft.ifft2(fw1), np.fft.ifft2(fw2))
        wv1, wv2 = scalararray*wv1, scalararray*wv2
    del(fw1)
    del(fw2)
    return np.array([wv1, wv2])

我需要做的是找到一个相当快的实现，它将使我的速度提高两倍，最好是最快。

我感兴趣的更笼统的问题是，如何使用最少的返回python的连接来加快这一步的速度。 我假设即使我加快了代码的特定段（例如标量数组乘法）的速度，在傅立叶变换时我仍然会从python回来并走，这会花费一些时间。 我有什么方法可以使用，例如numba或cython，而不会在循环中间使此“返回”到python？ 就个人而言，考虑到我已经在使用其他线程，因此我希望在单个线程上快速运行。

编辑：这是分析的结果，第一个用于4096x4096数组的时间为10个时间步，我需要将其放大为nt = 8000。

ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    1    0.099    0.099  432.556  432.556 <string>:1(<module>)
   40    0.031    0.001   28.792    0.720 fftpack.py:100(fft)
   40   45.867    1.147   68.055    1.701 fftpack.py:195(ifft)
   80    0.236    0.003   47.647    0.596 fftpack.py:46(_raw_fft)
   40    0.102    0.003    1.260    0.032 fftpack.py:598(_cook_nd_args)
   40    1.615    0.040   99.774    2.494 fftpack.py:617(_raw_fftnd)
   20    0.225    0.011   29.739    1.487 fftpack.py:819(fft2)
   20    2.252    0.113   72.512    3.626 fftpack.py:908(ifft2)
   80    0.000    0.000    0.000    0.000 fftpack.py:93(_unitary)
   40    0.631    0.016    0.820    0.021 fromnumeric.py:43(_wrapit)
   80    0.009    0.000    0.009    0.000 fromnumeric.py:457(swapaxes)
   40    0.338    0.008    1.158    0.029 fromnumeric.py:56(take)
  200    0.064    0.000    0.219    0.001 numeric.py:414(asarray)
    1  329.728  329.728  432.458  432.458 profiling.py:86(timepropagate)
    1    0.036    0.036  432.592  432.592 {built-in method builtins.exec}
   40    0.001    0.000    0.001    0.000 {built-in method builtins.getattr}
  120    0.000    0.000    0.000    0.000 {built-in method builtins.len}
  241    3.930    0.016    3.930    0.016 {built-in method numpy.core.multiarray.array}
    3    0.000    0.000    0.000    0.000 {built-in method numpy.core.multiarray.zeros}
   40   18.861    0.472   18.861    0.472 {built-in method numpy.fft.fftpack_lite.cfftb}
   40   28.539    0.713   28.539    0.713 {built-in method numpy.fft.fftpack_lite.cfftf}
    1    0.000    0.000    0.000    0.000 {built-in method numpy.fft.fftpack_lite.cffti}
   80    0.000    0.000    0.000    0.000 {method 'append' of 'list' objects}
   40    0.006    0.000    0.006    0.000 {method 'astype' of 'numpy.ndarray' objects}
    1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}
   80    0.000    0.000    0.000    0.000 {method 'pop' of 'list' objects}
   40    0.000    0.000    0.000    0.000 {method 'reverse' of 'list' objects}
   80    0.000    0.000    0.000    0.000 {method 'setdefault' of 'dict' objects}
   80    0.001    0.000    0.001    0.000 {method 'swapaxes' of 'numpy.ndarray' objects}
   40    0.022    0.001    0.022    0.001 {method 'take' of 'numpy.ndarray' objects}

我想我第一次做错了，使用time.time（）计算小型数组的时间差，并推断大型数组的结论。

Answer 1

如果大部分时间都花在了汉密尔顿乘法上，那么您可能想在那部分应用numba。 从NumPy中评估表达式时，最大的好处来自于删除所有需要的时间数组。

还请记住，阵列（4096、4096，c16）足够大，无法舒适地容纳在处理器缓存中。 单个矩阵将占用256 MiB。 因此，认为性能不太可能与操作完全相关，而与带宽相关。 因此，以仅对输入操作数执行一次传递的方式实现这些操作。 在numba中实现确实很简单。 注意：您只需要在numba中实现汉密尔顿表达式。

我还要指出，使用np.zeros的“预分配”似乎表明您的代码未遵循您的意图，例如：

    fw1 = ham11*fw1+ham12*fw2
    fw2 = ham12*fw1+ham22*fw2

实际上将为fw1，fw2创建新的数组。 如果您打算重用缓冲区，则可能要使用“ fw1 [：，：] = ...”。 否则，np.zeros只会浪费时间和内存。

您可能需要考虑将（wv1，wv2）合并为（2，4096，4096，c16）数组。 与（fw1，fw2）相同。 这样一来，您就可以依靠广播来处理“ scalararray”产品，从而使代码更加简单。 fft2和ifft2实际上会做正确的事（AFAIK）。