# How to make the following code work faster?

The problem

I would like to compute the Lorentz transformation $$\mathbf{p}'$$ of an arbitrary 3-momentum vector $$\mathbf{p}$$ given a boost characterized by some other momentum vector $$\mathbf{p}_{N}$$, and then calculate polar and azimuthal angles $$\theta',\phi'$$. The procedure is given below: $$\mathbf{p}' = \mathbf{p} + \gamma_{N}\mathbf{v}_{N}E + \Gamma \mathbf{v}_{N}(\mathbf{v}_{N}\cdot \mathbf{p})$$ Here, $$\quad |\mathbf{p}| = \sqrt{E^{2}-m^{2}}, \quad |\mathbf{p}_{N}| = \sqrt{E_{N}^{2}-m_{N}^{2}},\quad \mathbf{v}_{N} = \frac{\mathbf{p}_{N}}{E_{N}}, \quad \gamma_{N} = \frac{E_{N}}{m_{N}}, \quad \Gamma = \frac{\gamma_{N}-1}{v_{N}^{2}}$$ where in spherical coordinates $$\mathbf{p} = |\mathbf{p}|(s(\theta)c(\phi),s(\theta)s(\phi),c(\theta)), \quad \mathbf{p}_{N} = |\mathbf{p}_{N}|(s(\theta_{N})c(\phi_{N}),s(\theta_{N})s(\phi_{N}),c(\theta_{N}))$$ Finally, $$\theta' = \arccos\left( \frac{p'_{z}}{|\mathbf{p}'|}\right), \quad \phi' = \begin{cases}\arccos\left(\frac{p_{x}'}{\sqrt{p_{x}^{'2}+p_{y}^{'2}}} \right), \quad p_{y}' > 0, \\ -\arccos\left(\frac{p_{x}'}{\sqrt{p_{x}^{'2}+p_{y}^{'2}}} \right),\quad p_{y}' < 0 \end{cases}$$

The initial parameters here are $$\theta_{N},\phi_{N},E_{N},m_{N},m,E,\theta,\phi$$

The output must be

$$\mathbf{p}',\quad \theta', \quad \phi\$$

My attempt to resolve.

I have written a piece of code that computes all these parameters. Please find it below:

(*Position of the point in coordinate space in terms of spherical \
coordinates*)
θVal[x_, y_, z_] = ArcCos[z/Sqrt[x^2 + y^2 + z^2]];
ϕVal[x_, y_] =
If[y > 0, ArcCos[x/Sqrt[x^2 + y^2]], -ArcCos[x/Sqrt[x^2 + y^2]]];
(*N momentum, velocity and γ,Γ factors at its \
lab frame*)
pNvec[EN_, mN_, θN_, ϕN_] =
Sqrt[EN^2 - mN^2] {Sin[θN] Cos[ϕN],
Sin[θN] Sin[ϕN], Cos[θN]};
vvec[EN_, mN_, θN_, ϕN_] = +(
pNvec[EN, mN, θN, ϕN]/EN);
γFactor[EN_, mN_] = EN/mN;
Γfactor[EN_, mN_] =
Simplify[(γFactor[EN, mN] -
1)/((v /.
Solve[γFactor[EN, mN] == 1/Sqrt[1 - v^2], v])^2)[[1]]];
(*Momentum of X at the rest frame of N*)
pproductRestVec[EXrest_, mX_, θXrest_, ϕXrest_] =
Sqrt[EXrest^2 - mX^2] {Sin[θXrest] Cos[ϕXrest],
Sin[θXrest] Sin[ϕXrest], Cos[θXrest]};
(*Decay product's momentum and energy at N's lab frame*)
pproductLabVec[EN_, mN_, θN_, ϕN_, EXrest_,
mX_, θXrest_, ϕXrest_] =
Simplify[pproductRestVec[EXrest,
mX, θXrest, ϕXrest] + γFactor[EN, mN]*
vvec[EN, mN, θN, ϕN]*
EXrest + Γfactor[EN, mN]*
vvec[EN,
mN, θN, ϕN] (vvec[EN,
mN, θN, ϕN].pproductRestVec[EXrest,
mX, θXrest, ϕXrest])];
EproductLabVec[EN_, mN_, θN_, ϕN_, EXrest_,
mX_, θXrest_, ϕXrest_] =
Simplify[γFactor[EN,
mN] (EXrest +
vvec[EN, mN, θN, ϕN].pproductRestVec[EXrest,
mX, θXrest, ϕXrest])];
(*X product's θ and ϕ at lab frame*)
pproductLabVec3[EN_, mN_, θN_, ϕN_, EXrest_,
mX_, θXrest_, ϕXrest_] =
pproductLabVec[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest][[3]];
pproductLabVec2[EN_, mN_, θN_, ϕN_, EXrest_,
mX_, θXrest_, ϕXrest_] =
pproductLabVec[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest][[2]];
pproductLabVec1[EN_, mN_, θN_, ϕN_, EXrest_,
mX_, θXrest_, ϕXrest_] =
pproductLabVec[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest][[1]];
θproductLabVec =
Compile[{{EN, _Real}, {mN, _Real}, {θN, _Real}, {ϕN, \
_Real}, {EXrest, _Real}, {mX, _Real}, {θXrest, _Real}, \
{ϕXrest, _Real}},
ArcCos[pproductLabVec3[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]/((*Norm[pproductLabVec[EN,
mN,θN,ϕN,EXrest,
mX,θXrest,ϕXrest]]*)√(pproductLabVec1[EN,
mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]^2 +
pproductLabVec2[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]^2 +
pproductLabVec3[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]^2))]]
ϕproductLabVec =
Compile[{{EN, _Real}, {mN, _Real}, {θN, _Real}, {ϕN, \
_Real}, {EXrest, _Real}, {mX, _Real}, {θXrest, _Real}, \
{ϕXrest, _Real}}, ϕVal[
pproductLabVec1[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest],
pproductLabVec2[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]]]

Question

However, it turns out that it works quite slow:

ENvalue = 20;
mNvalue = 1.5;
θNvalue = 0.03;
ϕNvalue = Pi/3;
EXrestValue = mNvalue/2;
mXvalue = 0.01;
θXrestValue = Pi/8;
ϕXrestValue = 0.1;
AbsoluteTiming[ϕproductLabVec[ENvalue,
mNvalue, θNvalue, ϕNvalue, EXrestValue,
mXvalue, θXrestValue, ϕXrestValue]]*10^6
AbsoluteTiming[θproductLabVec[ENvalue,
mNvalue, θNvalue, ϕNvalue, EXrestValue,
mXvalue, θXrestValue, ϕXrestValue]]*10^6
AbsoluteTiming[
pproductLabVec[ENvalue, mNvalue, θNvalue, ϕNvalue,
EXrestValue, mXvalue, θXrestValue, ϕXrestValue]]*10^6

{75.1, 736183.}

{74.7, 39506.9}

{265.5, {564138., 511174., 1.92596*10^7}}

In situations when there are a lot of points (say, $$10^7$$), it is very crucial. Could you please point out what parts of the code may be improved in order to speed up it?

• You might find it helpful to simply use $\phi'=\arctan(p'_y/p'_x)$, which accounts for both the positive and negative case! Or, in mathematica, ArcTan[px, py], which accounts also for the $p'_x=0$ case. Commented Mar 27, 2021 at 0:53
• @John Taylor You made a typo in formula for $\bf {v}_N$ in Latex. Commented Mar 27, 2021 at 12:59
• @AlexTrounev : thanks! Commented Mar 27, 2021 at 14:19

You're using Compile in wrong way. I'd suggest reading this and this post as a start. The following is a quick fix for your code:

rule = Flatten[
DownValues /@ {pproductLabVec3, pproductLabVec2, pproductLabVec1, ϕVal}];
θproductLabVec =
Hold@Compile[{{EN, _Real}, {mN, _Real}, {θN, _Real}, {ϕN, _Real}, {EXrest, _Real},
{mX, _Real}, {θXrest, _Real}, {ϕXrest, _Real}},
ArcCos[pproductLabVec3[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]/(√(pproductLabVec1[EN,
mN, θN, ϕN, EXrest, mX, θXrest, ϕXrest]^2 +
pproductLabVec2[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]^2 +
pproductLabVec3[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]^2))]] //. rule // ReleaseHold;
ϕproductLabVec =
Hold@Compile[{{EN, _Real}, {mN, _Real}, {θN, _Real}, {ϕN, _Real}, {EXrest, _Real},
{mX, _Real}, {θXrest, _Real}, {ϕXrest, _Real}}, ϕVal[
pproductLabVec1[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest],
pproductLabVec2[EN, mN, θN, ϕN, EXrest,
mX, θXrest, ϕXrest]]] //. rule // ReleaseHold;

pproductLabVecCompiled =
Compile[{{EN, _Real}, {mN, _Real}, {θN, _Real}, {ϕN, _Real}, {EXrest, _Real},
{mX, _Real}, {θXrest, _Real}, {ϕXrest, _Real}}, #] &@
pproductLabVec[EN, mN, θN, ϕN, EXrest, mX, θXrest, ϕXrest];

Speed test:

ENvalue = 20;
mNvalue = 1.5;
θNvalue = 0.03;
ϕNvalue = Pi/3;
EXrestValue = mNvalue/2;
mXvalue = 0.01;
θXrestValue = Pi/8;
ϕXrestValue = 0.1;
AbsoluteTiming[
Do[ϕproductLabVec[ENvalue, mNvalue, θNvalue, ϕNvalue, EXrestValue,
mXvalue, θXrestValue, ϕXrestValue], {10^6}]]
(* {2.217114, Null} *)

AbsoluteTiming[
Do[θproductLabVec[ENvalue, mNvalue, θNvalue, ϕNvalue, EXrestValue,
mXvalue, θXrestValue, ϕXrestValue], {10^6}]]
(* {1.812695, Null} *)

AbsoluteTiming[
Do[pproductLabVecCompiled[ENvalue, mNvalue, θNvalue, ϕNvalue, EXrestValue,
mXvalue, θXrestValue, ϕXrestValue], {10^6}]]
(* {1.947791, Null} *)

The following takes approx. 4 seconds for 10^5 evaluations. A sped up could be obtained if the data is not fed one by one, but in bunches. Then instead of of calculating scalars, one could work with vectors.

boost[thn_, phn_, en_, mn_, m_, e_, ph_, th_] :=
Module[{mn2 = mn^2, vn, gamn, gam, pt, tht, pht, pn},
p = Sqrt[e^2 - m^2] {Sin[th] Cos[ph], Sin[th] Sin[ph], Cos[th]};
pn = Sqrt[en^2 - mn^2] {Sin[thn] Cos[phn], Sin[thn] Sin[phn],
Cos[thn]};
vn = pn/en;
gamn = en/mn;
gam = (gamn - 1)/vn.vn;

pt = p + gamn e vn + gam (vn.p) vn;
tht = ArcCos[pt[[3]]/Norm[pt]];
pht = ArcTan[pt[[1]], pt[[2]]];

{pt, tht, pht}
];

Here is a test case:

{m, mn} = RandomReal[1, {2}];
{e, en} = RandomReal[{1, 2}, {2}];
{th, thn} = RandomReal[Pi, {2}];
{ph, phn} = RandomReal[2 Pi, {2}];
p = RandomReal[1];

Do[boost[thn, phn, en, mn, m, e, ph, th], 10^5]; // Timing

(** {3.95313, Null} *)
• Is the convention in your code {EN, mN, θN, ϕN, EXrest, mX, θXrest, ϕXrest} -> {en, mn, thn, phn, e, m, th, ph} // Thread? If so, the output of your code doesn't match that of OP's. Commented Mar 27, 2021 at 11:12
• You have a typo in definition gamn, it should be gamn = en/mn. Also code by John Taylor has a typo in definition of $\Gamma$. Commented Mar 27, 2021 at 11:21
• @Alex Trounev Thank's, nobody is perfect. Commented Mar 27, 2021 at 12:16
• @Alex Perhaps $\mathbf{v}_{N}$ and $v_{N}$ are different? This should be clarified by OP, of course. Commented Mar 27, 2021 at 12:34
• @xzczd Sorry, his code has no typo, but he made typo in Latex with vn definition, it should be vn = pn/en. Commented Mar 27, 2021 at 12:57