How do I speed up a piece of python code which has a numpy function embedded in it?
<p>Here is the rate limiting function in my code   </p>

<pre><code>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])
</code></pre>

<p>What I would need to do is find a reasonably fast implementation that would allow me to go at twice the speed, preferably the fastest.</p>

<p>The more general question I'm interested in, is what way can I speed up this piece, using minimal possible connections back to python. I assume that even if I speed up specific segments of the code, say the scalar array multiplications, I would still come back and go from python at the Fourier transforms which would take time. Are there any ways I can use, say numba or cython and not make this "coming back" to python in the middle of the loops?
On a personal note, I'd prefer something fast on a single thread considering that I'd be using my other threads already.</p>

<p>Edit: here are results of profiling, the 1st one for 4096x4096 arrays for 10 time steps, I need to scale it up for nt = 8000. </p>

<pre><code>ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    1    0.099    0.099  432.556  432.556 &lt;string&gt;:1(&lt;module&gt;)
   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}
</code></pre>

<p>I think I've done it wrong the first time, using time.time() to calculate time differences for small arrays and extrapolating the conclusions for larger ones.</p>

<p>If most of the time is spent in the hamiltonian multiplication, you may want to apply numba on that part. The most benefit coming from removing all the temporal arrays that would be needed if evaluating expressions from within NumPy.</p>

<p>Bear also in mind that the arrays (4096, 4096, c16) are big enough to not fit comfortably in the processor caches. A single matrix would take 256 MiB. So think that the performance is unlikely to be related at all with the operations, but rather on the bandwidth. So implement those operations in a way that you only perform one pass in the input operands. This is really trivial to implement in <em>numba</em>. Note: You will only need to implement in <em>numba</em> the hamiltonian expressions.</p>

<p>I want also to point out that the "preallocations" using np.zeros seems to signal that your code is not following your intent as:</p>

<pre><code>    fw1 = ham11*fw1+ham12*fw2
    fw2 = ham12*fw1+ham22*fw2
</code></pre>

<p>will actually create new arrays for fw1, fw2. If your intent was to reuse the buffer, you may want to use "fw1[:,:] = ...". Otherwise the np.zeros do nothing but waste time and memory.</p>

<p>You may want to consider to join (wv1, wv2) into a (2, 4096, 4096, c16) array. The same with (fw1, fw2). That way code will be simpler as you can rely on broadcasting to handle the "scalararray" product. fft2 and ifft2 will actually do the right thing (AFAIK).</p>

python performance numpy cython numba