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 1

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) 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.