Trace of inverse-matrix product in scipy

Question

I'm looking to calculate the trace of an inverse-matrix product efficiently, i.e. Tr(A^-1 B). I use the inverse of A in many other places, so I have access to either the Cholesky decomposition or the explicit inverse.

The naive thing to do would be one of the following two:

np.trace(np.dot(invA, B))
np.trace(linalg.cho_solve(choA, B))

However, this seems wasteful, as it computes the full matrix A^-1 B at O(N^3), before only using its diagonal to calculate the trace.

Given the explicit inverse, the O(N^2) solution would be to do:

np.sum(invA.T * B)

Though this requires an explicit inverse, which is undesirable.

I think the ideal way to do it, would be to only calculate the diagonal elements of A^-1 B given the Cholesky decomposition and then simply sum. Is this possible using scipy? Or is there another way of calculating Tr(invA*B) in a numerically stable way, given the Cholesky decomposition?

Answer 1

The einsum suggested by @DSM intuitively seems like it would be the fastest way, as it would only calculate the required terms. Doing this manually in python would be slower and unless there is a specialised routine in numpy/scipy, a mathematical trick or without writing something in a lower level language, I can't see a better way. To check timings, we set up a dummy array,

import numpy as np
from scipy import linalg

invA = np.random.rand(100,100)
B = np.random.rand(100,100)
choA = linalg.cholesky(a)

Trying the naive suggestion,

%timeit np.sum(invA.T * B)
10000 loops, best of 3: 38.5 us per loop

The method using the inverse

%timeit np.sum(invA.T * B)
10000 loops, best of 3: 39.1 us per loop

Seems about the same. Finally using the einsum,

%timeit np.einsum("ij,ji", invA, B)
100000 loops, best of 3: 17.4 us per loop

which seems to be approximately twice as quick.