Search code examples
pythonnumpymatrix

Applying matrix multiplication in chain to several sets of values

Mathematically, I'm trying to calculate x^T A x, where x is an n-dimensional coordinate and A a n-dimensional square matrix. However, I'd like to efficiently calculate this for a set of coordinates. For example, in two dimensions:

import numpy as np
x = np.column_stack([[1,2,3,4,5],[6,7,8,9,0]])
A = np.array([[1,0],[0,2]])

print(x[0] @ A @ x[0])  # works
# How can I get efficiently an array of x[i] @ A @ x[i]?
y = [x[i] @ A @ x[i] for i in range(x.shape[0])]
Solution

  • Is this what you want?

    >>> (x @ A @ x.T).diagonal()
array([ 73, 102, 137, 178,  25])

    Or maybe:

    >>> (x @ A * x).sum(axis=1)
array([ 73, 102, 137, 178,  25])

    Good einsum suggestion from hpaulj in comments:

    >>> np.einsum('ij,jk,ik->i', x, A, x)
array([ 73, 102, 137, 178,  25])