更多pythonic方式迭代Numpy

Question

我是一名工科学生,我习惯于在Fortran中编写代码,但现在我正在尝试使用Numpy为Python编写更多的数据.

如果我需要使用来自多个数组的元素重复执行计算,那么我在Fortran中编写的内容将立即进行翻译

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

但我想知道这种方式是否更加pythonic,或以任何方式优先:

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]

感谢枚举我有索引,否则如果我不需要它我只能使用zip或itertools.izip.

有任何想法吗？代码如何在性能方面受到影响？有没有其他方法可以实现这一目标？

Answer 1

几乎所有的numpy操作都是按元素执行的.因此,不要编写显式循环,而是尝试k使用基于数组的公式进行定义:

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

例如,而不是

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

使用:

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)返回一个r与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])

制作这样的副本需要更多内存(大小相同r),但允许您进行快速numpy操作,而不是使用慢速Python循环.

有时为式k罐代替来定义r[:-1]和r[1:].注意r[:-1]并且r[1:]是r相同形状的切片.在这种情况下,你不需要任何额外的内存,因为基本的片r是所谓的意见的r,不是复印件.

我没有k在上面的示例中定义这种方式,因为k那样会有长度N-1而不是N,所以它会与原始代码产生的略有不同.