2
$\begingroup$

I need to compute the following sum as fast as possible: $$P_{ijkl}=\sum_{p,q,r,s}^nW_{sqpr}H_{js}G_{rl}A_{ipkq}$$

I came up with this code:

PP = Transpose[Flatten[H.W.G, {{1}, {4}, {2, 3}}].Flatten A, {{4, 2}, {3}, {1}}], {2, 4, 3, 1}];

which is approximately 50 times faster than the brute force implementation for dimensions n=3. I want to speed this computation as much as possible. Below I provide definitions of the ingredients. As can be seen, some matrices possess symmetries. You are free to use them, although I found it is hard. I am rather thinking in the direction of correct index ordering in the Flatten commands (there are several possibilities, but I do not know which one is better) and of compilation.

n = 3;
W = SymmetrizedArray[RandomReal[{0, 1}, {n, n, n, n}], {n, n, n, n}, 
                      {{{4, 2, 3, 1}, 1}, {{1, 3, 2, 4}, 1}, {{2, 4, 1, 3}, 1}}];
g = RandomComplex[{0, 1 + I}, {n, n}];
G = g + ConjugateTranspose[g];
H = IdentityMatrix[n] - G;
a = SymmetrizedArray[RandomComplex[{0, 1 + I}, {n, n, n, n}], {n, n, n, n}, {{2, 1, 4, 3}, 1}];
A = a + Conjugate[Transpose[a, {3, 4, 1, 2}]];

It would be perfect if someone can beat the run-time of my code (likely) or even improve its scaling with respect to n (unlikely).

Timing[P = 
   Table[Sum[
     W[[s, q, p, r]] A[[i, p, k, q]] H[[j, s]] G[[r, l]], {p, n}, {q, 
      n}, {r, n}, {s, n}], {i, n}, {j, n}, {k, n}, {l, n}];]
Timing[PP = 
   Transpose[Flatten[H.W.G, {{1}, {4}, {2, 3}}].Flatten[A, {{4, 2}, {3}, {1}}], {2, 4, 3, 1}];]
Norm[Flatten[PP - P]];
(*{0.507595, Null}*)
(*{0.014532, Null}*)
(*2.06669*10^-14*)
$\endgroup$
4
  • $\begingroup$ ArrayReshape is slightly faster than Flatten, but the downside is that you have to give the dimensions explicitly rather than just asking for e.g. the second and third index to be flattened. $\endgroup$ – ala10 Dec 12 '20 at 21:57
  • $\begingroup$ @ala10 Also Compile complains about complicated arguments in Flatten. Will it work with ArrayReshape? $\endgroup$ – yarchik Dec 12 '20 at 21:59
  • 1
    $\begingroup$ My understanding is that compiling won't help too much here. Dot already calls some low-level C libraries and parallelizes everything automatically. $\endgroup$ – ala10 Dec 12 '20 at 22:04
  • $\begingroup$ @ala10 Yes, I agree with you. However, I have a concern about the Flatten or ArrayReshape for that matter. I feel they are the bottleneck. $\endgroup$ – yarchik Dec 12 '20 at 22:08
1
$\begingroup$

Best to convert W and A to dense packed arrays with Normal; they are dense anyways and their symmetries do not compensate for not being able to use optimized dense matrix arithmetic. For n=6, this yields a 1000-fold speed-up:

n = 6;
W = SymmetrizedArray[
   RandomReal[{0, 1}, {n, n, n, n}], {n, n, n, 
    n}, {{{4, 2, 3, 1}, 1}, {{1, 3, 2, 4}, 1}, {{2, 4, 1, 3}, 1}}];
g = RandomComplex[{0, 1 + I}, {n, n}];
G = g + ConjugateTranspose[g];
H = IdentityMatrix[n] - G;
a = SymmetrizedArray[
   RandomComplex[{0, 1 + I}, {n, n, n, n}], {n, n, n, 
    n}, {{2, 1, 4, 3}, 1}];
A = a + Conjugate[Transpose[a, {3, 4, 1, 2}]];

nW = Normal[W];
nA = Normal[A];

PP = Transpose[
     Flatten[H.W.G, {{1}, {4}, {2, 3}}].Flatten[
       A, {{4, 2}, {3}, {1}}], {2, 4, 3, 1}]; // 
  AbsoluteTiming // First
nPP = Transpose[
     Flatten[H.nW.G, {{1}, {4}, {2, 3}}].Flatten[
       nA, {{4, 2}, {3}, {1}}], {2, 4, 3, 1}]; // 
  AbsoluteTiming // First
Max[Abs[PP - nPP]]

0.400711

0.000324

2.34371*10^-13

With a bit of more refactoring, I can get PP compute for n=40 in half a second. The preparatory computations take significantly longer, though. (And I was under the impression that SymmetrizedArray might have some memory related bug that made my kernel quit at times.)

n = 40;
W = Normal@SymmetrizedArray[
    RandomReal[{0, 1}, {n, n, n, n}],
    {n, n, n, n},
    {{{4, 2, 3, 1}, 1}, {{1, 3, 2, 4}, 1}, {{2, 4, 1, 3}, 1}}
    ];
G = # + ConjugateTranspose[#] &[RandomComplex[{0, 1 + I}, {n, n}]];
H = IdentityMatrix[n, WorkingPrecision -> MachinePrecision] - G;
A = # + Conjugate[Transpose[#, {3, 4, 1, 2}]] &[
   Normal@
    SymmetrizedArray[
     RandomComplex[{0, 1 + I}, {n, n, n, n}], {n, n, n, 
      n}, {{2, 1, 4, 3}, 1}]
   ];

PP = Transpose[
     Flatten[H.W.G, {{1}, {4}, {2, 3}}].Flatten[
       A, {{4, 2}, {3}, {1}}], {2, 4, 3, 1}]; // 
  AbsoluteTiming // First

0.489619

$\endgroup$

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.