I am trying to find the best way to compute only the diagonal part of a product of two (or three) matrices. I tried two different approaches: computing the standard matrix product and then extracting the diagonal part; extracting rows and columns from the external matrices and making a Table. The first approach is of course very costly as it computes many more elemenst than needed. And this becomes evident with two matrices

matrA = RandomReal[NormalDistribution[], {3000, 3000}];
matrB = RandomReal[NormalDistribution[], {3000, 3000}];
Timing[vec0 = Diagonal[matrA.matrB];]
Timing[Block[{matrBT = Transpose[matrB]}, vec1 = Table[matrA[[i]].matrBT[[i]], {i, 1, Length[matrA]}]];]

As expected, the first expression takes much longer than the second one. However, I was surprised to find out that with three matrices, this does not seem to be true anymore.

matrA = RandomReal[NormalDistribution[], {3000, 3000}];
matrB = RandomReal[NormalDistribution[], {3000, 3000}];
matrC = RandomReal[NormalDistribution[], {3000, 3000}];
Timing[vec0 = Diagonal[matrA.matrC.matrB];]
Timing[Block[{matrBT = Transpose[matrB]}, vec1 = Table[matrA[[i]].matrC.matrBT[[i]], {i, 1, Length[matrA]}]];]

Now surprisingly the first approach is much faster than the second.

It is clear that the first approach cannot be optimal. So what would you suggest to compute this kind of quantities?


How about

mA = RandomReal[NormalDistribution[], {3000, 3000}];
mB = RandomReal[NormalDistribution[], {3000, 3000}];
mC = RandomReal[NormalDistribution[], {3000, 3000}];

Timing[vec0 = Diagonal[mA.mC.mB];]
(*{7.515625, Null}*)

Timing[Total[Total[mA Transpose[mC.mB], {2}] - vec0]]
(*{3.968750, 3.8557602 10^-10}*)
| improve this answer | |
  • $\begingroup$ cool! Apparently it seems all the improvement comes by computing the product mC.mB separately... This seems to work pretty well: Timing[Block[{matrBT = Transpose[matrC.matrB]}, vec1 = Table[matrA[[i]].matrBT[[i]], {i, 1, Length[matrA]}]];] $\endgroup$ – abenassen Jul 5 '16 at 12:05

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.