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

你所展示的是'pythonic',因为它使用Python列表和迭代方法.唯一的用途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个数组的方程式.这在很大程度上取决于您应该学习的广播规则.

如果函数只接受标量输入,我们将不得不使用某种形式的迭代.这是一个经常提出的问题; 寻找[numpy] vectorize.

Answer 2

我想说这是执行迭代的最直接且普遍理解的方式。

另一种方法是迭代这些值并为给定的 (i, j) 对调用函数

import itertools

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

def do_work(i, j):
       M[i,j] = 1/((i+2*j+1)**2)

[do_work(i, j) for (i, j) in itertools.product(xrange(N), xrange(N))]

在这里，我只是使用 itertools.product 为可能的 (i, j) 值创建一个生成器，您也可以使用 for 循环。

for (i, j) in itertools.product(xrange(N), xrange(N)):
    M[i,j] = 1/((i+2*j+1)**2)