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


But what about the case with 3 matrices?

  • 3
    $\begingroup$ 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. $\endgroup$ – MPIchael Oct 28 '20 at 13:19
  • $\begingroup$ 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? $\endgroup$ – roi_saumon Oct 28 '20 at 14:11
  • $\begingroup$ Sorry, I meant $A$ and $B$ are $n\times k$ and $B$ is $n\times n$. $\endgroup$ – roi_saumon Oct 28 '20 at 14:23
  • $\begingroup$ 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:-) $\endgroup$ – MPIchael Oct 28 '20 at 14:40
  • $\begingroup$ Sorry I did a typo, again... :. $A$ and $C$ are $n\times k$ and $B$ is $n\times n$ $\endgroup$ – roi_saumon Oct 28 '20 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)

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.

  • $\begingroup$ Thanks, I will try this also $\endgroup$ – roi_saumon Oct 30 '20 at 18:46

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

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.

  • $\begingroup$ Thanks. Was numpy.tensordot(A, B@C,axis=2) okay though? $\endgroup$ – roi_saumon Oct 28 '20 at 14:57
  • $\begingroup$ 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. $\endgroup$ – MPIchael Oct 28 '20 at 15:54

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.