How do I optimize a simulation containing both a While loop and a Do loop?

Question

I have written a simulation of a rotating molecule that I was hoping to get help to optimize. Below you see the main function in the simulation. The function FlatPart calculates the angle θ and the angular velocity ω from an obtained analytical solution to a Langevin equation. The function OtherPart just changes the angle and angular velocity to discontinuously, i.e., just add a number to the existing value.

\[Delta] = 0;

prob[\[Gamma]_] := 
  ProbabilityDistribution[(\[Gamma] Exp[-\[Gamma] t]), {t, 
    0, \[Infinity]}];

\[CapitalOmega]0 = Sqrt[1 - (\[Chi]/2)^2];

\[Theta][t_, \[Theta]0_, \[Omega]0_] := (\[Theta]0 Cos[
       t \[CapitalOmega]0] + ((\[Theta]0 \[Chi] + 2 \[Omega]0) Sin[
       t \[CapitalOmega]0])/(2 \[CapitalOmega]0)) E^(-((t \[Chi])/2));

\[Omega][t_, \[Theta]0_, \[Omega]0_] := (
  E^(-((t \[Chi])/
    2)) (4 \[Omega]0 \[CapitalOmega]0 Cos[
       t \[CapitalOmega]0] - (2 \[Chi] \[Omega]0 + \[Theta]0 \
(\[Chi]^2 + 4 \[CapitalOmega]0^2)) Sin[t \[CapitalOmega]0]))/(
  4 \[CapitalOmega]0);

FlatPart[\[Theta]0_, \[Omega]0_] := Module[{},
  \[Theta]1 = \[Theta][t1, \[Theta]0, \[Omega]0];
  \[Omega]1 = \[Omega][t1, \[Theta]0, \[Omega]0];
  par = {\[Theta]1, \[Omega]1}
  ]

OtherPart[\[Theta]0_, \[Omega]0_] := Module[{},
  \[Omega]1 = \[Omega]0 + \[CapitalDelta]\[CapitalOmega];
  \[Theta]1 = \[Theta]0 + \[CapitalDelta]\[CapitalOmega] \
\[CapitalDelta]t;
  par = {\[Theta]1, \[Omega]1}
  ]

mod[\[Theta]_, \[Omega]_, \[Gamma]_, \
\[CapitalDelta]\[CapitalOmega]1_, \[Chi]h_] := Module[{},
  \[CapitalDelta]ts = RandomVariate[prob[\[Gamma]], 50];
  T = Accumulate[\[CapitalDelta]ts];
  
  \[Chi] = \[Chi]h;
  \[CapitalDelta]\[CapitalOmega] = \[CapitalDelta]\[CapitalOmega]1;
  \[Theta]0 = \[Theta];
  \[Omega]0 = \[Omega];
  t0 = 0;
  \[CapitalTheta] = {\[Theta]0};
  \[CapitalOmega] = {\[Omega]0};
  x = {t0};
  i = 2;
  t1 = T[[1]] - \[Delta];
  
  While[\[Theta]0 < \[Pi] && i < (Length[T] - 2),
   p = FlatPart[\[Theta]0, \[Omega]0];
   \[Theta]0 = p[[1]];
   \[Omega]0 = p[[2]];
   AppendTo[\[CapitalTheta], \[Theta]0];
   AppendTo[\[CapitalOmega], \[Omega]0];
   t0 = t1;
   AppendTo[x, t0];
   t1 = t1 + 2 \[Delta];
   \[CapitalDelta]t = t1 - t0;
   p = OtherPart[\[Theta]0, \[Omega]0];
   \[Theta]0 = p[[1]];
   \[Omega]0 = p[[2]];
   AppendTo[\[CapitalTheta], \[Theta]0];
   AppendTo[\[CapitalOmega], \[Omega]0];
   t0 = t1;
   AppendTo[x, t0];
   t1 = T[[i]] - \[Delta];
   i = i + 1;
   ];

   ListLinePlot[Transpose[{x[[1 ;;]], \[CapitalTheta][[1 ;;]]}], 
   AxesLabel -> {"t", "\[Theta]"}, PlotLegends -> {i}, PlotRange -> All]]

Manipulate[
 mod[\[Theta]0, \[Omega]0, \[Gamma], \[CapitalDelta]\[CapitalOmega], \
\[Chi]], {\[Theta]0, 0, 1}, {\[Omega]0, 0, 1}, {{\[Gamma], 16}, 0.01, 
  50}, {{\[CapitalDelta]\[CapitalOmega], 0.3}, 0, 1}, {{\[Chi], 0.01},
   0, 0.05}]

My problem arises when I want to run this simulation many times so that I can do some statistics on the simulation (since the problem is stochastic it has a different outcome every time I run the code). I do this in a simple Do loop (as seen below) but maybe this can be done in a more efficient way? With this code I can only run the simulation around 130 times but I would like to be able to run it at least 1000 times.

