I'm making a Python implementation of the Crank-Nicolson method to solve a nonlinear time-dependent Fokker-Planck equation:
$\partial_t \rho = -(F \partial_x \rho +\rho \partial_x F) +\frac{\sigma^2}{2}\partial^2_{x}\rho$
However the "mass" (the probability) is not being conserved when $\partial_xF$ is not $0$.
My scheme is the following: I'm considering a scheme with Considering a grid $x={x_0,.\dots,x_J}$ and a time steps $t={t_0,\dots,t_N}$, using central differences in space and forward difference in time, we get \begin{equation} \begin{aligned} \partial_t \rho &= \frac{\rho^{n+1}_{j}-\rho^{n}_{j}}{\Delta t}\\ F \partial_x \rho &= \beta F^{n+1}_{j}\frac{\rho^{n+1}_{j+1}-\rho^{n+1}_{j-1}}{2\Delta x}+(1-\beta)F^{n}_{j}\frac{\rho^{n}_{j+1}-\rho^{n}_{j-1}}{2\Delta x}\\ \rho \partial_x F &= \beta\rho^{n+1}_{j}\frac{F^{n+1}_{j+1}-F^{n+1}_{j-1}}{2\Delta x}+(1-\beta)\rho^{n}_{j}\frac{F^{n}_{j+1}-F^{n}_{j-1}}{2\Delta x}\\ \partial^2_x \rho &= \beta\frac{\rho^{n+1}_{j+1}-2\rho^{n+1}_{j}+\rho^{n+1}_{j-1}}{2\Delta x^2}+(1-\beta)\frac{\rho^{n}_{j+1}-2\rho^{n}_{j}+\rho^{n}_{j-1}}{2\Delta x^2} \end{aligned} \end{equation} The C-N scheme usually has $\beta=0.5$. In our case, $\beta=0$ recovers the forward scheme, and $\beta=1$ the implicit backwards scheme.
Defining $D=\frac{\sigma^2}{2}$, $\xi=\frac{D}{\Delta x}$, $\delta_c=\frac{\Delta t}{2\Delta x}$, and putting all together: \begin{equation} \begin{aligned} \rho^{n+1}_{j}-\rho^{n+1}_{j} &= - [\beta \delta_c F^{n+1}_{j} (\rho^{n+1}_{j+1}-\rho^{n+1}_{j-1})+(1-\beta)\delta_c F^{n}_{j} (\rho^{n}_{j+1}-\rho^{n}_{j-1})]\\ &-[\beta\delta_c\rho^{n+1}_{j}(F^{n+1}_{j+1}-F^{n+1}_{j-1})+(1-\beta)\delta_c\rho^{n}_{j}(F^{n}_{j+1}-F^{n}_{j-1})] \\ &+\xi\delta_c\beta(\rho^{n+1}_{j+1}-2\rho^{n+1}_{j}+\rho^{n+1}_{j-1})+\xi\delta_c(1-\beta)(\rho^{n}_{j+1}-2\rho^{n}_{j}+\rho^{n}_{j-1}) \end{aligned} \end{equation}
Separating $\rho^n$ and $\rho^{n+1}$ \begin{equation} \begin{aligned} \rho^{n+1}_{j} &+ \beta \delta_c F^{n+1}_{j} (\rho^{n+1}_{j+1}-F^{n+1}_{j-1})+\beta\delta_c\rho^{n+1}_{j}(F^{n+1}_{j+1}-F^{n+1}_{j-1})-\xi\delta_c\beta(\rho^{n+1}_{j+1}2\rho^{n+1}_{j}+\rho^{n+1}_{j-1}) =\\ &\rho^{n}_{j}-(1-\beta)\delta_c F^{n}_{j}(\rho^{n}_{j+1}-\rho^{n}_{j-1}) -(1-\beta)\delta_c\rho^{n}_{j}(F^{n}_{j+1}-F^{n}_{j-1}) +\xi\delta_c(1-\beta)(\rho^{n}_{j+1}-2\rho^{n}_{j}+\rho^{n+1}_{j-1}) \end{aligned} \end{equation}
Then the scheme can be written as \begin{equation} A \rho^{n+1}_{j}= B\rho^{n}_{j} \end{equation}
where \begin{equation} \begin{aligned} A &= I + \delta_c\beta A_{\mathrm{advection}}- \xi\delta_c\beta A_{\mathrm{difussion}}\\ B &=I - \delta_c(1-\beta) A_{\mathrm{advection}}+ \xi\delta_c(1-\beta) A_{\mathrm{difussion}}\\ \end{aligned} \end{equation}
with \begin{equation} A_{\mathrm{advection}}=\begin{bmatrix} F^{n+1}_{2}-F^{n+1}_{0} & F^{n+1}_{1} & 0 & \dots & 0 \\ -F^{n+1}_{1} & F^{n+1}_{3}-F^{n+1}_{1} & F^{n+1}_{2} & \dots & 0 \\ 0 & -F^{n+1}_{2} & F^{n+1}_{4}-F^{n+1}_{2} & \dots & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & \dots & F^{n+1}_{J+1}-F^{n+1}_{J-1} \end{bmatrix} \end{equation}
where $F^{n+1}_{0}$ and $F^{n+1}_{J+1}$ have to be chosen as a boundary condition, since it's outside the array. $B_{\mathrm{advection}}$ is the same as $A_{\mathrm{advection}}$ at time step $n$.
The diffusion matrices are: \begin{equation} A_{\mathrm{difussion}}=B_{\mathrm{difussion}}=\begin{bmatrix} -2 & 1 & 0 & \dots & 0 \\ 1 & -2 & 1 & \dots & 0 \\ 0 & 1 &-2 & \dots & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & \dots & -2 \end{bmatrix} \end{equation}
Therefore
\begin{equation} A=\begin{bmatrix} 1+\delta_c\beta(F^{n+1}_{2}-F^{n+1}_{0}+2\xi) & \delta_c\beta(F^{n+1}_{1}-\xi) & 0 & \dots \\ \delta_c\beta(-F^{n+1}_{1}-\xi) & 1+\delta_c\beta(F^{n+1}_{3}-F^{n+1}_{1}+2\xi) & \delta_c\beta(F^{n+1}_{2}-\xi) & \dots \\ 0 & \delta_c\beta(-F^{n+1}_{2}-\xi) & 1+\delta_c\beta(F^{n+1}_{4}-F^{n+1}_{2}+2\xi) & \dots \\ \vdots & \vdots & \vdots & \ddots \\ 0 & 0 & 0 & \dots \end{bmatrix} \end{equation}
\begin{equation} B=\begin{bmatrix} 1-\delta_c(1-\beta)(F^{n}_{2}-F^{n}_{0}+2\xi) & -\delta_c(1-\beta)(F^{n}_{1}-\xi) & 0 & \dots \\ -\delta_c(1-\beta)(-F^{n}_{1}-\xi) & 1-\delta_c(1-\beta)(F^{n}_{3}-F^{n}_{1}+2\xi) & -\delta_c(1-\beta)(F^{n}_{2}-\xi) & \dots \\ 0 & -\delta_c(1-\beta)(-F^{n}_{2}-\xi) & 1-\delta_c(1-\beta)(F^{n}_{4}-F^{n}_{2}+2\xi) & \dots \\ \vdots & \vdots & \vdots & \ddots \\ 0 & 0 & 0 & \dots \end{bmatrix} \end{equation}
The integration scheme needs to solve $\rho^{n+1}_{j}= A^{-1}B\rho^{n}_{j}$.
I'm using some Neumann boundary conditions such that the derivative at x[0],x[N] is the same as that at x[1],x[N-1].
I'm using the following code:
from __future__ import division
from scipy.sparse import spdiags
from scipy.sparse.linalg import spsolve
import numpy as np
import pylab
import matplotlib.pyplot as plt
plt.close('all')
def F_OB(x,t,v=-2,x0=0): "Ornstein uhlembeck field"
FM=np.zeros([len(x),len(t)])
FM=[v*(x-x0) for i in range(len(t))]
return FM
def F_drift(x,t,v=-2): "Constant drift field"
FM=np.zeros([len(x),len(t)])
FM=[np.array(v*np.ones(len(x)),dtype='float64') for i in range(len(t))]
return FM
def F_cos(x,t,v=0.01): "periodic forcing drift field"
FM=np.zeros([len(x),len(t)])
FM=[v*np.cos(t[i])*np.array(v*np.ones(len(x)),dtype='float64') for i in range(len(t))]
return FM
def makeCN_advection_difussion_matrices(x, dc,xi,FM,FM1,beta=0.5):
temp_F=np.zeros(len(x)) #gotta implement time evolution.
temp_F=[1*(FM[i+1]-FM[i-1]) for i in range(1,len(x)-1)]
Fn_diag=np.copy(FM)
Fn_diag[0]=2*(FM[1]-FM[0])
Fn_diag[-1]=2*(FM[-1]-FM[-2])
Fn_diag[1:-1]=temp_F
temp_F=[1*(FM1[i+1]-FM1[i-1]) for i in range(1,len(x)-1)]
Fn1_diag=np.copy(FM)
Fn1_diag[1:-1]=temp_F
Fn1_diag[0]=2*(FM1[1]-FM1[0])
Fn1_diag[-1]=2*(FM1[-1]-FM1[-2])
ones = np.ones(len(x))
Aid = spdiags( [ones], (0), len(x), len(x) )
Aadv = spdiags( [-FM1, Fn1_diag, FM1], (-1,0,1), len(x), len(x) )
Adif = spdiags( [ones, -2*ones, ones], (-1,0,1), len(x), len(x) )
A = Aid.tocsr() + beta*dc*Aadv.tocsr() - beta*xi*dc*Adif.tocsr()
Mid = spdiags( [ones], (0), len(x), len(x) )
Madv = spdiags( [-FM, Fn_diag, FM], (-1,0,1), len(x), len(x) )
Mdif = spdiags( [ones, -2*ones, ones], (-1,0,1), len(x), len(x) )
M = Mid.tocsr() - (1-beta)*dc*Madv.tocsr() + (1-beta)*xi*dc*Mdif.tocsr()
return A, M
s=0.1 # Set up basic constants
J = 3000 # total number of mesh points
xmin=-13
dt = 0.1 # time step
x = np.linspace(xmin,-xmin,J) # vertices
dx = np.abs(x[1]-x[0]) # space step
tf=100+dt#300
t=np.arange(0,tf,dt)
tp=t[:-2]
FM=F_OB(x,t,v=-0.01,x0=0) # FM=F_drift(x,t,v=0.1) # FM=F_cos(x,t,v=1)
dc=0.5*dt/dx
xi=0.5*s**2/dx
def Rinit2(x,mean, cov): #### Initial conditions (peak function)
from scipy.stats import multivariate_normal
# pos = np.dstack((x))
return multivariate_normal.pdf(x, mean, cov)
s0=s/3
pdf = Rinit2(x,mean=-0.5,cov=s0)
p0 = pdf
print('integral check: ', np.trapz(p0,x))
u_init = p0
A, M = makeCN_advection_difussion_matrices(x, dc,xi,FM[0],FM[1],beta=0.5) #Integration scheme
u = u_init
u[0]=u[1]+u[2]-u[3];u[-1]=u[-2]-u[-4]+u[-3]#Newmann boundary, the derivative is te same as the previuos or next point
ut=np.zeros([np.shape(u)[0],len(tp)])
ut[:,0]=u
for i in range(len(tp)):
u = spsolve(A, M * u)
u[0]=u[1]+u[2]-u[3];u[-1]=u[-2]-u[-4]+u[-3]#Newmann boundary, the derivative is te same as the previuos or next point
ut[:,i]=u
A, M = makeCN_advection_difussion_matrices(x, dc,xi,FM[i],FM[i+1],beta=0.5)
fig, ax = plt.subplots()
ax.pcolor(tp,x,ut)
ax.set_ylim([x[0],x[-1]])
mass=[np.trapz(ut[:,i],x) for i in range(len(tp))]
fig, ax = plt.subplots()
ax.plot(tp,mass)
from matplotlib.animation import FuncAnimation
fig, ax = plt.subplots()
line, = ax.plot(x, ut[:,0],color='C0')
def update(i):
line.set_ydata(ut[:,i])
return (line)
anim = FuncAnimation(fig, update, frames=range(len(tp)), interval=40)
It works fine for pure diffusion, pure advection, and constant advection with diffusion. But in the Ornstein-Uhlenbeck case the probability grows exponentially.