Efficient Langevin Equation Solver

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 $$\dot{X} = f(X) + \zeta(t)$$ where $f(X)$ is an arbitrary function and $\zeta(t)$ is a Gaussian white noise satisfying $$E(\zeta(t)) = 0, \qquad \text{and} \qquad E(\zeta(t) \zeta(t')) = \Gamma \delta(t-t')$$

To solve it we discretize time as $t = n dt$ and write $$X_{n+1} = X_{n} + f(X_n)dt + \sqrt{\Gamma dt}\xi_n$$ 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?

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

$$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}$$ 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 $$f_i = a x_i - bx_i^3 + c(x_{i-1}+x_{i+1})$$

In matrix notation I can write $$f(\vec{x}) = A\vec{x} - b \vec{x}^3$$ 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 $$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)$$ 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]}]

-
That's a very nicely written question. Wish more were like this... –  rm -rf Jul 26 '12 at 16:17
An immediate generalization to a system of equations would be nDLangevin[x0_, f_, covMat_, tf_, n_, m_: 1] := With[{nDim = Length[x0], mean = ConstantArray[0, Length[x0]], dt = N[tf/n], xx0 = Table[x0, {m}], nDf = Function[x, f[#] & /@ x]}, Transpose@NestList[# + dt nDf[#] + RandomVariate[MultinormalDistribution[mean, covMat], m] &, xx0, n]] but it's terribly slow. –  b.gatessucks Jul 26 '12 at 16:49
I don't think it can be parallelized, because each step depends on the previous step. Parallelization works well when the steps are independent of each other (e.g. Map, Do, etc.). I think your NestList approach is very clean and efficient. An equivalent formulation would be using memoization and recursive functions, but that is about 2x slower in my tests. All the functions used are compilable, but I'm not sure how to handle the case of arbitrary f. If f is known in advance, then you could easily compile it for that f. –  rm -rf Jul 26 '12 at 17:02
@R.M it can be parallelized if you calculate each realization independently (I ended up doing this when I was working on stochastic processes a few years back) –  acl Jul 26 '12 at 17:17
You can use JIT-compilation, but my tests show that this only improves speed by about 10-15 percents. The running time seems to be dominated by the random number generation, which, for your example, is done in chunks large enough so that compilation does not bring much speed improvement. So, your function is pretty efficient, both because you generate random numbers in large enough chunks, and because NestList auto-compiles. It may make more sense to compile it if, for example, you would use it with small m argument in a loop. –  Leonid Shifrin Jul 26 '12 at 17:29

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.

-
quizz: what package was used for figs 2 and 3 of this? –  acl Jul 26 '12 at 18:15
acl, your timing is the same as mine i Think. In the first one you compute two simulations (data1 and data2) and in the second one you compute only one. Both have similar run times, and given the simplicity of NestList + the ability to accept any function, I don't really see much advantage. –  Gabriel Landi Jul 26 '12 at 20:04
Yes you're right, I screwed up. Let me fix it later. The only advantage is that it is trivially parallelizable and it scales much easier (at least, for me). Try to solve the discretised KPZ equation (which is notorious for slow convergence) for large systems and thousands of runs and you will see what I mean. For obtaining intuition, though, I agree that the Nest approach is neater. In practice, anyway, for something as simple algorithmically as this, C is much faster and not very hard to code if performance is needed. –  acl Jul 26 '12 at 20:21
@acl is there a prize...? :-) –  sebhofer Jul 26 '12 at 20:47
@sebhofer a copy of the package :) –  acl Jul 26 '12 at 20:51

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.

-
Hi Jagra. I thought about using a single RandomVariate. But I couldn't really figure out how to efficiently Nest that. –  Gabriel Landi Jul 26 '12 at 17:36
I don't think so. This only generates the random numbers once. $r$ must be re-computed within each iteration. Ideally this would be done with r = RandomVariate[NormalDistribution[0,s],{m,n}]. But then you can't really nest this matrix; or at least I don't really know how. –  Gabriel Landi Jul 26 '12 at 17:54
Why not r := RandomVariate[NormalDistribution[0, s], m] ? –  b.gatessucks Jul 26 '12 at 17:59
@GabrielLandi You can probably use Fold then, and run an index through the pre-generated RandomVariate... –  rm -rf Jul 26 '12 at 18:06
@GabrielLandi -- Have you thought about using FoldList and then used something like r = RandomVariate[NormalDistribution[0, s], {m,m}]? –  Jagra Jul 26 '12 at 18:18