Fast and Listable Piecewise function

Question

I am writing a simulation where I have an array with ten columns and on the order of millions of rows, for which I need to iterate on the order of ten thousand times, making computation time a major issue. At the moment I am using the Map[] function to apply calculations on every row (every row is independent of each other), but I realized that using the listability of the basic functions is significantly faster.

However, my problem is that I have one piecewise function that is not listable. I also cannot define a function outside of the calculations (as f[x_]:=...), since that incurs a huge bottleneck (factor ~20 slower computations). I am therefore looking for any tips or advice on how to solve this.

As an example, this is one of the calculations that are performed:

Q=0.18;
Qfactor=1.*10^-10;
Qspace=0.135;
sigma2289=7.83717;
Qspace17sigma2=0.0292835;
sigma17=3.7995;
sigma=1.64676;
TWOPI=2.*Pi;

Map[TWOPI*(Q + #[[7]]*Qfactor - 
  If[#[[5]]^2 + #[[6]]^2 < sigma2289, 
   Qspace17sigma2*Sqrt[#[[5]]^2 + #[[6]]^2], 
   Qspace/(Sqrt[#[[5]]^2 + #[[6]]^2] - sigma17)/sigma]) &, 
   particleArray];

My issue is this part:

 If[#[[5]]^2 + #[[6]]^2 < sigma2289, 
   Qspace17sigma2*Sqrt[#[[5]]^2 + #[[6]]^2], 
   Qspace/(Sqrt[#[[5]]^2 + #[[6]]^2] - sigma17)/sigma]

The conditional part of that function appears to not be listable. However, if I remove this and compare the computational times of the Map[] version, versus the Listability version, you can see that there is a significant improvement:

in: AbsoluteTiming[particleArray[[All, 10]] = 
       Map[TWOPI*(Q + #[[3]]*Qfactor - 
         Qspace17sigma2*Sqrt[#[[1]]^2 + #[[2]]^2]) &, particleArray];]
out: (0.14621, Null)

in: AbsoluteTiming[particleArray[[All, 8]] = 
     TWOPI*(Q + particleArray[[All, 3]]*Qfactor - 
     Qspace17sigma2*Sqrt[particleArray[[All, 1]]^2 + particleArray[[All, 2]]^2]);]
out: (0.01042, Null)

You can generate particleArray using the following code, for testing:

heavyGaussian = MixtureDistribution[{0.75, 0.25},                   
  {MultinormalDistribution[{0, 0}, {{1, 0}, {0, 1}}], 
   MultinormalDistribution[{0, 0}, {{1.8, 0}, {0, 1.8}}]}];

particleArray = ParallelTable[
   Flatten[{RandomVariate[heavyGaussian],
     RandomVariate[
       MultinormalDistribution[{0, 3.14159}, {{9.*10^-8,0},{0,0.0194882}}]],
 0, 0, 0, 0, 0, 0}], {x, 1, 100000}];
particleArray[[All, 3]] *= 2.6*10^10;

I also have CUDA enabled and a decent GPU, so any tips related to that would also be appreciated.

Answer 1

You may use Compile as follows:

Q = 0.18;
Qfactor = 1.*10^-10;
Qspace = 0.135;
sigma2289 = 7.83717;
Qspace17sigma2 = 0.0292835;
sigma17 = 3.7995;
sigma = 1.64676;
TWOPI = 2.*Pi;

cf = With[{
   TWOPI = TWOPI, Q = Q, Qfactor = Qfactor, Qspace = Qspace, 
   sigma = sigma, sigma17 = sigma17, sigma2289 = sigma2289, 
   Qspace17sigma2 = Qspace17sigma2},
  Compile[{{X, _Real, 1}},
   TWOPI*(Q + Compile`GetElement[X, 7]*Qfactor - If[
       Compile`GetElement[X, 5]^2 + Compile`GetElement[X, 6]^2 < 
        sigma2289, 
       Qspace17sigma2 Sqrt[
         Compile`GetElement[X, 5]^2 + Compile`GetElement[X, 6]^2], 
       Qspace/(Sqrt[
            Compile`GetElement[X, 5]^2 + Compile`GetElement[X, 6]^2] -
            sigma17)/sigma
       ]
     ),
   CompilationTarget -> "C",
   RuntimeAttributes -> {Listable},
   Parallelization -> True,
   RuntimeOptions -> "Speed"
   ]
  ]

And here is a speed comparison. First we have to generate particleArray. Notice how I generate it: I use RandomVariate with a second argument to produce many random numbers at once. Moreover, I use ConstantArray[0., {n, 6}] to generate all the zero (as machine precision 0., not as eaxct 0) and merge everything with Join[#,2]&. This way, particleArray is a packed array an can also be processed faster, in general. (The array woud be coerced to machine precision reals and packed when the CompiledFunction cf is applied to it anyways.)

n = 100000;
particleArray = Join[
     RandomVariate[heavyGaussian, n],
     RandomVariate[MultinormalDistribution[{0,3.14159}, {{9.*10^-8, 0}, {0, 0.0194882}}], n],
     ConstantArray[0., {n, 6}],
     2
     ]; // AbsoluteTiming // First

a = Map[
     TWOPI*(Q + #[[7]]*Qfactor - If[
          #[[5]]^2 + #[[6]]^2 < sigma2289, 
          Qspace17sigma2*Sqrt[#[[5]]^2 + #[[6]]^2], 
          Qspace/(Sqrt[#[[5]]^2 + #[[6]]^2] - sigma17)/sigma
          ]
        ) &, particleArray
     ]; // AbsoluteTiming // First

b = cf[particleArray]; // AbsoluteTiming // First
Max[Abs[a - b]]

0.020113

0.063061

0.003697

0.