python - 在 Numpy 中迭代的更多 pythonic 方式

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]

感谢 enumerate 我有索引，否则如果我不需要它我可以只使用 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 的所谓 View ，而不是副本。

在上面的例子中我没有这样定义 k 因为 k 的长度是 N-1 而不是 N，因此它与您的原始代码生成的内容略有不同。