Integrators

本文档介绍了关于ordinary differential equations(ODEs) 和stochastic differential equations(SDEs) 常用的integrators。本文档借鉴和照搬了去年的demo list里面的许多内容，这里特别感谢去年的作者们！ 本文档主要以英文介绍为主（因为作者不太确定很多英文的名词的中文对应），希望各位能提供宝贵意见。

作者：花勐健以及去年demo list的作者们

审核：

1. 导入依赖

本文档使用的镜像为dmff_0.2.0-notebook。首先我们拉取DMFF最新版本代码，并安装一些必要的依赖。

可以通过import dmff的命令验证dmff已安装成功，在CPU节点上可能会看到一些warning信息，这是正常的。

2. ODE Integrators

2.1 Euler's method

Many commonly seen deterministic systems are governed by Newton's third law. An object's velocity and position as functions of time can be described by the following ordinary differential equations: $$\frac{d}{dt}{r}_t=v_t$$ $$\frac{d}{dt}{v}_t=a(t)=F(t)/m$$ where $r_t$ denotes the position of the object, $v_t$ is its time-dependent velocity, and $a(t)=F(t)/m$ is the acceleration calculated from Newton's third law with a time-dependent force $F(t)$ and a mass $m$.

Suppose we have a grid in time $t_0,t_1,\cdots ,t_n$, where we denote the mesh width/time step size by $\Delta t$. The Euler's method is the simplest way to integrate such ODE systems and the discretization is as follows:

$$r_{k+1}=r_k+v_k\Delta t$$ $$v_{k+1}=v_k+a(t_k)\Delta t$$

where $r_k,v_k$ are the position and velocity at time $t_k$.

With a Taylor expansion, we can show that the numerical error brought by this way of discretization in time is $O(\Delta t)$.

Moreover, it is worth noting that the Euler's method, and more generally, all explicit ODE integrators suffer from the problem of numerical stability. Here's a simple example that demonstrate the issue of numerical stability. Suppose we have an ODE system:

$$\frac{d}{dt}{y}_t=c{y}_t$$

where $c<0$ is a constant. The analytical solution of this ODE is $y_t=\exp (ct)y_0$, where $y_0$ is the initial condition. The numerical solution given by the Euler's method is $y_{k+1}=y_k+c\Delta t{y}_k=(1+c\Delta t)y_k$. It follows that $y_n=(1+c\Delta t{)}^n{y}_0$. Since $c<0$, $1+c\Delta t<1$. What if $\Delta t$ is so large that $1+c\Delta t<-1$? It is obvious that $y_n$ would blow up in this case. The implicit Euler's method would help avoid this issue. Let's take the same ODE $\frac{d}{dt}{y}_t=c{y}_t$ as the example. The discretization of the implicit Euler's method is also very simple:

$$y_{k+1}=y_k+c\Delta t{y}_{k+1}$$

We don't know the right-hand side but we can work out an equivalent expression:

$$y_{k+1}=(1-c\Delta t{)}^{-1}{y}_k$$

We observe that the numerical stability is no longer a problem for the implicit Euler's method because $(1-c\Delta t{)}^{-1}\in (0,1)$ as long as $\Delta t>0$.

In the block below, we implement the explicit Euler's method for general 1D systems governed by Newton's third law.

2.2 verlet method and velocity-verlet method

Since the Euler's method is only first-order accurate (i.e. the error ~ $O(\Delta t)$), we often want to use a high-order accuracy method such that we can use a larger time step size $\Delta t$ without loss of accuracy. There are many methods of such (e.g. Runge-Kutta methods, multistep methods). If the readers want to know more about the families of high-order methods, we recommend this book: https://faculty.washington.edu/rjl/fdmbook/. Here, We introduce the verlet and the velocity-verlet method, which are more relevant to MD problems. Now, we consider our system of ODEs governed by Newton's third law:

$$\frac{d}{dt}{r}_t=v_t$$ $$\frac{d}{dt}{v}_t=a(t)=F(t)/m$$

