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，而不是让这个＆＃34;回来＆＃34;在循环中间的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的“preallocations”似乎表示您的代码没有遵循您的意图：

    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）。

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

1 个答案: