Numpy中的向量化 - 广播

Question

我有一段Python代码，里面有以下几个部分：

我有一个叫做 intensities 的向量，它大概是这样的：

array([ 1142.,  1192.,  1048., ...,    29.,    18.,    35.])

我还有一个叫做 x 的向量，它看起来是这样的：

array([   0,    1,    1, ..., 1060, 1060, 1061])

接着，我有一个for循环，用来填充另一个向量 radialDistribution，代码是这样的：

for i in range(1000):
    radialDistribution[i] = sum(intensities[np.where(x == i)]) / len(np.where(x == i)[0])

问题是，这个过程需要20秒才能完成……所以我想要把它变得更快，也就是进行向量化。但我对Numpy中的广播功能还不太熟悉，网上也没找到太多相关的信息……所以我需要你们的帮助。

我尝试过这个方法，但没有成功：

i= np.ogrid[:1000]
intensities[i] = sum(sortedIntensities1D[np.where(sortedDists1D == i)]) / len(np.where(sortedDists1D == i)[0])

你能告诉我应该去哪里学习Numpy的向量化操作吗？

非常感谢你们的宝贵帮助！

Answer 1

在我之前提到的关于 itertools.groupby 的解决方案基础上，这里有一个适用于两个小数组的解决方案。

import numpy as np
import itertools
intensities = np.arange(12,dtype=float)
x=np.array([1,0,1,2,2,1,0,0,1,2,1,0]) # general, not sorted or consecutive

首先是一个调整过的 bincount 解决方案，适用于不连续的值。

# using bincount
# if 'x' are not consecutive
J=np.bincount(x)>0
print np.bincount(x,weights=intensities)[J]/np.bincount(x)[J]

接下来是一个 groupby 的解决方案。

# using groupby;
# sort if need
I=np.argsort(x)
x=x[I]
intensities=intensities[I]

# make a record array for use by groupby
xi=np.zeros(shape=x.shape, dtype=[('intensities',float),('x',int)])
xi['intensities']=intensities
xi['x']=x

g=itertools.groupby(xi, lambda z:z['x'])
xx=np.array([np.array([z[0] for z in y[1]]).mean() for y in g])
print xx

这里有一个简洁的 numpy 解决方案，使用了 np.unique 的 return_index 选项和 np.split。注意，x 需要先排序。对于大数组，我对速度不太乐观，因为在 unique 和 split 中会有额外的迭代，除了列表推导外。

[values, index] = np.unique(x, return_index=True)
[y.mean() for y in np.split(intensities, index[1:])]

Answer 2

这里有一种使用广播的方法：

# arrays need to be at least 2D for broadcasting
x = np.atleast_2d(x)

# create vector of indices
i = np.atleast_2d(np.arange(x.size))

# do the vectorized calculation
bool_eq = (x == i.T)
totals = np.sum(np.where(bool_eq, intensities, 0), axis=1)
rD = totals / np.sum(bool_eq, axis=1)

这个方法用了两次广播：一次是在操作 x == i.T 时，另一次是在调用 np.where 时。不幸的是，上面的代码运行得非常慢，甚至比最初的代码还慢。主要的问题出在 np.where 上，不过我们可以通过对布尔数组和强度值进行乘法运算来加速这个过程（同样是通过广播）：

totals = np.sum(bool_eq*intensities, axis=1)

这实际上和矩阵与向量的乘法是一样的，所以我们可以写成：

totals = np.dot(intensities, bool_eq.T)

最终的结果是比原来的代码快（至少在中间数组的内存使用成为限制因素之前），但根据其他答案的建议，你可能还是用迭代的方法更好。

编辑：使用 np.einsum 的方法在我的测试中更快：

totals = np.einsum('ij,j', bool_eq, intensities)

Answer 3

这里有我在numpy中实现的group_by功能。它的概念上和pandas的解决方案类似；不过这个不需要用到pandas，我觉得它应该成为numpy的核心部分。

使用这个功能，你的代码看起来会像这样：

radialDistribution = group_by(x).mean(intensities)

而且执行起来非常快。

另外，看看最后定义的test_radial函数，它可能会更接近你的最终目标。

Answer 4

如果你的 x 向量是从0开始的连续整数，那么你可以直接这样做：

radialDistribution = np.bincount(x, weights=intensities) / np.bincount(x)