Simulating molecular dynamics efficiently

Question

The following website offers a very nice molecular dynamics simulation: http://physics.weber.edu/schroeder/md/. It is pretty neat and quite a few physical phenomena can be described from that small scale, see the author's flyer. The author mentions that the code can run smoothly in a browser with 1000 particules and a 2000 steps per second (only displaying a few dozens per second, of course). I wanted to simulate something similar with Mathematica but am facing efficiency issues. I am not interested in features such as live manipulation of particules, I mostly want to choose initial conditions, compute the time evolution, and visualize the results later.

This is my code:

n = 200; (* number of particules *)
h = 0.01; (* time step *)
size = 50.; (* size of the box *)

Clear[x, a]
(* initial conditions *)
x[0] = Table[{2 Mod[i, size] - size, Floor[i/50] - size}, {i, n}];
SeedRandom[1234];
v[0] = 0.1*size*RandomReal[{-1, 1}, {n, 2}];
x[1] = x[0] + h*v[0];
(* Verlet integration scheme *)
mem : x[i_] := mem = 2 x[i - 1] - x[i - 2] + h^2*a[i - 1];
mem : a[i_] := mem = Table[forceParticules[x[i], j, epsilon], {j, n}] 
                     +  Table[forceWalls[x[i], j, 100], {j, n}];

Explanations: the particules are initial stacked at the bottom of a square box, with random initial velocities. As in the applet, the time integration scheme is a Verlet integration scheme. The acceleration at time i+1, denoted by a[i], comprises two things:

• the force from surrounding particles (forceParticules). The trick is to consider only the influence of surrounding particles, at distance of at most epsilon, to reduce computation time. The surrounding particles exert of force deriving from a Lennard-Jones potential.
• the force exerted by the boundary of the box (forceWalls).

This is the definition I used, and both are called hundreds of hundreds of times.

