More pythonic way to iterate in Numpy
Question: I am an engineering student and I'm accustomed to write code in Fortran, but
now I'm trying to get more into Python for my numerical recipes using Numpy.

If I needed to perform a calculation repeatedly using elements from several
arrays, the immediate translation from what I'd write in Fortran would be

    
    
    k = np.zeros(N, dtype=np.float)
    u = ...
    M = ...
    r = ...
    for i in xrange(N):
      k[i] = ... # Something with u[i], M[i], r[i] and r[i - 1], for example
    

But I was wondering if this way is more pythonic, or preferrable in any way:

    
    
    for i, (k_i, u_i, M_i, r_i) in enumerate(zip(k, u, M, r)):
      k_i = ... # Something with u_i, M_i, r_i and r[i - 1]
    

Thanks to enumerate I have the index, otherwise if I don't need it I could use
just zip or itertools.izip.

Any ideas? How is the code affected in terms of performance? Is there any
other way to accomplish this?


Answer: Almost all numpy operations are performed element-wise. So instead of writing
an explicit loop, try defining `k` using an array-based formula:

    
    
    r_shifted = np.roll(x, shift = 1)
    k = ... # some formula in terms of u, M, r, r_shifted
    

For example, instead of

    
    
    import numpy as np
    
    N=5
    k = np.zeros(N, dtype=np.float)
    u = np.ones(N, dtype=np.float)
    M = np.ones(N, dtype=np.float)
    r = np.ones(N, dtype=np.float)
    for i in xrange(N):
      k[i] = u[i] + M[i] + r[i] + r[i-1]
    print(k)  
    # [ 4.  4.  4.  4.  4.]
    

use:

    
    
    r_shifted = np.roll(r, shift = 1)
    k = u + M + r + r_shifted
    print(k)
    # [ 4.  4.  4.  4.  4.]
    

* * *

[np.roll(r, shift =
1)](http://docs.scipy.org/doc/numpy/reference/generated/numpy.roll.html)
returns a new array of the same size as `r`, with `r_shifted[i] = r[i-1]` for
`i = 0, ..., N-1`.

    
    
    In [31]: x = np.arange(5)
    
    In [32]: x
    Out[32]: array([0, 1, 2, 3, 4])
    
    In [33]: np.roll(x, shift = 1)
    Out[33]: array([4, 0, 1, 2, 3])
    

Making a copy like this requires more memory (of the same size as `r`) but
allows you to do fast numpy operations instead of using a slow Python loop.

* * *

Sometimes the formula for `k` can instead be defined in terms of `r[:-1]` and
`r[1:]`. Note `r[:-1]` and `r[1:]` are slices of `r` and of the same shape. In
this case, you don't need any extra memory since basic slices of `r` are so-
called _views_ of `r`, not copies.

I didn't define `k` this way in the example above because then `k` would have
had length `N-1` instead of `N`, so it would have been slightly different than
what your original code would have produced.

