对 numpy 数组的所有值执行操作，引用 i 和 j答案

【问题标题】：Performing operations on all values of a numpy array, referencing i and j对 numpy 数组的所有值执行操作，引用 i 和 j
【发布时间】：2012-06-04 23:52:21
【问题描述】：

我试图通过对二维数组应用操作来提高 numpy 性能，问题是数组中每个元素的值取决于该元素的 i,j 位置。

显然，最简单的方法是使用嵌套的 for 循环，但我想知道是否有更好的方法通过引用 np.indices 或类似的东西？这是我的“愚蠢”代码：

for J in range(1025):
    for I in range(1025):
        PSI[I][J] = A*math.sin((float(I+1)-.5)*DI)*math.sin((float(J+1)-.5)*DJ)
        P[I][J] = PCF*(math.cos(2.*float(I)*DI)+math.cos(2.*float(J)*DJ))+50000.

【问题讨论】：

没有人明确解决这个问题，但是在索引 numpy 数组的地方使用 python for 循环正在失去用 C 编写 numpy 的速度优势。

标签： python arrays performance numpy indexing

【解决方案1】：

由于您在两个数组之间进行乘法运算，因此您可以在使用 arange 获取 sin/cos 的数组之后使用 outer 函数。

类似这样的东西（使用 numpy 的三角函数，因为它们是矢量化的）

PSI_i = numpy.sin((arange(1,1026)-0.5)*DI)
PSI_j = numpy.sin((arange(1,1026)-0.5)*DJ)
PSI = A*outer(PSI_i, PSI_j)

P_i = numpy.cos(2.*arange(1,1026)*DI)
P_j = numpy.cos(2.*arange(1,1026)*DJ)
P = PCF*outer(P_i, P_j) + 50000

如果您的环境是使用from numpy import * 或from pylab import * 设置的，那么在您的三角函数之前您不需要那些numpy. 前缀。我将它们保留在其中以将它们与 math 区分开来，这不适用于这种方法。

【讨论】：

对于这个特殊的问题，我认为这是最好的解决方案（+1）
这是正确的方法，但如果您使用的是numpy.sin，您可能应该使用numpy.outer。
@tillsten - 不管它有什么价值，您都可以使用outer 进行减法、除法、幂等。例如。 numpy.subtract.outer(a, b)、numpy.power.outer(a, b)等
@tillsten - 它只适用于ufuncs，所以它仍然是有限的（广播更普遍）。不过，这是一个方便的技巧。很高兴你发现它有用！

【解决方案2】：

您可以通过索引获取索引值的网格：

I,J=np.indices(PSI.shape)
#All constants set to one
PSI2=np.sin(I+1-.5)*np.sin(J+1-.5)
print PSI-PSI2 # should be zero.

我用 ipython 做了一些计时：

import numpy as np
import math
A = 1
P = 1
DI = 1
DJ = 1

def a():
    PSI=np.zeros((1025,1025))
    for J in range(1025):
        for I in range(1025):
            PSI[I][J] = A*math.sin((float(I+1)-.5)*DI)*math.sin((float(J+1)-.5)*DJ)
%timeit a()

def b():
    PSI=np.zeros((1025,1025))
    for I,J in np.ndindex(*PSI.shape):
        PSI[I,J] = A*math.sin((float(I+1)-.5)*DI)*math.sin((float(J+1)-.5)*DJ)        
%timeit b()

def c():
    I,J=np.indices((1025, 1025))
    P2=A*np.sin((I+1-.5)*DI)*np.sin((J+1-.5)*DJ)    
%timeit c()

def d():
    PSI_i = np.sin((np.arange(1,1026)-0.5)*DI)
    PSI_j = np.sin((np.arange(1,1026)-0.5)*DJ)
    PSI = A*np.outer(PSI_i, PSI_j)    
%timeit d()

结果在我的机器上一点也不奇怪：

1 loops, best of 3: 1.75 s per loop
1 loops, best of 3: 3.51 s per loop
10 loops, best of 3: 77.1 ms per loop
100 loops, best of 3: 7.16 ms per loop

【讨论】：

您实际上并没有调用这些函数。你需要%timeit a()等，而不是%timeit a。
正确计时，我得到 np.outer 的解决方案大约需要 0.009 秒，索引的解决方案大约需要 0.12 秒，嵌套循环需要 1.4 秒，np.ndindex 需要高达 3.28 秒.我想知道为什么np.ndindex 这么慢——我原以为它会出现在其他解决方案之间，但无论如何我已经删除了我的答案；）。（我想这说明了我们分析的原因）。
@mgilson，ndindex 这么慢的原因是它返回了一个 python 迭代器。 a = np.ndindex((10, 10)); print a.next(); print a.next().
@senderle -- 很有趣。我以为它会返回一个生成器，但我想这就是我没有仔细阅读文档的结果。

【解决方案3】：

试试 numpy 的 ndenumerate 函数，它返回值和索引：

>>> a
array([[5, 5, 5],
       [1, 2, 3]])


>>> for index, value in numpy.ndenumerate(a):
...     print index, value

(0, 0) 5
(0, 1) 5
(0, 2) 5
(1, 0) 1
(1, 1) 2
(1, 2) 3

【讨论】：

这就是为体面大小的 Python 列表做的方式。但是在使用 NumPy 时，这是要避免的，既不破坏未装箱+打包数据和 C 函数的性能优势，又因为您通常将它们用于 huge 数量的数据（OP 的示例已经一百万个元素）。