For the time step $n$, the Taylor expansion of $r_{n+1}$ is $$r_{n+1}=r_n+\Delta t\frac{d}{dt}{r}_n+\frac{1}{2}(\Delta t{)}^2\frac{d^2}{d{t}^2}{r}_n+\frac{1}{6}(\Delta t{)}^3\frac{d^3}{d{t}^3}{r}_n+O(\Delta t^4)$$

From our ODEs, we know that $\frac{d}{dt}{r}_n=v_n$, $\frac{d^2}{d{t}^2}{r}_n=\frac{d}{dt}{v}_n=a(t_n)$, and $\frac{d^3}{d{t}^3}{r}_n=a^{\prime }(t_n)$. Hence, $$r_{n+1}=r_n+\Delta t{v}_n+\frac{1}{2}a(t_n)(\Delta t{)}^2+\frac{1}{6}(\Delta t{)}^3{a}^{\prime }(t_n)+O(\Delta t^4)$$

Similarly, if we expand $r_{n-1}$ with Taylor series, we would get

$$r_{n-1}=r_n-\Delta t{v}_n+\frac{1}{2}a(t_n)(\Delta t{)}^2-\frac{1}{6}(\Delta t{)}^3{a}^{\prime }(t_n)+O(\Delta t^4)$$

Therefore, adding these two Taylor expansions yields

$$r_{n-1}+r_{n+1}=2r_n+a(t_n)(\Delta t{)}^2+O(\Delta t^4)$$

and this gives us a numerical scheme, which people often call the verlet method,

$$r_{n+1}=2r_n-r_{n-1}+a(t_n)(\Delta t{)}^2$$

which is third-order accurate. Note that the velocity-verlet method does not require we update the velocity at every time step and we do not need to keep track of $v_1,v_2,\cdots ,v_n$. We can compute $r_{n+1}$ directly as long as we have access to $r_n,r_{n-1}$ and we can evaluate $a(t_n)$.

As opposed to the verlet method, we also have the velocity-verlet method, which requires us to update the velocity at every step. The scheme is as follows:

$$r_{n+1}=r_n+v_n\Delta t+\frac{1}{2}a(t_n)(\Delta t{)}^2$$ $$v_{n+1}=v_n+\frac{(a(t_n)+a(t_{n+1}))\Delta t}{2}$$

The velocity is updated via a central difference scheme with error $O(\Delta t^2)$. So, this error will be become $O(\Delta t^3)$ when it comes to the update of the position. Therefore, the position is updated with a local truncation error of $O(\Delta t^3)$.

Both the verlet and the velocity-verlet can achieve better numerical accuracy than the simple explicit and implicit Euler methods. But we should also note here that the time step size for the verlet and the velocity-verlet methods is also restricted by the numerical stability but they both provide much better numerical stability than the explicit Euler's method. So, in practice, they are the methods that people often use.

2.3 Numerical Example

In the below blocks, we use a simple 1D harmonic oscillator example to illustrate the performance of the methods we introduced above. The energy of the system we are considering is obviously conserved but it is not conserved in the numerical solutions because of discretization errors. Here, we do a convergence study on this. Let our time step sizes be $t_k=2^{-2-k}$ and the potential energy we get from numerical solver integrating till $t=1$ with time step size $t_k$ be $E_k$. The potential energy $r_k$ only depends on the position $r_k$, so we use it instead of the total energy to measure the convergence.

Then, the relative error is defined by $|E_{k+1}-E_k|$ and a plot of which can indicate the speed of convergence of the numerical solver. We make a plot of such for the three integrators we introduced above for the 1D harmonic oscillator and we see that the verlet method and the velocity-verlet method achieve almost the same accuracy and they are all much better than the forward Euler method.

