Recovering Matlab performance scaling via vectorization

Question

I'm trying to run a simple particle trajectory simulation, in which I propagate ~10,000 particles through a potential energy surface. In other programs like Matlab, I of course get a huge speedup when I vectorize my code such that all particles take time steps together, and I apply Newton's laws to entire vectors holding the positions/ velocities. In Mathematica, I'm trying to use this same method but see that the simulation time scales exactly linearly with the number of particles I put in. Is this typical, or should Mathematica also benefit from vectorized code?

singlestep[rv_, dt_, dUx_, dUy_, dUz_, nf_] :=

 Module[{fx, fy, fz, xx, yy, zz, xnew, ynew, znew, vx, vy, vz, vxnew, 
   vynew, vznew},
  xx = rv[[;; , 1]]; yy = rv[[;; , 2]]; zz = rv[[;; , 3]];
  vx = rv[[;; , 4]]; vy = rv[[;; , 5]]; vz = rv[[;; , 6]];
  Do[
   fx = -dUx /@ Transpose@{xx, yy, zz};
   fy = -dUy /@ Transpose@{xx, yy, zz};
   fz = -dUz /@ Transpose@{xx, yy, zz};
   xx += vx* dt + 1/2*fx/M*dt^2;
   yy += vy *dt + 1/2*fy/M*dt^2;
   zz += vz* dt + 1/2*fz/M*dt^2;
   vx += fx/M*dt;
   vy += fy/M*dt;
   vz += fz/M*dt,
   {nn, 1, nf}];
  xnew = xx; ynew = yy; znew = zz;
  vxnew = vx; vynew = vy; vznew = vz;
  {xnew, ynew, znew, vxnew, vynew, vznew}]

nf is a variable referring to how many "fine" steps are taken (but for this test I am setting it to 1). dUx/y/z are interpolating functions that contain the force components as a function of position. rv is an array that contains the 3 positions and 3 velocity components from the previous steps, so for each particle is size 1x6. As I add more particles by increasing the size of rv from 1x6 to nx6, I see the code slow down exactly proportional to n.

Answer 1

In Mathematica, I'm trying to use this same method but see that the simulation time scales exactly linearly with the number of particles I put in.

What do you expect, bro? Increasing the number of flops linearly will inevitably lead to linear growth in (pure) computation time. Maybe you see some sublinear scaling of the runtime in Matlab because of the hidden JIT compiler which may introduce a certain overhead.

Back to the code. Currently, I see three issues that may lead to slowdown: 1. Excessive use of Transpose. This coordinate-wise processing of vectors is so matlabish. It makes code unreadible and seldomly leads to any improvements. 2. Map in the form of /@, and 3. the Do (but that is really less an issue).

As you do not provide dUx, dUy, dUz, I use the gradient of a potential in the following:

This is the code:

Quiet[Block[{x, xx},
   xx = Table[x[[i]], {i, 1, 3}];
   f = x \[Function] Evaluate[-Exp[-x.x/2]];
   gradf = x \[Function] Evaluate[D[f[xx], {xx, 1}]];

   dUx = x \[Function] Evaluate[D[f[xx], xx[[1]]]];
   dUy = x \[Function] Evaluate[D[f[xx], xx[[2]]]];
   dUz = x \[Function] Evaluate[D[f[xx], xx[[3]]]];

   With[{code = gradf[xx]},
    cgradf = Compile[{{x, _Real, 1}},
      code,
      CompilationTarget -> "C",
      RuntimeAttributes -> {Listable},
      Parallelization -> True
      ]
    ];
   ]];

step = r \[Function] With[{f = -cgradf[r[[1]]]}, 
    {r[[1]] + r[[2]] dt + (0.5 dt^2/M) f, r[[2]] + (dt/M) f}
   ];

This is some example data

n = 100000;
p = RandomReal[{-1, 1}, {n, 3}];
v = RandomReal[{-1, 1}, {n, 3}];
M = 1.;
dt = 0.1;

And this is the actual test with timings

stepcount = 100;
rv = Join[p, v, 2];
bla = singlestep[rv, dt, dUx, dUy, dUz, stepcount]; // AbsoluteTiming
blubb = Nest[step, {p, v}, stepcount]; // AbsoluteTiming
Max[Abs[Transpose@Join[blubb[[1]], blubb[[2]], 2] - bla]]

{5.26767, Null}

{0.719926, Null}

7.34968*10^-14