numpy矢量化的双python for循环

Question

V是（n，p）numpy阵列，通常尺寸为n~10，p~20000

我现在的代码看起来像

A = np.zeros(p)
for i in xrange(n):
    for j in xrange(i+1):
        A += F[i,j] * V[i,:] * V[j,:]

我如何重写这个以避免双循环for python？

Answer 1

虽然Isaac的答案似乎很有希望，因为它删除了那两个嵌套for循环，你必须创建一个中间数组M，它是原始n数组大小的V倍。 Python for循环并不便宜，但内存访问也不是免费的：

n = 10
p = 20000
V = np.random.rand(n, p)
F = np.random.rand(n, n)

def op_code(V, F):
    n, p = V.shape
    A = np.zeros(p)
    for i in xrange(n):
        for j in xrange(i+1):
            A += F[i,j] * V[i,:] * V[j,:]
    return A

def isaac_code(V, F):
    n, p = V.shape
    F = F.copy()
    F[np.triu_indices(n, 1)] = 0
    M = (V.reshape(n, 1, p) * V.reshape(1, n, p)) * F.reshape(n, n, 1)
    return M.sum((0, 1))

如果你现在两个都试骑：

In [20]: np.allclose(isaac_code(V, F), op_code(V, F))
Out[20]: True

In [21]: %timeit op_code(V, F)
100 loops, best of 3: 3.18 ms per loop

In [22]: %timeit isaac_code(V, F)
10 loops, best of 3: 24.3 ms per loop

因此删除for循环会让您花费8倍减速。这不是一件好事......此时你甚至可能想要考虑一个需要3ms进行评估的函数是否需要进一步优化。如果你这样做，可以使用np.einsum：

进行一些改进

def einsum_code(V, F):
    n, p = V.shape
    F = F.copy()
    F[np.triu_indices(n, 1)] = 0
    return np.einsum('ij,ik,jk->k', F, V, V)

现在：

In [23]: np.allclose(einsum_code(V, F), op_code(V, F))
Out[23]: True

In [24]: %timeit einsum_code(V, F)
100 loops, best of 3: 2.53 ms per loop

因此，大约20％的速度提升了代码，这些代码很可能不像你的for循环那么可读。我认为不值得......

Answer 2

关于这一点的困难部分是你只想用j <= i取得元素的总和。如果不是那样，那么你可以做以下事情：

M = (V.reshape(n, 1, p) * V.reshape(1, n, p)) * F.reshape(n, n, 1)
A = M.sum(0).sum(0)

如果F是对称的（如果F[i,j] == F[j,i]），那么您可以利用上面M的对称性，如下所示：

D = M[range(n), range(n)].sum(0)
A = (M.sum(0).sum(0) - D) / 2.0 + D

也就是说，这实际上不是矢量化的理想选择，因为你有n << p所以你的for - 循环不会对这个计算的速度产生太大影响。

编辑：正如比尔在下面所说，您可以确保首先将您不想使用的F元素设置为零，然后M.sum(0).sum(0) 1}}结果将是你想要的。

Answer 3

表达式可以写成

因此您可以使用np.newaxis - construct：

将其加总

na = np.newaxis
X = (np.tri(n)*F)[:,:,na]*V[:,na,:]*V[na,:,:]
X.sum(axis=1).sum(axis=0)

这里构造了一个3D数组X[i,j,p]，然后将2个第一轴相加，得到一维数组A[p]。 Additionaly F乘以三角矩阵以根据问题限制求和。