Efficient Langevin Equation Solver

Question

This question is not about good algorithms for solving stochastic differential equations. It is about how to implement simple codes in Mathematica efficiently exploiting Mathematica's programming methodology. (Hopefully, this may be useful in a stochastic processes course, for instance).

A simple Langevin Eq. in a single random variable $X$ with additive noise reads \begin{equation} \dot{X} = f(X) + \zeta(t) \end{equation} where $f(X)$ is an arbitrary function and $\zeta(t)$ is a Gaussian white noise satisfying \begin{equation} E(\zeta(t)) = 0, \qquad \text{and} \qquad E(\zeta(t) \zeta(t')) = \Gamma \delta(t-t') \end{equation}

To solve it we discretize time as $t = n dt$ and write \begin{equation} X_{n+1} = X_{n} + f(X_n)dt + \sqrt{\Gamma dt}\xi_n \end{equation} where $\xi_n \sim N(0,1)$.

Here is my best implementation thus far:

Langevin[x0_, f_, G_, tf_, n_, m_: 1] := 
    With[{dt = N[tf/n], s = N[Sqrt[ tf G/n]], xx0 = Table[x0, {m}]}, 
    Transpose@NestList[ # + dt f[#] + RandomVariate[NormalDistribution[0, s], m] &, xx0, n]];

It takes as input a initial condition $x_0$, a function $f[x]$, the spectral density $\Gamma$ (here written as $G$), the final integration time $t_f$ and the number of integration points $n$. The time step is then $dt = t_f/n$. It also takes an optional argument $m$ corresponding to the number of realisations.

The output consists of $m$ vectors $(X_0, X_1, X_2, \ldots,X_n )$ representing the stochastic realisations.

Here is this program applied to the famous bi-stable potential given by $V(x) = -\frac{x^2}{2} + \frac{x^4}{4}$, so that $f(x) = - V'(x) = x-x^3$. It simulates a cold ($\Gamma=0.1$, in data1) and a hot ($\Gamma=1$, in data2) condition:

First@AbsoluteTiming[
    data1 = Langevin[0, -#^3 + # &, 0.1, 10, 10^3, 2000];
    data2 = Langevin[0, -#^3 + # &, 1, 10, 10^3, 2000];]
0.317665

To analyse the steady state I discard some initial points (80% in this example). This shows how the particle remains distributed close to the potential minima when it's cold, but spread out when it's hot:

Show[
  Histogram[{Flatten[data1[[All, 800 ;; 1000]]], Flatten[data2[[All, 800 ;; 1000]]]}, Automatic, "PDF"], 
  Plot[-z^2/2 + z^4/4, {z, -1.8, 1.8}, PlotStyle -> Red], 
  AxesOrigin -> {0, 0}, PlotRange -> {-0.3, 1.2}]

Now to the questions:

1. Any immediate improvements on this function?
2. Is there a better way through a different approach
3. Can I compile this function as is to gain speed?
4. What about parallelisation?

A follow up would be to extend all this to systems of Langevin equations, replacing $X$, $f$ and $\zeta$ by vector valued functions. But, then, we loose the advantage of computing many realisations at once within the same NestList. I'll think more about this problem and if I come up with any ideas I'll update the question.

Thank you all in advance and I hope this may be of use to other researchers as well.

Note: here is an example using the idea of @R.M.: generate all random numbers at once and use an index through the iteration to move along:

LangevinBad[x0_, f_, G_, tf_, n_, m_: 1] := Block[{i = 1},  
   With[{dt = N[tf/n], r = RandomVariate[ NormalDistribution[0, N[Sqrt[ tf G/n]]], {n, m}], xx0 = Table[x0, {m}]}, 
   Transpose@NestList[ # + dt f[#] + r[[i++]] &, xx0, n]]];

Maybe my coding is no good, but this version is really bad. Actually; Nest probably has an internal variable to keep track of what iteration step it is, but I have no idea if it is possible to access that.

---

ACL's version

@ACL came up with a really efficient code, which I copy here for completeness.

(* This was originally called l4 by ACL *)

LangevinACL[fn_] := With[{f = fn}, 
    Compile[{{x0, _Real}, {G, _Real}, {tf, _Real}, {n, _Integer}}, 
    Module[{dt, s, state, r}, dt = N[tf/n];
    s = N[Sqrt[tf G/n]];
    state = ConstantArray[0., n];
    state[[1]] = x0;
    r = RandomVariate[NormalDistribution[0, s], n];
    Do[state[[nc]] = state[[nc - 1]] + dt*f@state[[nc - 1]] + r[[nc - 1]], {nc, 2, n}];
state], CompilationTarget -> "C"]]

Then to compile for a given function use

ll = LangevinACL[(# - #^3) &];
AbsoluteTiming[dat = Table[ll[0, .1, 10, 10^3], {2000}];]

This code is always faster then the originally posted and allows for easy parallelisation.

---

Vector Equations

In vector equations there are two possibilities; either all particles have the same fluctuating properties, in which case we usually write $E(\zeta_i(t) \zeta_j(t')) = \Gamma\delta_{i,j} \delta(t-t')$ for the components of the fluctuating vector; or each particle has a specific fluctuation: $E(\zeta_i(t) \zeta_j(t')) = \Gamma_{i,j} \delta(t-t')$, where $\Gamma_{i,j}$ are the entries of a covariance matrix.

Here are two implementations of the former (all equations with the same fluctuation).

The first is a simple variation of the original code, as suggested by @ACL, so that instead of computing several realisations at once, each function call evaluates only a single realisation, but for a vector system:

LangevinVec[x0_, f_, G_, tf_, n_] := 
  With[{dt = N[tf/n], s = N[Sqrt[ tf G/n]], m = Length@x0}, 
  NestList[ # + dt f[#] + RandomVariate[NormalDistribution[0, s], m] &, x0, n]]; 

Everything is exactly as in Langevin, except that on input $x_0$ should be an array of numbers. Note also that there is no failsafe to check if the function $f$ has the correct dimensionality! (it should be a mapping from $\mathbb{R}^m\rightarrow\mathbb{R}^m$, where $m$ is the length of $x_0$).

The second implementation is again motivated by @ACL's code:

LangevinVecACL[fn_] := 
  With[{f = fn}, 
  Compile[{{x0, _Real, 1}, {G, _Real}, {tf, _Real}, {n, _Integer}}, 
  Module[{dt, s, state, r, m},
  m = Length@x0;
  dt = N[tf/n];
  s = N[Sqrt[tf G/n]];
  state = ConstantArray[0., {n, m}];
  state[[1]] = x0;
  r = RandomVariate[NormalDistribution[0, s], {n, m}];
  Do[state[[nc]] = state[[nc - 1]] + dt*f@state[[nc - 1]] + r[[nc - 1]], {nc, 2, 
  n}];
  state], CompilationTarget -> "C"]]

Now to applications. Here is a model of ferromagnetism reminiscent of the 1D Ising system. There are $m$ random variables in $\vec{x} = (x_1,x_2,\ldots,x_m)$ representing spins in a linear chain of atoms. The interaction potential is given by

\begin{equation} V(\vec{x}) = - \sum_{i=1}^m (\frac{a x_i^2}{2} - \frac{b x_i^4}{4}) - c \sum_{i=1}^m x_i x_{i+1} \end{equation} This refers to a bi-stable potential (as in the previous example) for each variable representing the magnetic order, plus a harmonic-type interaction between them. The corresponding force is \begin{equation} f_i = a x_i - bx_i^3 + c(x_{i-1}+x_{i+1}) \end{equation}

In matrix notation I can write \begin{equation} f(\vec{x}) = A\vec{x} - b \vec{x}^3 \end{equation} where $\vec{x}^3$ stands for $(x_1^3,x_2^3,\ldots)$ and $A$ is an $m\times m$ tridiagonal matrix with of the form \begin{equation} A = \left( \begin{array}{ccccc} a & c & 0 & 0 & c \\ c & a & c & 0 & 0 \\ 0 & c & a & c & 0 \\ 0 & 0 & c & a & c \\ c & 0 & 0 & c & a \end{array} \right) \end{equation} Note that I am using periodic boundary conditions $x_{m+1}=x_1$ and thence the c's in the upper-right and lower-left corners.

Here is $f(x)$ in Mathematica

m = 100; a = 2.0; b = 3.0; c = 3;
A = SparseArray[{
    {m, 1} -> c, {1, m} -> c, 
    Band[{1, 1}] -> a, Band[{2, 1}] -> c, Band[{1, 2}] -> c}, 
    {m, m}];
f[x_] := A.x - b x^3

The choice of parameters is somewhat arbitrary and perhaps this definition of $f(x)$ is not the fastest due to the dot product.

I will use ACL's version LangevinVecACL, which is faster. So I first compile it

llvec = LangevinVecACL[f];

Here are two data sets for $\Gamma = 0.01$ (pretty cold) and $\Gamma = 10$ (pretty hot).

x0 = ConstantArray[0.0, m];
AbsoluteTiming[
   data1 = llvec[x0, .01, 4, 10^4];
   data2 = llvec[x0, 10, 4, 10^4];
]

The following code shows the steady-state distribution of a single realisation

GraphicsGrid[{Map[
  ListPlot[#, PlotRange -> {{-1, m + 1}, {Floor@Min@#, Ceiling@Max@#}}, 
    Filling -> Axis, Frame -> True, BaseStyle -> 14, 
    FrameLabel -> {"Position", "Magnetization"}] &,
  {Last@data1, Last@data2}]}, ImageSize -> {600}]

As can be seen, at cold temperatures the system tends to divide itself into domains with all spins chunked either "up" or "down"; conversely, at high temperatures the domain configuration is clearly degraded.

The following function animates the time evolution of the system.

animateSpinChain[data_] :=  Animate[ListPlot[data[[i]], 
   PlotRange -> {{-1, m + 1}, {Floor[Min[data]], Ceiling[Max[data]]}},
   Filling -> Axis], {i, 1, Length@data, Floor[Length@data/100]}]

Answer 1

The first thing I'd try is to compile this, and also split it up so each realization is done independently (to allow easy parallelization). Taking the mexican hat potential you mentioned in the question (see later for general potential):

    l3 = Compile[
       {
        {x0, _Real},
        {G, _Real},
        {tf, _Real},
        {n, _Integer}
        },
       Module[{dt, s, state, r}, dt = N[tf/n];
        s = N[Sqrt[tf G/n]];
        state = ConstantArray[0., n];
        state\[LeftDoubleBracket]1\[RightDoubleBracket] = x0;
        r = RandomVariate[NormalDistribution[0, s], n];
        Do[
         state\[LeftDoubleBracket]nc\[RightDoubleBracket] = 
          state\[LeftDoubleBracket]nc - 1\[RightDoubleBracket] + 
              dt*(# - #^3) &@
            state\[LeftDoubleBracket]nc - 1\[RightDoubleBracket] + 
           r\[LeftDoubleBracket]nc - 1\[RightDoubleBracket],
         {nc, 2, n}
         ];
        state
        ],
       CompilationTarget \[Rule] "C"
       ];

(just paste it and it will look better). Note that I produce the noise once, put it in a list, and then just access each element as needed, to avoid the overhead of calling RandomVariate many times (I may have messed something up here, I haven't checked if the moments are the same as yours for the same potential, but I think it's OK).

This seems slightly faster than the original version (without parallelization):

AbsoluteTiming[data1 = Langevin[0, -#^3 + # &, 0.1, 10, 10^3, 2000];]
(*0.245488*)

dat = Table[l3[1, .1, 10, 10^3], {2000}]; // AbsoluteTiming
(*{0.153759, Null}*)

and it can be easily parallelized (but I have not tried as my laptop has only two cores; you'd need a lot of cores for this to be worth it).

Generalizing to discretised partial SDEs (or systems of Langevin equations) is straightforard with this approach (I could provide code if you want). Obviously, you sacrifice the generality of accepting any f to be able to compile, which is a disadvantage.

Finally, I would suggest you avoid NestList for this. It's clean, but if you want to do this for a largish system of SDEs, and want eg 10^6 steps with eg 10^4 realizations, there's no way you'll be able to produce them at once, as you do. It simply doesn't scale. It's also more difficult to keep only the last N results, if you want to wait some time for the system to equilibrate.

EDIT: To compile this with arbitrary potential f, use

l4[fn_] := With[
  {f = fn},
  Compile[
   {
    {x0, _Real},
    {G, _Real},
    {tf, _Real},
    {n, _Integer}
    },
   Module[{dt, s, state, r}, dt = N[tf/n];
    s = N[Sqrt[tf G/n]];
    state = ConstantArray[0., n];
    state\[LeftDoubleBracket]1\[RightDoubleBracket] = x0;
    r = RandomVariate[NormalDistribution[0, s], n];
    Do[
     state\[LeftDoubleBracket]nc\[RightDoubleBracket] = 
      state\[LeftDoubleBracket]nc - 1\[RightDoubleBracket] + 
       dt*f@state\[LeftDoubleBracket]nc - 1\[RightDoubleBracket] + 
       r\[LeftDoubleBracket]nc - 1\[RightDoubleBracket],
     {nc, 2, n}
     ];
    state
    ],
   CompilationTarget \[Rule] "C"
   ]
  ]

and then

ll = l4[(# - #^3) &]
AbsoluteTiming[dat = Table[ll[0, .1, 10, 10^3], {2000}];]
(*{0.138822, Null}*)

However, Nest tries to automatically compile its first argument if it is nested over a number of times (100 by default) visible in SystemOptions[CompileOptions], so this may not be worth much effort for runs as small as these.

Answer 2

I get a very small, but consistent improvement in time to calculate data1 and data2 by replacing With with Module:

Langevin[x0_, f_, G_, tf_, n_, m_: 1] := Module[{dt, s, xx0},
   dt = N[tf/n];
   s = N[Sqrt[tf G/n]];
   xx0 = Table[x0, {m}]; 
   Transpose@
    NestList[# + dt f[#] + 
       RandomVariate[NormalDistribution[0, s], m] &, xx0, n]];

First@AbsoluteTiming[
  data1 = Langevin[0, -#^3 + # &, 0.1, 10, 10^3, 2000];
  data2 = Langevin[0, -#^3 + # &, 1    , 10, 10^3, 2000];]

but on my machine it only gives the difference between: 0.493533 and 0.436566

You can speed up the generation of the histograms using ParallelMap:

AbsoluteTiming[
 histograms = 
  ParallelMap[
   Histogram[#, Automatic, "PDF"] &, {Flatten[data1[[All, 800 ;; 1000]]],
    Flatten[data2[[All, 800 ;; 1000]]]}];
 Show[
  histograms[[1]],
  histograms[[2]],
  Plot[-z^2/2 + z^4/4, {z, -1.8, 1.8},
   PlotStyle -> Red],
  AxesOrigin -> {0, 0},
  PlotRange -> {-0.3, 1.2}
  ]]

but I don't think 3 or more than kernels will help over 2 (I'd love to be wrong on this).

I wonder if pre-calculating all of your RandomVariates might gain you a bit. I'll try something and update this answer if I can block out some time.

...

Responding to the OP's comment below and given that I may have missed what they meant in the comment did you mean something like the following when you wrote:

I thought about using a single RandomVariate. But I couldn't really figure out how to efficiently Nest that.:

Langevin2[x0_, f_, G_, tf_, n_, m_: 1] := Module[{dt, s, xx0, r},
   dt = N[tf/n];
   s = N[Sqrt[tf G/n]];
   xx0 = Table[x0, {m}];
   r = RandomVariate[NormalDistribution[0, s], m];
   Transpose[
    NestList[
     # + dt f[#] + r &,
     xx0,
     n]
    ]
   ];

Faster, but given your comment and the original question I'm not certain it does what you need.