更加Pythonic的Numpy迭代方式

Question

更加Pythonic的Numpy迭代方式

3

我是一名工程学生，习惯使用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

但是我想知道这种方式是否更符合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。

有什么想法吗？在性能方面，代码会受到影响吗？还有其他完成此任务的方法吗？

- astrojuanlu

2个回答

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

- Adrian Ratnapala

事实上，使用这种方法无法恢复先前的元素，例如（这就是我使用enumerate的原因，如所示） - astrojuanlu

非常正确。当我需要对序列执行复杂操作时，有时会编写一个独立的函数生成器，显式地产生结果。然后，我的调用代码可以将其打包成一个NumPy数组（如果它想要的话）。 - Adrian Ratnapala

-1，因为在numpy中首选的方式是逐元素操作。请参见unutbu的答案。 - bmu

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- unutbu · Accepted Answer

几乎所有的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 大小相同的新数组，对于 i = 0, ..., N-1，r_shifted[i] = r[i-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 循环。

有时可以使用 r[:-1] 和 r[1:] 来定义 k 的公式。注意，r[:-1] 和 r[1:] 是 r 的切片，形状相同。在这种情况下，您不需要任何额外的内存，因为 r 的基本切片是所谓的视图，而不是副本。

我在上面的示例中没有以这种方式定义 k，因为那样的话，k 的长度将为 N-1 而不是 N，所以它会略微不同于您原来的代码产生的结果。