更Pythonic的Numpy迭代方法

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

不过我在想，这样写是否更符合Python的风格，或者在某种情况下更好：

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

我喜欢列表推导式

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

Answer 2

几乎所有的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，所以结果会和你原来的代码稍微有点不同。