sampl[\[Theta]0_, \[Omega]0_, \[Gamma]_, \[CapitalDelta]\
\[CapitalOmega]_, \[Chi]_, n_] := Module[{},
  Tend = {};
  Tnot = 0;
  Do[
   mod[\[Theta]0, \[Omega]0, \[Gamma], \
\[CapitalDelta]\[CapitalOmega], \[Chi]];
   If[i < (Length[T] - 2),
    AppendTo[Tend, t0], Tnot += 1
    ]
   
   , n];
  tmean = Mean[Tend];
  Histogram[Tend, {.025}]
  ]

sampl[0, 0, 16, 0.8, 0.01, 1000]

Thank you!

Answer 1

Okay, so there were two main bottlenecks in your code. First of all, the LineListPlot should be removed unless you actually want the plots for every iteration.

Secondly, you iterate over \[CapitalDelta]ts = RandomVariate[prob[\[Gamma]], 50];

which is very slow. It is better to generate all the random numbers you need first, and then iterate. I thus generate a list that is large enough for all iterations in the sampl[] function, and then I send a subset of 50 values to the mod[] function for every call.

I have also improved some other things

• Reap and Sow instead of AppendTo*
• changing constant numbers to non-constant
• avoid assigning values to variables when not necessary

however the performance gain from these changes is very small in comparison. Here is a faster solution that should do the same thing (please cross-check the calculations):

\[Delta] = 0;
prob[\[Gamma]_] := 
  ProbabilityDistribution[(\[Gamma] Exp[-\[Gamma] t]), {t, 
    0, \[Infinity]}];
\[CapitalOmega]0 = Sqrt[1 - (\[Chi]/2.)^2];
\[Theta][t_, \[Theta]0_, \[Omega]0_] := (\[Theta]0 Cos[
       t \[CapitalOmega]0] + ((\[Theta]0 \[Chi] + 2. \[Omega]0) Sin[
         t \[CapitalOmega]0])/(2. \[CapitalOmega]0)) E^(-((t \[Chi])/
        2.));
\[Omega][t_, \[Theta]0_, \[Omega]0_] := (E^(-((t \[Chi])/
          2.)) (4. \[Omega]0 \[CapitalOmega]0 Cos[
         t \[CapitalOmega]0] - (2. \[Chi] \[Omega]0 + \[Theta]0 (\
\[Chi]^2 + 4. \[CapitalOmega]0^2)) Sin[
         t \[CapitalOmega]0]))/(4. \[CapitalOmega]0);
FlatPart[\[Theta]0_, \[Omega]0_] := Module[{},
  Return[{\[Theta][t1, \[Theta]0, \[Omega]0], \[Omega][
     t1, \[Theta]0, \[Omega]0]}]]
OtherPart[\[Theta]0_, \[Omega]0_] := Module[{},
  Return[{\[Theta]0 + \[CapitalDelta]\[CapitalOmega] \
\[CapitalDelta]t, \[Omega]0 + \[CapitalDelta]\[CapitalOmega]}]]

mod[\[Theta]_, \[Omega]_, \[Gamma]_, \
\[CapitalDelta]\[CapitalOmega]1_, \[Chi]h_, dts_] := Module[{},
  T = Accumulate[dts];
  \[Chi] = \[Chi]h;
  \[CapitalDelta]\[CapitalOmega] = \[CapitalDelta]\[CapitalOmega]1;
  p = {\[Theta], \[Omega]};
  t0 = 0;
  i = 2;
  t1 = T[[1]] - \[Delta];
  res = Reap[
    While[p[[1]] < \[Pi] && i < (Length[T] - 2),
      p = FlatPart[p[[1]], p[[2]]];
      t0 = t1;
      Sow[{p[[1]], p[[2]], t0}];
      t1 = t1 + 2. \[Delta];
      \[CapitalDelta]t = t1 - t0;
      p = OtherPart[p[[1]], p[[2]]];
      t0 = t1;
      Sow[{p[[1]], p[[2]], t0}];
      t1 = T[[i]] - \[Delta];
      i = i + 1;];
    ];
  Return[res[[2, 1]]];
  ]

sampl[\[Theta]0_, \[Omega]0_, \[Gamma]_, \[CapitalDelta]\
\[CapitalOmega]_, \[Chi]_, n_] := Module[{}, Tend = {};
  Tnot = 0;
  dts = RandomVariate[prob[\[Gamma]], 50*n];
  Tend = Reap[
    Do[mod[\[Theta]0, \[Omega]0, \[Gamma], \[CapitalDelta]\
\[CapitalOmega], \[Chi], dts[[(j - 1)*50 + 1 ;; j*50]]];
      If[i < (Length[T] - 2), Sow[t0], Tnot += 1], {j, 1, n}];
    ];
  Tend = Tend[[2, 1]];
  tmean = Mean[Tend];
  Histogram[Tend, {.025}, 
   PlotLabel -> {": mean time to reach \[Theta]=\[Pi]" tmean}]
  ]

When calling sampl using the following code:

AbsoluteTiming[sampl[0, 0, 16, 0.8, 0.01, 100]] 

