title: How do I speed up a piece of python code which has a numpy function embedded in it?
body: <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>