Question

我有一个k*n矩阵 X 和一个k*k矩阵 A。对于X的每一列，我想计算标量

X[:, i].T.dot(A).dot(X[:, i])

（或在数学上，Xi' * A * Xi）。

目前，我有一个for循环：

out = np.empty((n,))
for i in xrange(n):
    out[i] = X[:, i].T.dot(A).dot(X[:, i])

但是由于n很大，如果可能的话，我想更快地做到这一点（即使用一些 NumPy 函数而不是循环）。

11

6

0

Answer 1

这似乎做得很好：(X.T.dot(A)*X.T).sum(axis=1)

编辑：这有点快。np.einsum('...i,...i->...', X.T.dot(A), X.T)。如果X和A是 Fortran 连续的，则两者都可以更好地工作。

Answer 2

您可以使用numpy.einsum：

np.einsum('ji,jk,ki->i',x,a,x)

这会得到相同的结果,让我们看看它是否快得多:

看起来dot仍然是最快的选择，特别是因为它使用线程 BLAS，而不是在一个内核上运行的einsum。

import perfplot
import numpy as np
def setup(n):
    k = n
    x = np.random.random((k, n))
    A = np.random.random((k, k))
    return x, A
def loop(data):
    x, A = data
    n = x.shape[1]
    out = np.empty(n)
    for i in range(n):
        out[i] = x[:, i].T.dot(A).dot(x[:, i])
    return out
def einsum(data):
    x, A = data
    return np.einsum('ji,jk,ki->i', x, A, x)
def dot(data):
    x, A = data
    return (x.T.dot(A)*x.T).sum(axis=1)
perfplot.show(
    setup=setup,
    kernels=[loop, einsum, dot],
    n_range=[2**k for k in range(10)],
    logx=True,
    logy=True,
    xlabel='n, k'
    )