r0 = 1; (* diameter of one particle *)
epsilon = 10; (* radius of influence *)
potential[r_] = 4 * ((r0/r)^12 - (r0/r)^6);
dpotential[r_] = potential'[r];
forceParticules[xi_, i_, epsilon_] := Module[{},
  close = Nearest[Drop[xi, {i}], xi[[i]], {100, epsilon}];
  vecs = # - xi[[i]] & /@ close;
  output = Total[dpotential[Norm[#]]*#/Norm[#] & /@ vecs]
  ]

stiffness[x_] = 
  Piecewise[{{-100*(x + size), x < -size}, {-100*(x - size), 
     x > size}}, 0];
forceWalls[xi_, i_, raideur_] := {stiffness[xi[[i, 1]]], 
  stiffness[xi[[i, 2]]]}

It then suffices to evaluate the x[i] and view the result:

 Table[x[i], {i, 500}];
 Animate[ListPlot[x[i], PlotRange -> size*{{-1, 1}, {-1, 1}}, 
   AspectRatio -> 1, PlotStyle -> PointSize[1/size]], {i, 1, 100, 1}]

My problem is that the code runs very slowly (13 seconds for 500 steps with 50 particles). I am miles away from the efficiency of the applet.

Answer 1

Okay, here is a way to compute the forces much faster: We create a CompiledFunction (called getForces). It eats a list of points in the plane and spits out the net force onto the first point of the list; here the second to last points are supposed to be those points that are so close to the first one that they exert a force onto it.

size = 50.;(*size of the box*)
box = {{-size, size}, {-size, size}};
r0 = 1.;(*diameter of one particle*)

Quiet[Block[{r, x1, x2, y1, y2, xx, yy, force, potential},
   xx = {x1, x2};
   yy = {y1, y2};
   potential = r \[Function] 1/4 ((r0/r)^2 - (r0/r));
   force = -D[potential[Sqrt[Dot[xx - yy, xx - yy]]], {xx, 1}];
   With[{
     f1code = N@force[[1]], f2code = N@force[[2]], slope = 100.,
     a1 = N@box[[1, 1]], b1 = N@box[[1, 2]], a2 = N@box[[2, 1]], 
     b2 = N@box[[2, 2]]
     },
    getForces = Compile[{{X, _Real, 2}},
      Block[{x1, x2, y1, y2, f1, f2},
       x1 = Compile`GetElement[X, 1, 1];
       x2 = Compile`GetElement[X, 1, 2];
       f1 = slope (Ramp[a1 - x1] - Ramp[x1 - b1]);
       f2 = slope (Ramp[a2 - x2] - Ramp[x2 - b2]);
       Do[
        y1 = Compile`GetElement[X, i, 1];
        y2 = Compile`GetElement[X, i, 2];
        f1 += f1code;
        f2 += f2code;
        , {i, 2, Length[X]}
        ];
       {f1, f2}
       ],
      CompilationTarget -> "C",
      RuntimeAttributes -> {Listable},
      Parallelization -> True
      ]
    ]]];

The following is a (not so pure) function that computes the $n \times 2$ matrix xnew consisting of the new point positions; it uses the matrix x of positions at a given time instance and the matrix xold which represents the positions at the preceeding time instance. xold is handled as side effect which makes it not too nicely, but this way, we can use it in NestList later.

step = x \[Function] (
    xnew =  2. x - xold +  h^2 (1./m) getForces[Nearest[x, x, {\[Infinity], epsilon}]];
    xold = x;
    xnew
    );

Setting up the remaining parameters and the initial conditions...

SeedRandom[1234];
n = 200;(*number of particules*)
h = 0.01;(*time step*)
epsilon = 10.;(*radius of influence*)
(*initial conditions*)
x0 =   N@Table[{2 Mod[i, size] - size, Floor[i/50] - size}, {i, n}];
v0 = 0.1 size RandomReal[{-1, 1}, {n, 2}];
x1 = x0 + h v0;(*Verlet integration scheme*)
m = ConstantArray[1., n];(*particle masses*)

... and now, we let it run:

timesteps = 10000;
xold = x0;
x = x1;
data = Join[{x0}, NestList[step, x1, timesteps]]; // AbsoluteTiming

{4.58476, Null}

Aha, roughly 2000 iterations per second. That's at least not so much worse than the JavaScript implementation...

And here a visualization:

frames = Map[X \[Function] Graphics[{PointSize[1/size ],
      Point[X]}, PlotRange -> box, PlotRangePadding -> 0.1 size
     ], data];
Export["a.gif", frames[[1 ;; -1 ;; 20]]]

Note that I use a different potential (less singular) in order to get it running. Of course, you can play with the parameters...

Further improvements

In the meantime, I did some hand tuning of the compiled code. The new positions get now computed completely within the CompiledFunction. That's the first new thing. The second is that for each point also all current positions, the preceeding position of this point, and an index list ilist with the indices for the relevant points are handed over. The first index in ilist markes the point for which we want to compute the new position; the other entries mark the points in the region of interaction. This way, we can recycle the index lists ilists that Nearest will produce every now and then in the main loop (see below).

size = 50.;(*size of the box*)
box = {{-size, size}, {-size, size}};
r0 = 1.;(*diameter of one particle*)
Quiet[
 Block[{x1, x2, y1, y2, xx, yy, force, potential, r, r2},
  xx = {x1, x2};
  yy = {y1, y2};
  potential = r \[Function] 4 ((r0/r)^12 - (r0/r)^6);
  force = Simplify[
    -D[potential[Sqrt[Dot[xx - yy, xx - yy]]], {xx, 1}] 
    /. (x1 - y1)^2 + (x2 - y2)^2 -> r2
    ] /. Sqrt[r2] -> r;

  With[{
    f1code = N@force[[1]], f2code = N@force[[2]], slope = 100.,
    a1 = N@box[[1, 1]], b1 = N@box[[1, 2]], a2 = N@box[[2, 1]], 
    b2 = N@box[[2, 2]]
    },
   getStep = 
    Compile[{{X, _Real, 2}, {Xold, _Real, 1}, {ilist, _Integer, 
       1}, {factor, _Real}},
     Block[{x1, x2, y1, y2, f1, f2, r, r2, j, i},
      j = Compile`GetElement[ilist, 1];
      x1 = Compile`GetElement[X, j, 1];
      x2 = Compile`GetElement[X, j, 2];
      f1 = slope (Ramp[a1 - x1] - Ramp[x1 - b1]);
      f2 = slope (Ramp[a2 - x2] - Ramp[x2 - b2]);
      Do[
       i = Compile`GetElement[ilist, k];
       y1 = Compile`GetElement[X, i, 1];
       y2 = Compile`GetElement[X, i, 2];
       r2 = (x1 - y1)^2 + (x2 - y2)^2;
       r = Sqrt[r2];
       f1 += f1code;
       f2 += f2code;
       , {k, 2, Length[ilist]}
       ];
      {
       2. x1 - Compile`GetElement[Xold, 1] + factor f1,
       2. x2 - Compile`GetElement[Xold, 2] + factor f2
       }
      ],
     CompilationTarget -> "C",
     RuntimeAttributes -> {Listable},
     Parallelization -> True,
     RuntimeOptions -> "Speed"
     ]
   ]

  ]];

Preparing some constants...

SeedRandom[1234];
n = 1000;(*number of particules*)
h = 0.005;(*time step*)

epsilon = 3.;(*radius of influence*)
(*initial conditions*)

x0 = Developer`ToPackedArray[
   N@Table[{2 Mod[i, size] - size, Floor[i/50] - size}, {i, n}]];
v0 = 0.1 size RandomReal[{-1, 1}, {n, 2}];
x1 = x0 + h v0;

m = ConstantArray[1., n];(*particle masses*)
factors = (h^2./m);
timesteps = 10000;
skip = 10; (*Nearest gets called only every 10th time iteration*)

Main loop

We write directly into a preallocated array (which makes no difference since NestList is cleverly implemented). The major change is that we request only the vertex indices from Nearest (with x -> Automatic). In contrast to the coordinates, these won't change significantly for several iterations. So we have to call Nearest less than once per time iteration. Developer`ToPackedArray seems to be needed because the (listable) getStep is applied to ragged lists (which cannot be packed) and so the output is not packed.

data = ConstantArray[0., {timesteps + 1, n, 2}];
data[[1]] = x0;
data[[2]] = x1;
xold = x0;
x = x1;
ilists = Nearest[x -> Automatic, x, {\[Infinity], epsilon}];
Do[
   If[Mod[iter, skip] == 0,
    ilists = Nearest[x -> Automatic, x, {\[Infinity], epsilon}];
    ];
   data[[iter]] = xnew = Developer`ToPackedArray[getStep[x, xold, ilists, factors]];
   xold = x;
   x = xnew;
   ,
   {iter, 3, timesteps + 1}]; // AbsoluteTiming

{4.37231, Null}

The timing is much better than 2000 steps per second now.

I have also prepared a somewhat nicer visualization:

velocities = Sqrt[(Join[{v0}, Differences[data]/h]^2).ConstantArray[1., {2}]];
colorcoords = Rescale[Clip[velocities, {-100, 100}]];
frames = Table[
   Graphics[{
     PointSize[0.5 r0/size],
     Transpose[{ColorData["TemperatureMap"] /@ colorcoords[[i]], 
       Point /@ data[[i]]}],
     }, PlotRange -> box, PlotRangePadding -> 0.1 size, 
    Background -> Black
    ],
   {i, 1, Length[data], 100}];

Manipulate[frames[[i]],{i, 1, Length[frames], 1}]

Answer 2

To expand on @HenrikSchumacher's comment, compare:

r1 = Table[
    Nearest[Drop[x[1],{i}], x[1][[i]], {100,10}],
    {i,200}
];//AbsoluteTiming

nf = Nearest[x[1]]; //AbsoluteTiming
r2 = nf[x[1], {100, 10}][[All, 2;;]]; //AbsoluteTiming

r1 === r2

{0.004298, Null}

{0.000029, Null}

{0.000463, Null}

True

Addendum

Also, note that you can have the NearestFunction return distances as well:

r1 = Table[
    Norm[#-x[1][[i]]]& /@ Nearest[Drop[x[1],{i}], x[1][[i]], {100,10}],
    {i,200}
]; //AbsoluteTiming

nf = Nearest[x[1]->"Distance"];//AbsoluteTiming
r2 = nf[x[1], {100, 10}][[All, 2;;]]; //AbsoluteTiming

r1 === r2

{0.021323, Null}

{0.000049, Null}

{0.000326, Null}

True