计算矢量v矩阵的"v ^ TA v"

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函数而不是循环).

Answer 1

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

编辑:这快一点.np.einsum('...i,...i->...', X.T.dot(A), X.T).如果X和AFortran连续,两者都能更好地工作.

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'
    )