Here, we should emphasize that the actual error in the numerical solution at the final time (i.e. the global error) is often different in order from the local truncation error, which is what we analyzed from Taylor expansions. In practice, both the verlet and the velocity-verlet methods have a global error of $O(\Delta t^2)$. This difference is due to the fact that we have assumed in our Taylor expansions that $\frac{d}{dt}r(t_n)=v_n$ and $\frac{d}{dt}v(t_n)=a(x_n)$, which only hold when $n=0$ (i.e. initial conditions). The error accumulates in time and this makes these relations no longer hold. This also explains why the verlet method and the velocity-verlet method achieve almost the same accuracy in the end whereas our analysis shows that the verlet method has a local truncation error of approximately $O(\Delta t^4)$ in the position whereas the velocity-verlet should have a local truncation error of approximately $O(\Delta t^3)$ in the position.

When we perform the analysis, we consider $a(t)$ as a function of $t$, but in this example, we have $a(t)=a(r_t)$. With this difference, the third-order term in the Taylor expansion is now

$$\frac{1}{6}(\Delta t{)}^3\frac{d}{dt}a(x(t_n))=\frac{1}{6}(\Delta t{)}^3{a}^{\prime }(x(t_n))\frac{d}{dt}x(t_n)\ne \frac{1}{6}(\Delta t{)}^3{a}^{\prime }(x_n)v_n$$

2. Simple SDE Integrators

We will put more emphasis on the following sections on SDE integrators rather than the ODE integrators because we more often deal with SDEs instead of ODEs in MD simulations. Suppose we have an one-dimensional SDE as follows:

$$d{X}_t=b(t,X_t)dt+\sigma (t,X_t)d{W}_t$$

where $b(t,X_t)$ is often called the drift term, $\sigma (t,X_t)$ is often called the diffusion term, and $W_t$ denotes a standard Brownian motion/Wiener measure.

We often write stochastic processes as the form of SDEs but they are only defined as the short-hand for the integral equation

$$X_t=X_0+{\int }_0^{t}b(t,X_s)ds+{\int }_0^t\sigma (s,X_s)d{W}_s$$

The integral ${\int }_0^t\sigma (s,X_s)d{W}_s$ is defined as the Ito integral. Unlike ODEs, the SDE themselves do not have a rigorous mathematical definition because the Brownian motion is almost surely nowhere differentiable. Therefore, $d{W}_t$ is undefined as a mathematical object.

Now, we are going to introduce several SDE integrators, including the Euler-Maruyama scheme, which is the counterpart of the Euler method we introduced as an ODE integrator. We will also introduce the Milstein scheme, which is less frequently used than the Euler-Maruyama scheme.

2.1 Euler-Maruyama Scheme

The Euler-Maruyama scheme is very simple. Suppose we know $X_t$ at step n (i.e. time $t_n$), which we denote by $X_n$, then

$$X_{n+1}=X_n+b(t_n,X_n)\Delta t+\sigma (t_n,X_n)\xi$$

where $\Delta t$ denotes the time step size and $\xi \sim \mathcal{N}(0,\Delta t)$ is a sample from the Gaussian distribution with mean $0$ and variance $\Delta t$. This can be derived from the stochastic Ito-Taylor expansion.

2.2 Milstein Scheme

Using the notation in the previous section, we state the Milstein scheme, which is less stable but more accurate than the Euler-Maruyama scheme, for the time-homogeneous SDEs (i.e. $b,\sigma$ only depend on $X_t$)as follows:

$$X_{n+1}=X_n+b(X_n)\Delta t+\sigma (X_n)\xi +\frac{1}{2}\sigma (X_n){\sigma }^{\prime }(X_n)({\xi }^2-\Delta t)$$

where $\Delta t$ denotes the time step size and $\xi \sim \mathcal{N}(0,\Delta t)$ is a sample from the Gaussian distribution with mean $0$ and variance $\Delta t$. Notice that the first two terms on the right-hand side coincide with the Euler-Maruyama scheme. Here the last term on the right-hand side is a correction term, which can be derived from the stochastic Ito-Taylor expansion.

3. Langevin Dynamics & Its numerical Integrators

We consider the Hamiltonian system governed by the following equations:

$$d\mathbf{q}=M^{-1}\mathbf{p}dt$$

$$d\mathbf{p}=(-\nabla U(\mathbf{q})-\gamma \mathbf{p})dt+\sqrt{2k_{B}T\gamma M}d{\mathbf{W}}_t$$

