在 Numpy 中迭代的更多 Pythonic 方式

Question

我是一名工科学生，我习惯于用 Fortran 编写代码，但现在我正在尝试更多地使用 Python 来使用 Numpy 来编写我的数值配方。

如果我需要使用多个数组中的元素重复执行计算，那么我在 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] 代表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，所以它会与您的原始代码产生的略有不同.

Answer 2

我喜欢列表推导

k = [ x ** y for x, y in zip(some_array, some_other_array) ]

其他人喜欢map

map( lambda x, y : x*y , zip(some_array, some_other_array) )

将两个数组相乘并返回一个列表或生成器。 （当然你还有其他方法可以在 numpy 中完成该特定任务。）如果你想将它转换回数组，你可以这样做

k = array( [ x ** y for x, y in zip(some_array, some_other_array) ] )