NumPy 加速设置二维数组的元素答案

【问题标题】：NumPy speed up setting elements of 2D ArrayNumPy 加速设置二维数组的元素
【发布时间】：2015-01-23 04:58:47
【问题描述】：

我正在尝试在 Python 中有效地索引二维数组，但问题是它真的很慢。

这是我尝试过的（简化示例）：

xSize = veryBigNumber
ySize = veryBigNumber
a = np.ones((xSize,ySize))
N = veryBigNumber  
const = 1

for t in range(N):

  for i in range(xSize):
    for j in range(ySize):
      a[i,j] *= f(i,j)*const   # f(i,j) is an arbitrary function of i and j.

现在我想用更有效的方法替换嵌套循环。我该怎么做？

【问题讨论】：

你需要告诉我们更多关于你想要做什么的信息 - 应该能够避免为一个简单的公式声明一个巨大的数组。如果您确实需要“一个大小接近最大可用 RAM 的数组”，那将是非常有限的，您应该说明该要求。您还可以查看磁盘支持的数据结构。看看熊猫。
@smci：我会试着在一秒钟内把它说得更清楚。你能扭转标题的变化吗？需要改进的不是创作。
@smci: 粗略的，以后会用到！如果没有，上面的代码就没有任何意义，不是吗？

标签： arrays performance python-2.7 numpy vectorization

【解决方案1】：

您的二维数组可以使用以下加法生成：

np.arange(200)[:,np.newaxis] + np.arange(200)

这种类型的矢量化操作可能非常快：

>>> %timeit np.arange(200)[:,np.newaxis] + np.arange(200)
1000 loops, best of 3: 178 µs per loop

这个方法在不只限于加法。我们可以将上述操作中的两个数组作为任意universal function（通常缩写为ufunc）的参数。

例如：

>>> np.multiply(np.arange(5)[:,np.newaxis], np.arange(5))
array([[ 0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4],
       [ 0,  2,  4,  6,  8],
       [ 0,  3,  6,  9, 12],
       [ 0,  4,  8, 12, 16]])

NumPy 为所有基本算术运算和一些更有趣的运算内置了 ufunc。如果您需要更奇特的功能，NumPy allows you to make your own ufunc。

编辑：快速解释此方法中发生的广播；你可以这样想……

np.arange(5) 生成一维数组，如下所示：

array([0, 1, 2, 3, 4])

代码np.arange(5)[:,np.newaxis] 将第二个维度（列）添加到范围，生成这个二维数组：

array([[0],
       [1],
       [2],
       [3],
       [4]])

要使用np.multiply 创建最终的 5x5 数组（尽管我们可以使用任何 ufunc 或二进制算术运算），NumPy 将第二个数组中的 0 与第一个数组中的每个元素相乘，形成一行这个：

[ 0,  0,  0,  0,  0]

然后它将第二个数组中的第二个元素1 与第一个数组相乘，生成这一行：

[ 0,  1,  2,  3,  4]

这一直持续到我们得到最终的 5x5 矩阵。

【讨论】：

我以 'i+j' 为例。它可以是“i”和“j”的任何其他函数。因此，我寻求嵌套循环的通用替代品。
这看起来出奇的好^^。您介意解释一下np.multiply(np.arange(5)[:,np.newaxis], np.arange(5)) 的作用吗？我无法弄清楚。尤其是[:,np.newaxis]这部分！谢谢！

【解决方案2】：

您可以使用indices 例程：

b=np.indices(a.shape)
a=b[0]+b[1]

时间安排：

%%timeit
    ...: b=np.indices(a.shape)
    ...: c=b[0]+b[1]
1000 loops, best of 3: 370 µs per loop


%%timeit
for i in range(200):
  for j in range(200):
     a[i,j] = i + j

100 loops, best of 3: 10.4 ms per loop

【讨论】：

好的，这个解决方案有效。但是，它需要三倍于完成任务所需的内存。我需要的是数组“a”的大小接近最大可用 RAM。然后，我认为这个解决方案行不通。抱歉，因为我没有足够的声誉，所以无法投票赞成您的答案:(

【解决方案3】：

由于您的输出矩阵a 是矩阵F 与元素f_ij = f(i,j) * const 的N 的元素幂，因此您的代码可以简化为

F = np.empty((xSize, ySize))
for i in range(xSize):
    for j in range(ySize):
        F[i,j] = f(i,j) * const

a = F ** n

如果函数f(i,j) 是矢量化的，您可以将F 矩阵的创建换成更高效的方法：

xmap, ymap = numpy.meshgrid(range(xSize), range(ySize))
F = f(xmap, ymap) * const

【讨论】：