# Using three Table functions together makes the code really slow

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.

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.

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


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}