where the bold characters are vector-valued functions. In the above equations, $\mathbf{q},\mathbf{p},U(\mathbf{q})$ denote the position, momentum, and the potential energy of a particle. $T$ denotes the temperature, $k_B$ denotes the Boltzmann constant, and $\gamma$ denotes the friction coefficient.

Question: by definition, what is the Hamiltonian of the system?

Answer: $$H(\mathbf{q},\mathbf{p})=U(\mathbf{q})+\frac{1}{2}{\mathbf{p}}^T{M}^{-1}\mathbf{p}$$

We can rewrite this system of equations as

$$d\left[\begin{array}{c}\mathbf{q}\\ \mathbf{p}\end{array}\right]=\left[\begin{array}{c}{M}^{-1}\mathbf{p}\\ 0\end{array}\right]dt+\left[\begin{array}{c}0\\ -U(\mathbf{q})\end{array}\right]dt+\left[\begin{array}{c}0\\ -\gamma \mathbf{p}dt+\sigma M^{1/2}d{\mathbf{W}}_t\end{array}\right]$$

where we let $\sigma =\sqrt{2k_{B}T\gamma }$ to simplify the notations.

We label the three terms on the right-hand side as $A=\left[\begin{array}{c}{M}^{-1}\mathbf{p}\\ 0\end{array}\right]dt$, $B=\left[\begin{array}{c}0\\ -U(\mathbf{q})\end{array}\right]dt$, and $O=\left[\begin{array}{c}0\\ -\gamma \mathbf{p}dt+\sigma M^{1/2}d{\mathbf{W}}_t\end{array}\right]$. Suppose we have the following three equations:

$$d\left[\begin{array}{c}\mathbf{q}\\ \mathbf{p}\end{array}\right]=A,d\left[\begin{array}{c}\mathbf{q}\\ \mathbf{p}\end{array}\right]=B,d\left[\begin{array}{c}\mathbf{q}\\ \mathbf{p}\end{array}\right]=O$$

From time step $n$ to $n+1$, to the best of our knowledge, these three equations can be solved for one step using the following update:

$$\left[\begin{array}{c}{\mathbf{q}}_{n+1}\\ {\mathbf{p}}_{n+1}\end{array}\right]=\left[\begin{array}{c}{\mathbf{q}}_n+M^{-1}{\mathbf{p}}_n\Delta t\\ {\mathbf{p}}_n\end{array}\right],\left[\begin{array}{c}{\mathbf{q}}_{n+1}\\ {\mathbf{p}}_{n+1}\end{array}\right]=\left[\begin{array}{c}{\mathbf{q}}_n\\ {\mathbf{p}}_n-U({\mathbf{q}}_n)\Delta t\end{array}\right],\left[\begin{array}{c}{\mathbf{q}}_{n+1}\\ {\mathbf{p}}_{n+1}\end{array}\right]=\left[\begin{array}{c}{\mathbf{q}}_n\\ e^{-\gamma \Delta t}{\mathbf{p}}_n+\sqrt{k_{B}T(1-e^{-2\gamma \Delta t})}{M}^{1/2}{\mathbf{R}}_n\end{array}\right]$$

where ${\mathbf{R}}_n$ is a Gaussian vector with mean $0$ and covariance $I$. Now, we have splitted the Langevin equations into three pieces and we know how to solve each one of them.

Since the Langevin equations can be written as

$$d\left[\begin{array}{c}\mathbf{q}\\ \mathbf{p}\end{array}\right]=A+B+O$$

we then know how to solve the Langevin equations with the above solutions to the three splitted equations and there are many numerical schemes we can come up with.

For example, we can first solve A for an half time step, then solve B with ${\mathbf{q}}_{n+1/2}$ obtained from solving A, then solve A again with ${\mathbf{p}}_{n+1/2}$ we get from solving B, and finally solve O with ${\mathbf{q}}_{n+1}$. This method is short-handed as ABAO. Similarly, one can invent many schemes of such, such as BABO, ABOBA, BAOAB. Writing out these schemes (e.g. BABO, ABOBA, BAOAB) will leave as an exercise to the readers.

