python - 在 numpy 中通过数组嵌套 for 循环的更好替代方案？

Question

我经常需要遍历数组并对每个条目执行一些操作，其中操作可能取决于索引和条目的值。这是一个简单的例子。

import numpy as np

N=10
M = np.zeros((N,N))

for i in range(N):
    for j in range(N):
        M[i,j] = 1/((i+2*j+1)**2)

是否有更短、更清晰或更 Pythonic 的方式来执行此类任务？

Answer 1

从使用 Python 列表和迭代方法的意义上说，您展示的是“pythonic”。 numpy 的唯一用途是分配值，M{i,j] =。列表不采用那种索引。

要充分利用 numpy，创建索引网格或数组，并一次计算所有值，无需显式循环。例如，在您的情况下:

In [333]: N=10
In [334]: I,J = np.ogrid[0:10,0:10]
In [335]: I
Out[335]: 
array([[0],
       [1],
       [2],
       [3],
       [4],
       [5],
       [6],
       [7],
       [8],
       [9]])
In [336]: J
Out[336]: array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
In [337]: M = 1/((I + 2*J + 1)**2)
In [338]: M
Out[338]: 
array([[ 1.        ,  0.11111111,  0.04      ,  0.02040816,  0.01234568,
         0.00826446,  0.00591716,  0.00444444,  0.00346021,  0.00277008],
 ...
       [ 0.01      ,  0.00694444,  0.00510204,  0.00390625,  0.00308642,
         0.0025    ,  0.00206612,  0.00173611,  0.00147929,  0.00127551]])

ogrid 是构建可以一起“广播”的数组集的几种方法之一。 meshgrid 是另一个常用函数。

在您的情况下，方程式适用于 2 个像这样的数组。这在很大程度上取决于您应该研究的广播规则。

如果函数只接受标量输入，我们将不得不使用某种形式的迭代。这是一个常见的 SO 问题；搜索 [numpy] vectorize。