# Compute $tr(A^TBC)$ in Python

I have to compute the trace of the product between three matrices $$A,B,C$$ in python, i.e. I have to compute $$tr(A^TBC)$$ and I was wondering what was a good way to do it in Python(here $$A^T$$ is the transposed of $$A$$). If I had just 2 matrices $$A, B$$ and I wanted to compute $$tr(A^TB)$$ I know I could do something as

numpy.tensordot(A,B,axis=2)

But what about the case with 3 matrices?

• Welcome to scicomp! How large are your matrices? Is this performance critical, or are you asking for an elegant algorithm? What do you know about the three matrices? Are they all square, are they sparse, real/complex? The order in which you multiply the inner matrices is associative, so you can look at whether you save any execution time switching the order of calculation around. Oct 28, 2020 at 13:19
• Thank you. The matrices are $n\times k$ where typically $n>>k$ and could be big I think. They contain real numbers but they don't have any particular structure. Then I guess it is correct to do numpy.tensordot(A, B@C,axis=2) right? Oct 28, 2020 at 14:11
• Sorry, I meant $A$ and $B$ are $n\times k$ and $B$ is $n\times n$. Oct 28, 2020 at 14:23
• I'm sorry. If A and B are $n \times k$, B can't at the same time be $n \times n$. You have me confused:-) Oct 28, 2020 at 14:40
• Sorry I did a typo, again... :. $A$ and $C$ are $n\times k$ and $B$ is $n\times n$ Oct 28, 2020 at 14:55

Another nice way to do this sort of thing is numpy.einsum, which allows you to express the multiplication in index notation:

>>> A=np.array([[1,2,3],[4,5,6]])
>>> B=np.array([[1,2],[3,4]])
>>> C=np.array([[1,2,3],[4,5,6]])
>>> A.T@B@C
array([[ 85, 116, 147],
[113, 154, 195],
[141, 192, 243]])
>>> np.einsum('ij,jk,kl->il',A.T,B,C)
array([[ 85, 116, 147],
[113, 154, 195],
[141, 192, 243]])
>>> np.einsum('ij,jk,ki->',A.T,B,C)
482


This has a nice correspondence with how we would write this equation mathematically: $$\mathrm{Tr}(A^TBC)=\sum_{ijk}(A^T)_{ij}B_{jk}C_{ki}$$ Einsum also has an optimize keyword, which has it search for the fastest contraction order. This could make the multiplication a lot faster depending on the dimensions of your matrices.

I have not quite understood what the dimensions of your matrices are (see comments). but the following code may give you a start. It is not efficient or particularly elegant, but you can improve from there:

   import numpy as np
A = np.matrix('1 2; 3 4')
B = np.matrix('1 2; 3 4')
C = np.matrix('1 2; 3 4')
result = (np.matmul(A.transpose(),np.matmul(B,C)))
print(result)
print(np.trace(result))


Generally speaking, the people here will be able to help better if we get all the info and a bit of context for your problem, like a small sample problem etc.

• Thanks. Was numpy.tensordot(A, B@C,axis=2) okay though? Oct 28, 2020 at 14:57
• You'd miss the transposition, but that should also work. You can just try it out for a couple of examples to get a feel for it and compare different realizations. python lets you change things quick. Oct 28, 2020 at 15:54