How to make the following code work faster?

Question

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. — thorimur, Mar 27 '21 at 00:53
@John Taylor You made a typo in formula for $\bf {v}_N$ in Latex. — Alex Trounev, Mar 27 '21 at 12:59

score 4 · Accepted Answer · answered Mar 27 '21 at 10:52

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

Daniel Huber · Answer 2 · 2021-03-27T13:27:39.110

3

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

edited Mar 27 '21 at 13:27

answered Mar 27 '21 at 09:57

Daniel Huber

51,463
1
23
57

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. – xzczd Mar 27 '21 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$. – Alex Trounev Mar 27 '21 at 11:21
@Alex Trounev Thank's, nobody is perfect. – Daniel Huber Mar 27 '21 at 12:16
@Alex Perhaps $\mathbf{v}{N}$ and $v{N}$ are different? This should be clarified by OP, of course. – xzczd Mar 27 '21 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. – Alex Trounev Mar 27 '21 at 12:57
@DanielHuber There is second typo in your code with definition vn, it should be vn=pn/en. – Alex Trounev Mar 27 '21 at 13:04
@Alex gosh, you are again correct. It's some time since I did relativity. – Daniel Huber Mar 27 '21 at 13:29
@DanielHuber Now it is ok (+1). I am also pass relativity in the last century. :) – Alex Trounev Mar 27 '21 at 14:09

How to make the following code work faster?

2 Answers2

Linked