I had 2.7 seconds before removing ListLinePlot, 0.7 seconds before extracting the RandomVariate out from the mod[] function and finally ending at 0.04 seconds.

The code still requires quite a bit of cleaning I think. One of my main difficulties with it was that you are mixing global and local variables everywhere, and the path is not clear what depends on what and where the things are changed. I would suggest to really only define globally what should be global, outside of the loops, and then call functions with local variables. I did not attempt to sort this out because I am afraid that I would destroy some of the calculations, so I leave that to you.

* reap and sow might not be faster or only negligibly faster than AppendTo for small and simple lists. This seems to be the case here, but if you want to scale up, reap and sow is really the way to go.

Answer 2

The main culprit is not AppendTo", but the superfluous "ListLinePlot". Here is the improved version with "Reap" that take approx. 9 sec.

\[Delta] = 0;

prob[\[Gamma]_] := 
  ProbabilityDistribution[(\[Gamma] Exp[-\[Gamma] t]), {t, 
    0, \[Infinity]}];

\[CapitalOmega]0 = Sqrt[1 - (\[Chi]/2)^2];

\[Theta][t_, \[Theta]0_, \[Omega]0_] := (\[Theta]0 Cos[
       t \[CapitalOmega]0] + ((\[Theta]0 \[Chi] + 2 \[Omega]0) Sin[
         t \[CapitalOmega]0])/(2 \[CapitalOmega]0)) E^(-((t \[Chi])/
        2));

\[Omega][t_, \[Theta]0_, \[Omega]0_] := (E^(-((t \[Chi])/
          2)) (4 \[Omega]0 \[CapitalOmega]0 Cos[
         t \[CapitalOmega]0] - (2 \[Chi] \[Omega]0 + \[Theta]0 \
(\[Chi]^2 + 4 \[CapitalOmega]0^2)) Sin[
         t \[CapitalOmega]0]))/(4 \[CapitalOmega]0);

FlatPart[\[Theta]0_, \[Omega]0_] := 
 Module[{}, \[Theta]1 = \[Theta][t1, \[Theta]0, \[Omega]0];
  \[Omega]1 = \[Omega][t1, \[Theta]0, \[Omega]0];
  par = {\[Theta]1, \[Omega]1}]

OtherPart[\[Theta]0_, \[Omega]0_] := 
 Module[{}, \[Omega]1 = \[Omega]0 + \[CapitalDelta]\[CapitalOmega];
  \[Theta]1 = \[Theta]0 + \[CapitalDelta]\[CapitalOmega] \
\[CapitalDelta]t;
  par = {\[Theta]1, \[Omega]1}]

mod[\[Theta]_, \[Omega]_, \[Gamma]_, \
\[CapitalDelta]\[CapitalOmega]1_, \[Chi]h_] := 
 Module[{}, \[CapitalDelta]ts = RandomVariate[prob[\[Gamma]], 50];
  T = Accumulate[\[CapitalDelta]ts];
  \[Chi] = \[Chi]h;
  \[CapitalDelta]\[CapitalOmega] = \[CapitalDelta]\[CapitalOmega]1;
  \[Theta]0 = \[Theta];
  \[Omega]0 = \[Omega];
  t0 = 0;
  i = 2;
  t1 = T[[1]] - \[Delta];
  {\[CapitalTheta], \[CapitalOmega], x} = Reap[
     Sow[\[Theta]0, 1];
     Sow[\[Omega]0, 2];
     Sow[t0, 3];
     While[\[Theta]0 < \[Pi] && i < (Length[T] - 2), 
      p = FlatPart[\[Theta]0, \[Omega]0];
      \[Theta]0 = p[[1]];
      \[Omega]0 = p[[2]];
      Sow[\[Theta]0, 1];
      Sow[\[Omega]0, 2];
      t0 = t1;
      Sow[t0, 3];
      t1 = t1 + 2 \[Delta];
      \[CapitalDelta]t = t1 - t0;
      p = OtherPart[\[Theta]0, \[Omega]0];
      \[Theta]0 = p[[1]];
      \[Omega]0 = p[[2]];
      Sow[\[Theta]0, 1];
      Sow[\[Omega]0, 2];
      t0 = t1;
      Sow[t0, 3];
      t1 = T[[i]] - \[Delta];
      i = i + 1;];
     ][[2]];
  ]

sampl[\[Theta]0_, \[Omega]0_, \[Gamma]_, \[CapitalDelta]\[CapitalOmega]_, \[Chi]_, n_] := Module[{}, Tend = {};
  Tnot = 0;
  Tend = Reap[
     Do[mod[\[Theta]0, \[Omega]0, \[Gamma], \[CapitalDelta]\[CapitalOmega], \[Chi]];
       If[i < (Length[T] - 2), Sow[t0], Tnot += 1], n];
     ][[2]];
  tmean = Mean[Tend];
  Histogram[Tend, {.025}]];

sampl[0, 0, 16, 0.8, 0.01, 1000] // Timing