As a side note, ABOBA and BAOAB are often referred to as the Position-Verlet method and the Velocity-Verlet method and they perserve the variance of position.

Exercise: Can you try to modify the code above and write down the OABAO integrator? Then, using the code of simulating the one-dimensional harmonic oscillator provided below to test your integrator by finding an invariant/preserved quantity and monitoring whether it is invariant/preserved during the simulation.

4. Overdamped Langevin Dynamics

An interesting special case of Langevin dynamics is obtained by considering the large $\gamma$ limit, which is often referred to as the overdamped limit. The resulting equation is called the overdamped Langevin dynamics. In the limit of large $\gamma$, the change in $\mathbf{p}$ damps very quickly and we can assume that the acceleration is negligible, which yields the following system of equations:

$$d\mathbf{q}=M^{-1}\mathbf{p}dt$$

$$0=(-\nabla U(\mathbf{q})-\gamma \mathbf{p})dt+\sqrt{2k_{B}T\gamma M}d{\mathbf{W}}_t$$

From the second equation, we obtain that

$$\mathbf{p}dt=-{\gamma }^{-1}\nabla U(\mathbf{q})dt+\sqrt{2k_{B}TM{\gamma }^{-1}}d{\mathbf{W}}_t$$

We combine it with the first equation and we get

$$d\mathbf{q}=-{\gamma }^{-1}{M}^{-1}\nabla U(\mathbf{q})dt+\sqrt{2k_{B}T{M}^{-1}{\gamma }^{-1}}d{\mathbf{W}}_t$$

Eliminating $\gamma$ and $M$ by rescaling the equation yields

$$d\mathbf{q}=-\nabla U(\mathbf{q})dt+\sqrt{2k_{B}T/M}d{\mathbf{W}}_t$$

This is also sometimes referred to as Brownian dynamics. The overdamped Langevin dynamics have an equilibriunm Gibbs distribution $\rho (\mathbf{q})\propto \exp (-U(\mathbf{q})/(k_{B}T))$, which is also referred to as the canonical ensemble, and people often integrate the above SDE to sample from $\rho (\mathbf{q})$. This sampling technique is called Langevin dynamics in the machine learning community and it is widely used to sample from energy-based models and diffusion models. Another less commonly used name for this method is called the stochastic gradient descent algorithm.

5. Nose-Hoover Dynamics

Nose-Hoover dynamics provide an alternative to the Langevin dynamics for sampling from the canonical measure $\rho (\mathbf{p},\mathbf{q})\propto \exp (-H(\mathbf{p},\mathbf{q})/(k_{B}T))$, where $H(\mathbf{p},\mathbf{q})$ denotes the system Hamiltonian, with deterministic paths. For any test function $\varphi (\mathbf{p},\mathbf{q})$, we want to compute the integral

$${\int }_{\mathcal{D}}\varphi (\mathbf{p},\mathbf{q})\rho (\mathbf{p},\mathbf{q})d{\omega }_{\mathbf{q}}{\omega }_{\mathbf{p}}$$

where $\mathcal{D}$ denotes the domain and ${\omega }_{\mathbf{q}},{\omega }_{\mathbf{p}}$ are the coordinates of $\mathbf{q}$ and $\mathbf{p}$. One way to compute this integral in an extended space by introducing new variables $\xi$ in addition to $(q,p)$ and additional equations of motion to drive them. The combined system is carefully designed to preserve an extended invariant distribution $$\tilde{\rho }(\mathbf{q},\mathbf{p},\xi )\propto \rho (\mathbf{q},\mathbf{p})\times \hat{\rho }(\xi )$$

where the auxiliary density $\hat{\rho }(\xi )$ has a simple form (e.g. a multivariate Gaussian). Then, it follows that

$$\int {\int }_{\mathcal{D}}\varphi (\mathbf{p},\mathbf{q})\tilde{\rho }(\mathbf{q},\mathbf{p},\xi )d{\omega }_{\mathbf{q}}{\omega }_{\mathbf{p}}d\xi \propto {\int }_{\mathcal{D}}\varphi (\mathbf{p},\mathbf{q})\rho (\mathbf{p},\mathbf{q})d{\omega }_{\mathbf{q}}{\omega }_{\mathbf{p}}$$

