Using three Table functions together makes the code really slow

Question

I have a list with shape {5,90,6000} and the result is of shape {5,90,2000}. Following is my code:

finalResultEnergy = 
      Table[Table[
        Table[Sqrt[
          momentum[[index, j, i]]^2 + momentum[[index, j, i + 1]]^2 + 
           momentum[[index, j, i + 2]]^2 + m^2], {i, 1, 
          Dimensions[momentum][[3]], 3}], {j, 1, 
         Dimensions[momentum][[2]], 1}], {index, 1, 
        Dimensions[momentum][[1]], 1}];

momentum is of shape {5,90,6000}. The code is really slow and I don't know why. This is really fast in python.

score 3 · Answer 1 · answered Dec 03 '17 at 04:27

Here are two improvements. I'm sure there are more, but these seem to be fairly significant (although I can't be sure since I never waited long enough to see the timing on your nested Tables).

The first is simply to use a single table to construct:

singletable[momentumarray_] := Table[
   Sqrt[momentumarray[[index, j, i]]^2 + 
     momentumarray[[index, j, i + 1]]^2 + 
     momentumarray[[index, j, i + 2]]^2 + m^2],
   {index, Dimensions[momentumarray][[1]]},
   {j, Dimensions[momentumarray][[2]]},
   {i, 1, Dimensions[momentumarray][[3]], 3}];

The second is to use Map and Apply (@@@):

mapply[momentumarray_] := 
  Map[Sqrt[#1^2 + #2^2 + #3^2 + m^2] & @@@ # &, 
   Map[Partition[#, 3] &, momentumarray, {2}], {2}];

I've put your three table setup into the function threetable.

Here's some tests on a smaller array:

momentum = RandomReal[1, {5, 45, 900}];
AbsoluteTiming[res1 = threetable[momentum];]
AbsoluteTiming[res2 = singletable[momentum];]
AbsoluteTiming[res3 = mapply[momentum];]
res1 == res2 == res3

(* {7.7918, Null}
   {0.514314, Null}
   {0.382995, Null}
   True  *)

So both singletable and mapply produce the same output and offer a significant speed up. Leaving out threetable for tests on the entire momentum array, we get:

momentum = RandomReal[1, {5, 90, 6000}];
AbsoluteTiming[res2 = singletable[momentum];]
AbsoluteTiming[res3 = mapply[momentum];]
res2 == res3

(* {7.02987, Null}
   {5.2053, Null}
   True *)

score 2 · Accepted Answer · answered Dec 03 '17 at 04:25

No need to use nested tables (see also here). The documentation for Table gives:

Table[expr,{i,imin,imax},{j,jmin,jmax},…] gives a nested list. The list associated with i is outermost.

momentum = RandomReal[{-1, 1}, {5, 90, 6000}];

With nested tables:

AbsoluteTiming[
 finalResultEnergyNested = Table[Table[Table[
      Sqrt[
       momentum[[index, j, i]]^2 + momentum[[index, j, i + 1]]^2 + 
        momentum[[index, j, i + 2]]^2 + m^2],
      {i, 1, Dimensions[momentum][[3]], 3}],
     {j, 1, Dimensions[momentum][[2]], 1}],
    {index, 1, Dimensions[momentum][[1]], 1}
    ];
 ]

(*{205.798, Null}*)

With non-nested tables:

AbsoluteTiming[
 finalResultEnergy = Table[
   Sqrt[momentum[[index, j, i]]^2 + momentum[[index, j, i + 1]]^2 + 
     momentum[[index, j, i + 2]]^2 + m^2],
   {index, 1, Dimensions[momentum][[1]], 1},
   {j, 1, Dimensions[momentum][[2]], 1},
   {i, 1, Dimensions[momentum][[3]], 3}
   ];
 ]

(*{37.4112, Null}*)

A third option:

AbsoluteTiming[
 finalResultEnergyAlt = 
   Sqrt[Map[Plus @@@ Partition[#, 3] &, momentum^2, {2}] + m^2];
 ]

(*{16.1016, Null}*)

These are just some ideas (not sure if there are better practices), but hope that helps a bit.

Henrik Schumacher · Answer 3 · 2017-12-03T08:47:15.373

Your code can easily be vectorized so that no loop constructs are needed.

Here is some random data.

momentum = RandomReal[{-1, 1}, {5, 90, 6000}];
m = 1.;

This is the actual code:

finalResultEnergy = 
   With[{dims = {
      Dimensions[momentum][[1]],
      Dimensions[momentum][[2]],
      Floor[Dimensions[momentum][[3]]/3], 3}
      },
    With[{c = ArrayReshape[momentum, dims]},
     Sqrt[(c c).ConstantArray[1., {3}]+ m^2]
     ]
    ]; // RepeatedTiming

{0.019, Null}

This is roughly 100 times faster than mapapply by @aardvark2012.

c = ArrayReshape[momentum, dims] transforms your {5, 90, 6000} array into an array of dimensions {5, 90, 2000, 3}. The order of the elements in memory doesn't change, so this is fairly efficient. But still, it makes up half of the running time. The actual code is this

Sqrt[(c c).ConstantArray[1., {3}]+ m^2]

Note that multiplication c c and the Sqrt in the end are performed elementwise.

Final remark

I mentioned that ArrayReshape is not for free. So you might consider to store your momenta in an array of dimensions {5, 90, 2000, 3} in the first place:

momentum = RandomReal[{-1, 1}, {5, 90, 2000, 3}];
m = 1.;
finalResultEnergy = 
   Sqrt[(momentum momentum).ConstantArray[1., {3}] + m^2]; // RepeatedTiming

{0.0069, Null}

Using three Table functions together makes the code really slow

3 Answers3