5
$\begingroup$

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.

$\endgroup$
2
$\begingroup$

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.

$\endgroup$
3
$\begingroup$

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 *)
$\endgroup$
2
$\begingroup$

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}

$\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.