With the ergodicity assumption of the coupled system, we can then calculate the integral with respect to the Gibbs distribution by direct averaging along trajectories of the extended system. However, it is likely that the assumption of ergodicity fails and there are some known cases where this deterministic way of computing the integral leads to incorrect results. Still, this class of deterministic methods has been often used in molecular simulation.

Nose-Hoover dynamics is perhaps the simplest one of the class of deterministic methods we summarized above. Let $N_d$ denotes the number of degrees of freedom of the system. The extended system of Nose-Hoover dynamics takes the form

$$\frac{d\mathbf{q}}{dt}=M^{-1}\mathbf{p}$$ $$\frac{d\mathbf{p}}{dt}=-\nabla U(\mathbf{q})-\xi \mathbf{p}$$ $$\mu \frac{d\xi }{dt}={\mathbf{p}}^T{M}^{-1}\mathbf{p}-N_d{k}_{B}T$$

When the temperature of the system, measured as the average kinetic energy per degree of freedom, is higher than the prescribed temperature $k_{B}T$, the equation for $\xi$ will have a positive right-hand side and it will ensure that $\xi$ increases, and $\xi$ will eventually become positive. This will damp/cool the physical variables through equation for $\mathbf{p}$; eventually the kinetic energy will fall below the target value, and the right-hand side of the last equation will become negative. By fluctuating in this way, $\xi$ typically controls the time-averaged kinetic energy of the system such that it is very close to $N_d{k}_{B}T/2$, which is what the kinetic energy should be at the equilibrium. $\mu$ is a parameter that can be arbitrarily chosen to improve the performance of the dynamics.

This Nose-Hoover Dynamics preserves an extended measure with density of the form

$$\rho (\mathbf{q},\mathbf{p})\times \exp (-mu{\xi }^2/(2k_{B}T))$$

which we will not prove here. To monitor the performance of numerical methods, we also observe that the Nose-Hoover system preserves the quantity

$$I(\mathbf{q},\mathbf{p},\xi )=H(\mathbf{q},\mathbf{p})+\frac{\mu {\xi }^2}{2}+N_d{k}_{B}T\log s$$

where $s$ solves the ODE $\frac{ds}{dt}=\xi$. In the spirit of the splitting method we present when introducing numerical integrators for the Langevin equation, we here provide a splitting method for solving for the evolution of the Nose-Hoover dynamics. Suppose we want to integrate from time step $n$ to $n+1$ with time step size $\Delta t$, then the scheme is as follows

$${\mathbf{q}}_{n+1/2}={\mathbf{q}}_n+\frac{\Delta t}{2}{M}^{-1}{\mathbf{p}}_n$$

$${\mathbf{p}}_{n+1/2}=e^{-\Delta t{\xi }_n/2}{\mathbf{p}}_n-\frac{1-e^{-\Delta t{\xi }_n/2}}{{\xi }_n}\nabla U({\mathbf{q}}_{n+1/2})$$

$${\xi }_{n+1}={\xi }_n+\Delta t{\mu }^{-1}({\mathbf{p}}_{n+1/2}^T{M}^{-1}{\mathbf{p}}_{n+1/2}-N_d{k}_{B}T)$$

$${\mathbf{p}}_{n+1}=e^{-\Delta t{\xi }_{n+1}/2}{\mathbf{p}}_{n+1/2}-\frac{1-e^{-\Delta t{\xi }_{n+1}/2}}{{\xi }_{n+1}}\nabla U({\mathbf{q}}_{n+1/2})$$

$${\mathbf{q}}_{n+1}={\mathbf{q}}_{n+1/2}+\frac{\Delta t}{2}{M}^{-1}{\mathbf{p}}_{n+1}$$

6. Numerical Example (1D Harmonic Oscillator)

7. Numerical Example (2D Lennard-Jones Fluid)