Notebook

PDE solver in julia 0.6.0 by wanglj

Solving the standard diffusion equation.¶

$\begin{equation} \frac{\partial \rho(x, t)}{\partial t} = D \frac{\partial^2 \rho(x, t)}{\partial x^2} \end{equation}$

Finite difference method for time and space¶

Finite difference of time:

$\begin{equation} \frac{\partial\rho}{\partial t} (x_i,t_j) = \frac{\rho(x_i, t_j + \Delta t) - \rho(x_i, t_j)}{\Delta t} + \mathcal{O}(\Delta t) \end{equation}$

Finite difference of space:

$\begin{equation} \frac{\partial^2\rho}{\partial x^2} (x_i,t_j) = \frac{\rho(x_i+\Delta x,t_j) - 2\rho(x_i,t_j) + \rho(x_i-\Delta x,t_j)}{dx^2} + \mathcal{O}(\Delta x^2) \end{equation}$

Explicit scheme (forward time)¶

Apply equation(2) and (3) to equation(1), and let $\rho_{i,j} = \rho(x_i,t_j)$ , we get the expression of forward time evolution

$\begin{align} &\ \frac{\rho_{i,j+1} - \rho_{i,j}}{\Delta t} = D \frac{\rho_{i+1,j} - 2\rho_{i,j} + \rho_{i-1,j}}{\Delta x^2} \\ &\Rightarrow \rho_{i,j+1} = (1 - 2\lambda) \rho_{i,j} + \lambda (\rho_{i+1,j} + \rho_{i-1, j}) \end{align}$

where

$\lambda=\frac{D \cdot \Delta t}{\Delta x^2}$ . Rewrite this in matrix form,

$\begin{equation} \vec{\rho}_{i, j+1} = A \vec{\rho}_{i, j} + R(\text{accounts for boundary conditions}) \end{equation}$

Supplemented by initial condition: $\rho_{i,0} = \rho(x_i)$ , and Boundary conditions: $\rho_{0,j} = \rho_0$ , $\rho_{m,j} = \rho_m$ , we can solve it like this:

$\begin{equation} \begin{bmatrix} \rho_{1,j+1} \\ \rho_{2,j+1} \\ \vdots \\ \rho_{m-1,j+1} \end{bmatrix} = \begin{bmatrix} 1 - 2 \lambda & \lambda & 0 & \\ \lambda & 1 - 2 \lambda & \lambda & \\ & \ddots & \ddots & \ddots \\ & 0 & \lambda & 1 - 2 \lambda \end{bmatrix} \begin{bmatrix} \rho_{1,j} \\ \rho_{2,j} \\ \vdots \\ \rho_{m-1,j} \end{bmatrix} + \lambda \begin{bmatrix} \rho_0 \\ 0 \\ \vdots \\ \rho_m \end{bmatrix} \end{equation}$

Sovling an example system with initial value of $\rho_{center}=1.0$ at center and boundary conditions with $\rho_L = \rho_R = 0$ .

In [2]:

# Define constants and calculate standard dt.
const nx = 100
const dx = 1.e-11
const D = 1.69e-9
const steps = 100
dt = 0.5 * dx^2 /(2 * D)

# Set grid and initial value(including boundary conditions) for dx. 
grid = zeros(Float64, nx, steps)
grid[50, 1] = 1.0
Lbound = 0.0
Rbound = 0.0
grid[1, :] = Lbound
grid[nx, :] = Rbound

# Set up tridiagonal and sparse matrix for A and R.
lbd = D * dt / dx^2 
matAsp = SymTridiagonal(fill(1 - 2 * lbd, nx - 2), fill(lbd, nx - 3))
matRsp = spzeros(nx - 2)
matRsp[1] = Lbound
matRsp[nx - 2] = Rbound

# Calculate and plot. 
for i = 2:steps
    grid[2:nx-1, i] = matAsp * grid[2:nx-1, i - 1] + lbd * matRsp
end

using Plots
#plotly(size=(600,400))
#plot(grid[:,1], title="t")
gr()
anim = @animate for i=1:steps
    plot(grid[:, i])
    plot!(ylims = (0, 1), yticks = 0:0.1:1, title="t = $(round(i * dt * 1.E12, 3)) ps")
    end every 10
gif(anim, "pics/forward.gif", fps = 5)

WARNING: using Plots.grid in module Main conflicts with an existing identifier.
INFO: Saved animation to /home/wanglj/PDE/pics/forward.gif

Out[2]:

Implicit scheme (backward time)¶

In implicit scheme, time evolves backwards.

$\begin{align} &\ \frac{\rho_{i,j} - \rho_{i,j-1}}{\Delta t} = D \frac{\rho_{i+1,j} - 2\rho_{i,j} + \rho_{i-1,j}}{\Delta x^2} \\ &\Rightarrow (1 + 2\lambda) \rho_{i,j} - \lambda (\rho_{i+1,j} + \rho_{i-1, j}) = \rho_{i,j-1} \end{align}$

where

$\lambda=\frac{D \cdot \Delta t}{\Delta x^2}$ . Rewrite this in matrix form,

$\begin{equation} A \vec{\rho}_{i, j} - R = \vec{\rho}_{i, j - 1} \end{equation}$

Supplemented by initial condition: $\rho_{i,0} = \rho(x_i)$ , and Boundary conditions: $\rho_{0,j} = \rho_0$ , $\rho_{m,j} = \rho_m$ , we can solve it like this:

$\begin{equation} \begin{bmatrix} 1 + 2 \lambda & - \lambda & 0 & \\ - \lambda & 1 + 2 \lambda & - \lambda & \\ & \ddots & \ddots & \ddots \\ & 0 & - \lambda & 1 + 2 \lambda \end{bmatrix} \begin{bmatrix} \rho_{1,j} \\ \rho_{2,j} \\ \vdots \\ \rho_{m-1,j} \end{bmatrix} - \lambda \begin{bmatrix} \rho_0 \\ 0 \\ \vdots \\ \rho_m \end{bmatrix} = \begin{bmatrix} \rho_{1,j-1} \\ \rho_{2,j-1} \\ \vdots \\ \rho_{m-1,j-1} \end{bmatrix} \end{equation}$

Sovling an example system with initial value of $\rho_{center}=1.0$ at center and boundary conditions with $\rho_L = \rho_R = 0$ .

In [3]:

const nx = 100
const dx = 1.e-11
const D = 1.69e-9
const steps = 100
dt = 0.5 * dx^2 /(2 * D)

# Set grid and initial value(including boundary conditions) for dx. 
grid = zeros(Float64, nx, steps)
grid[50, 1] = 1.0
Lbound = 0.0
Rbound = 0.0
grid[1, :] = Lbound
grid[nx, :] = Rbound

# Set up tridiagonal and sparse matrix for A and R.
lbd = D * dt / dx^2
matRsp = spzeros(nx - 2)
matRsp[1] = Lbound
matRsp[nx - 2] = Rbound
# Calculate and plot. 
ds = zeros(nx - 3)
d = zeros(nx - 2)
for i = 2:steps
    bmatrix = grid[2:nx-1, i-1] + lbd * matRsp
    ds .= -lbd
    d .= 1 + 2 * lbd
    grid[2:nx-1, i] = LAPACK.ptsv!(d, ds, bmatrix)
end

anim = @animate for i=1:steps
    plot(grid[:, i])
    plot!(ylims = (0, 1), yticks = 0:0.1:1, title="t = $(round(i * dt * 1.E12, 3)) ps")
    end every 10
gif(anim, "pics/backward.gif", fps = 5)

INFO: Saved animation to /home/wanglj/PDE/pics/backward.gif

Out[3]:

Crank-Nicolson scheme¶

In Crank-Nicolson scheme, we take the average of explicit scheme and implicit scheme.

where

$\lambda=\frac{D \cdot \Delta t}{\Delta x^2}$ . Rewrite this in matrix form,

$\begin{equation} A \vec{\rho}_{i, j + 1} - R = B \vec{\rho}_{i, j} + R \end{equation}$

Supplemented by initial condition: $\rho_{i,0} = \rho(x_i)$ , and Boundary conditions: $\rho_{0,j} = \rho_0$ , $\rho_{m,j} = \rho_m$ , we can solve it like this:

$\begin{equation} \begin{bmatrix} 2 + 2 \lambda & - \lambda & 0 & \\ - \lambda & 2 + 2 \lambda & - \lambda & \\ & \ddots & \ddots & \ddots \\ & 0 & - \lambda & 2 + 2 \lambda \end{bmatrix} \begin{bmatrix} \rho_{1,j} \\ \rho_{2,j} \\ \vdots \\ \rho_{m-1,j} \end{bmatrix} - \lambda \begin{bmatrix} \rho_0 \\ 0 \\ \vdots \\ \rho_m \end{bmatrix} = \begin{bmatrix} 2 - 2 \lambda & \lambda & 0 & \\ \lambda & 2 - 2 \lambda & \lambda & \\ & \ddots & \ddots & \ddots \\ & 0 & \lambda & 2 - 2 \lambda \end{bmatrix} \begin{bmatrix} \rho_{1,j} \\ \rho_{2,j} \\ \vdots \\ \rho_{m-1,j} \end{bmatrix} + \lambda \begin{bmatrix} \rho_0 \\ 0 \\ \vdots \\ \rho_m \end{bmatrix} \end{equation}$

Sovling an example system with initial value of $\rho_{center}=1.0$ at center and boundary conditions with $\rho_L = \rho_R = 0$ .

In [4]:

const nx = 100
const dx = 1.e-11
const D = 1.69e-9
const steps = 100
dt = 0.5 * dx^2 /(2 * D)

# Set grid and initial value(including boundary conditions) for dx. 
grid = zeros(Float64, nx, steps)
grid[50, 1] = 1.0
Lbound = 0.0
Rbound = 0.0
grid[1, :] = Lbound
grid[nx, :] = Rbound

# Set up tridiagonal and sparse matrix for A and R.
lbd = D * dt / dx^2
matBsp = SymTridiagonal(fill(2 - 2 * lbd, nx - 2), fill(lbd, nx - 3))
matRsp = spzeros(nx - 2)
matRsp[1] = Lbound
matRsp[nx - 2] = Rbound
# Calculate and plot. 
ds = zeros(nx - 3)
d = zeros(nx - 2)
for i = 2:steps
    bmatrix = matBsp * grid[2:nx-1, i-1] + 2 * lbd * matRsp
    ds .= -lbd
    d .= 2 + 2 * lbd
    grid[2:nx-1, i] = LAPACK.ptsv!(d, ds, bmatrix)
end

anim = @animate for i=1:steps
    plot(grid[:, i])
    plot!(ylims = (0, 1), yticks = 0:0.1:1, title="t = $(round(i * dt * 1.E12, 3)) ps")
    end every 10
gif(anim, "pics/crank-nicolson.gif", fps = 5)

INFO: Saved animation to /home/wanglj/PDE/pics/crank-nicolson.gif

Out[4]:

Solving diffusion equation with a potential profile.¶

$\begin{equation} \frac{\partial \rho(x, t)}{\partial t} = D \frac{\partial}{\partial x}[ \frac{\partial \rho(x, t)}{\partial x} + \frac{1}{k_BT} \frac{\partial U(x)}{\partial x}\rho(x,t)] = D \frac{\partial^2 \rho(x, t)}{\partial x^2} + \frac{D}{k_BT} \cdot U^{\prime}(x) \cdot \frac{\partial \rho(x, t)}{\partial x} + \frac{D}{k_BT} U^{\prime \prime}(x) \rho(x, t) \end{equation}$

The finite difference in Crank-Nicolson scheme for the above equation is:

$\begin{align} &\ \frac{\rho_{i,j + 1} - \rho_{i,j}}{\Delta t} = \frac{D}{2} ( \frac{\rho_{i+1,j} - 2\rho_{i,j} + \rho_{i-1,j}}{\Delta x^2} + \frac{\rho_{i + 1,j + 1} - 2\rho_{i,j + 1} + \rho_{i - 1,j + 1}}{\Delta x^2} ) + \frac{D \cdot U^{\prime}(x_i)}{2 k_BT} (\frac{\rho_{i + 1,j + 1} - \rho_{i - 1,j + 1}}{2 \Delta x} + \frac{\rho_{i + 1,j} - \rho_{i - 1, j}}{2 \Delta x}) + \frac{D \cdot U^{\prime \prime}(x_i)}{2 k_BT} (\rho_{i,j + 1} + \rho_{i,j}) \\ &\Rightarrow (2 + 2\lambda + \tau_i) \rho_{i,j + 1} + (\eta_i - \lambda)\rho_{i+1,j + 1} - (\lambda + \eta_i) \rho_{i-1, j + 1} = (2 - 2\lambda - \tau_i) \rho_{i,j} + (\lambda - \eta_i) \rho_{i+1,j} + (\lambda + \eta_i) \rho_{i-1, j} \end{align}$

where $\lambda=\frac{D \cdot \Delta t}{\Delta x^2}$ , $\eta_i = - \frac{D \cdot \Delta t}{2 k_BT \Delta x} \cdot U^{\prime}(x_i)$ , $\tau_i = - \frac{D \cdot \Delta t}{k_BT} \cdot U^{\prime \prime}(x_i)$ .

Rewrite this in matrix form,

$\begin{equation} A \vec{\rho}_{i, j + 1} - R = B \vec{\rho}_{i, j} + R \end{equation}$

Supplemented by initial condition: $\rho_{i,0} = \rho(x_i)$ , and Boundary conditions: $\rho_{0,j} = \rho_0$ , $\rho_{m,j} = \rho_m$ , we can solve it like this:

$\begin{equation} \begin{bmatrix} 2 + 2\lambda + \tau_1 & \eta_1 - \lambda & 0 & \\ - (\lambda + \eta_2) & 2 + 2\lambda + \tau_2 & \eta_2 - \lambda & \\ & \ddots & \ddots & \ddots \\ & 0 & - (\lambda + \eta_{m-1}) & 2 + 2\lambda + \tau_{m-1} \end{bmatrix} \begin{bmatrix} \rho_{1,j} \\ \rho_{2,j} \\ \vdots \\ \rho_{m-1,j} \end{bmatrix} + \begin{bmatrix} - (\lambda + \eta_1) \rho_0 \\ 0 \\ \vdots \\ (\eta_{m - 1} - \lambda) \rho_m \end{bmatrix} = \begin{bmatrix} 2 - 2\lambda - \tau_1 & \lambda - \eta_1 & 0 & \\ \lambda + \eta_2 & 2 - 2\lambda - \tau_2 & \lambda - \eta_2 & \\ & \ddots & \ddots & \ddots \\ & 0 & \lambda + \eta_{m - 1} & 2 - 2\lambda - \tau_{m-1} \end{bmatrix} \begin{bmatrix} \rho_{1,j} \\ \rho_{2,j} \\ \vdots \\ \rho_{m-1,j} \end{bmatrix} + \begin{bmatrix} (\lambda + \eta_1) \rho_0 \\ 0 \\ \vdots \\ (\lambda - \eta_{m - 1}) \rho_m \end{bmatrix} \end{equation}$

Sovling an example system with initial value of $\rho=1.0$ at $x=30 \unicode{x212B}$ and boundary conditions with $\rho_L = \rho_R = 0$ using the potential profile of $n_w = 0$ .

In [5]:

using DataFrames
df = readtable("csv/new0_1d.csv")

# Fit potential curve and its derivative from csv data file using Scipy.
using PyCall
@pyimport importlib.machinery as machinery
loader = machinery.SourceFileLoader("CsvFit","/home/wanglj/PDE/potential.py")
cf = loader[:load_module]("CsvFit")
fit = cf[:Potential]("csv/new0_1d.csv")
fit2 = cf[:Potential]("csv/new1_extended2nd_1d.csv")
display(plot([df[:z].*1.e9, df[:z].*1.e9], [fit[:f](df[:z]), fit2[:f](df[:z])], title="potential", label=["new 0","new 1"]))
display(plot([df[:z].*1.e9, df[:z].*1.e9], [fit[:grad](df[:z]).*1.e-9, fit2[:grad](df[:z]).*1.e-9], title="first derivative"))
display(plot([df[:z].*1.e9, df[:z].*1.e9], [fit[:grad2](df[:z]).*1.e-18, fit2[:grad2](df[:z]).*1.e-18], title="second derivative"))

WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
in #readtable#84 at /home/wanglj/.julia/v0.6/DataFrames/src/dataframe/io.jl
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
in #readtable#84 at /home/wanglj/.julia/v0.6/DataFrames/src/dataframe/io.jl
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
WARNING: Compat.UTF8String is deprecated, use String instead.
  likely near In[5]:2
in builddf at /home/wanglj/.julia/v0.6/DataFrames/src/dataframe/io.jl

In [6]:

const nx = 100
const dx = 1.e-10
const D = 1.69e-9
const steps = 100
const kBT = 0.593 # kcal/mol
dt = 0.5 * dx^2 /(2 * D)

# Set grid and initial value(including boundary conditions) for dx. 
grid = zeros(Float64, nx, steps)

# Set η and τ array
η = [ - D * dt * fit[:grad](dx * i + 12.5e-10)[1] / (2 * kBT * dx) for i in 1:nx - 2]
τ = [ - D * dt * fit[:grad2](dx * i + 12.5e-10)[1] / kBT for i in 1:nx - 2]

grid[30, 1] = 1.0
Lbound = 0.0
Rbound = 0.0
grid[1, :] = Lbound
grid[nx, :] = Rbound

# Set up tridiagonal and sparse matrix for B and R.
λ = D * dt / dx^2
matBsp = Tridiagonal(fill(λ, nx - 3), fill(2 - 2 * λ, nx - 2), fill(λ, nx - 3))
matBsp[1, 1] -= τ[1]
matBsp[1, 2] -= η[1]
for i = 2:nx - 3
    matBsp[i, i] -= τ[i]
    matBsp[i, i-1] += η[i]
    matBsp[i, i+1] -= η[i]
end
matBsp[nx-2, nx-2] -= τ[end]
matBsp[nx-2, nx-3] += η[end]

matRsp = spzeros(nx - 2)
matRsp[1] = Lbound * 2 * (λ + η[1])
matRsp[nx - 2] = Rbound * 2 * (λ - η[end])
# Calculate and plot. 
dg = fill(2 + 2 * λ, nx - 2) # Diagonal elements
doff = fill(- λ, nx - 3) # Offdiagonal elements
dl = doff - η[2:end]
ds = dg + τ
du = doff + η[1:end - 1]
# Solve tridiagonal system.
for i = 2:steps
    bmatrix = matBsp * grid[2:nx-1, i-1] + matRsp
#     dl = doff - η[2:end]
#     ds = dg + τ
#     du = doff + η[1:end - 1]
    grid[2:nx-1, i] = LAPACK.gtsv!(copy(dl), copy(ds), copy(du), bmatrix)
end

anim = @animate for i=1:steps
    plot(grid[:, i])
    plot!(ylims = (0, 1), yticks = 0:0.1:1, title="t = $(round(i * dt * 1.E12, 3)) ps")
    end every 10
gif(anim, "pics/potential-crank-nicolson.gif", fps = 5)

WARNING: redefining constant dx
INFO: Saved animation to /home/wanglj/PDE/pics/potential-crank-nicolson.gif

Out[6]:

Solving coupled diffusion equations with potential profiles and reactions.¶

When reaction terms are added to the equation(16), equations corresponding to different $n_w$ become coupled equations:

$\begin{equation} \begin{cases} \frac{\partial \rho_0(x, t)}{\partial t} = D \frac{\partial}{\partial x}[ \frac{\partial \rho_0(x, t)}{\partial x} + \frac{1}{k_BT} \frac{\partial U_0(x)}{\partial x}\rho_0(x, t)] - [k_{0 \to 1} + k_0^{wf}(x)]\rho_0(x, t) + k_{1 \to 0}\rho_1(x, t) \\ \frac{\partial \rho_1(x, t)}{\partial t} = D \frac{\partial}{\partial x}[ \frac{\partial \rho_1(x, t)}{\partial x} + \frac{1}{k_BT} \frac{\partial U_1(x)}{\partial x}\rho_1(x, t)] - [k_{1 \to 2} + k_{1 \to 0} + k_1^{wf}(x)]\rho_1(x, t) + k_{2 \to 1}\rho_2(x, t) + k_{0 \to 1}\rho_0(x, t) \\ \frac{\partial \rho_2(x, t)}{\partial t} = D \frac{\partial}{\partial x}[ \frac{\partial \rho_2(x, t)}{\partial x} + \frac{1}{k_BT} \frac{\partial U_2(x)}{\partial x}\rho_2(x, t)] - [k_{2 \to 3} + k_{2 \to 1} + k_2^{wf}(x)]\rho_2(x, t) + k_{3 \to 2}\rho_3(x, t) + k_{1 \to 2}\rho_1(x, t) \\ \vdots \\ \vdots \\ \frac{\partial \rho_{n-1}(x, t)}{\partial t} = D \frac{\partial}{\partial x}[ \frac{\partial \rho_{n-1}(x, t)}{\partial x} + \frac{1}{k_BT} \frac{\partial U_{n-1}(x)}{\partial x}\rho_{n-1}(x, t)] - [k_{n-1 \to n} + k_{n-1 \to n-2} + k_{n-1}^{wf}(x)]\rho_{n-1}(x, t) + k_{n \to n-1}\rho_n(x, t) + k_{n-2 \to n-1}\rho_{n-2}(x, t) \\ \frac{\partial \rho_n(x, t)}{\partial t} = D \frac{\partial}{\partial x}[ \frac{\partial \rho_n(x, t)}{\partial x} + \frac{1}{k_BT} \frac{\partial U_n(x)}{\partial x}\rho_n(x, t)] - [k_{n \to n-1} + k_n^{wf}(x)]\rho_n(x, t) + k_{n-1 \to n}\rho_{n-1}(x, t) \end{cases} \end{equation}$

Let $\rho_{i,j}^k = \rho_k(x_i,t_j)$ , the finite difference in Crank-Nicolson scheme for $n_w = k$ is:

$\begin{align} &\ \frac{\rho_{i,j + 1}^k - \rho_{i,j}^k}{\Delta t} = \frac{D}{2} ( \frac{\rho_{i+1,j}^k - 2\rho_{i,j}^k + \rho_{i-1,j}}{\Delta x^2} + \frac{\rho_{i + 1,j + 1} - 2\rho_{i,j + 1} + \rho_{i - 1,j + 1}}{\Delta x^2} ) + \frac{D \cdot U_k^{\prime}(x_i)}{2 k_BT} (\frac{\rho_{i + 1,j + 1}^k - \rho_{i - 1,j + 1}^k}{2 \Delta x} + \frac{\rho_{i + 1,j}^k - \rho_{i - 1, j}^k}{2 \Delta x}) + \frac{D \cdot U_k^{\prime \prime}(x_i)}{2 k_BT} (\rho_{i,j + 1}^k + \rho_{i,j}^k) - \frac{k_{k \to k-1} + k_{k \to k+1} + k_k^{wf}(x_i)}{2} (\rho_{i,j+1}^k + \rho_{i,j}^k) + \frac{k_{k-1 \to k}}{2} (\rho_{i,j+1}^{k-1} + \rho_{i,j}^{k-1}) + \frac{k_{k+1 \to k}}{2} (\rho_{i,j+1}^{k+1} + \rho_{i,j}^{k+1})\\ &\Rightarrow (2 + 2\lambda + \tau_i^k) \rho_{i,j + 1}^k + (\eta_i - \lambda)\rho_{i+1,j + 1}^k - (\lambda + \eta_i^k) \rho_{i-1, j + 1}^k + \xi_i^k \rho_{i,j + 1}^k - \zeta^k \rho_{i,j + 1}^{k+1} - \kappa^k \rho_{i,j + 1}^{k-1} = (2 - 2\lambda - \tau_i^k) \rho_{i,j}^k + (\lambda - \eta_i^k) \rho_{i+1,j}^k + (\lambda + \eta_i^k) \rho_{i-1, j}^k - \xi_i^k \rho_{i,j + 1}^k + \zeta^k \rho_{i,j + 1}^{k+1} + \kappa^k \rho_{i,j + 1}^{k-1} \end{align}$

where

$\begin{align} &\lambda=\frac{D \cdot \Delta t}{\Delta x^2} \\ &\eta_i^k = - \frac{D \cdot \Delta t}{2 k_BT \Delta x} \cdot U_k^{\prime}(x_i) \\ &\tau_i^k = - \frac{D \cdot \Delta t}{k_BT} \cdot U_k^{\prime \prime}(x_i) \\ &\xi_i^k = (k_{k \to k+1} + k_{k \to k-1} + k_k^{wf}(x_i)) \Delta t \\ &\zeta^k = k_{k+1 \to k} \Delta t \\ &\kappa^k = k_{k-1 \to k} \Delta t \end{align}$

We further define following $(m-1) \times (m-1)$ matrices:

$\begin{align} &\boldsymbol{A}_k = \begin{bmatrix} 2 + 2\lambda + \tau_1^k & \eta_1^k - \lambda & 0 & \\ - (\lambda + \eta_2^k) & 2 + 2\lambda + \tau_2^k & \eta_2^k - \lambda & \\ & \ddots & \ddots & \ddots \\ & 0 & - (\lambda + \eta_{m-1}^k) & 2 + 2\lambda + \tau_{m-1}^k \end{bmatrix} \\ &\boldsymbol{B}_k = \begin{bmatrix} 2 - 2\lambda - \tau_1^k & \lambda - \eta_1^k & 0 & \\ \lambda + \eta_2^k & 2 - 2\lambda - \tau_2^k & \lambda - \eta_2^k & \\ & \ddots & \ddots & \ddots \\ & 0 & \lambda + \eta_{m - 1}^k & 2 - 2\lambda - \tau_{m-1}^k \end{bmatrix} \\ &\boldsymbol{\Xi}_k = \begin{bmatrix} \xi_1^k & & & \\ & \xi_2^k & & \\ & & \ddots & \\ & & & \xi_{m-1}^k \end{bmatrix} \\ &\boldsymbol{Z}_k = \begin{bmatrix} \zeta^k & & & \\ & \zeta^k & & \\ & & \ddots & \\ & & & \zeta^k \end{bmatrix} \\ &\boldsymbol{K}_k= \begin{bmatrix} \kappa^k & & & \\ & \kappa^k & & \\ & & \ddots & \\ & & & \kappa^k \end{bmatrix} \\ &\boldsymbol{P}_{k, j} = \begin{bmatrix} \rho_{1,j}^k \\ \rho_{2,j}^k \\ \vdots \\ \rho_{m-1,j}^k \end{bmatrix} \\ &\boldsymbol{R}_{k} = \begin{bmatrix} (\lambda + \eta_1^k) \rho_0^k \\ 0 \\ \vdots \\ (\lambda - \eta_{m - 1}^k) \rho_m^k \end{bmatrix} \end{align}$

Rewrite in matrix form,

$\begin{equation} \begin{bmatrix} \boldsymbol{A}_0 + \boldsymbol{\Xi}_0 & -\boldsymbol{Z}_0 & 0 & \\ - \boldsymbol{K}_1 & \boldsymbol{A}_1 + \boldsymbol{\Xi}_1 & -\boldsymbol{Z}_1 & \\ & \ddots & \ddots & \ddots \\ & 0 & - \boldsymbol{K}_n & \boldsymbol{A}_n + \boldsymbol{\Xi}_n \end{bmatrix} \begin{bmatrix} \boldsymbol{P}_{0, j} \\ \boldsymbol{P}_{1, j} \\ \vdots \\ \boldsymbol{P}_{n, j} \end{bmatrix} - \begin{bmatrix} \boldsymbol{R}_{0} \\ \boldsymbol{R}_{1} \\ \vdots \\ \boldsymbol{R}_{n} \end{bmatrix} = \begin{bmatrix} \boldsymbol{B}_0 - \boldsymbol{\Xi}_0 & \boldsymbol{Z}_0 & 0 & \\ \boldsymbol{K}_1 & \boldsymbol{B}_1 - \boldsymbol{\Xi}_1 & \boldsymbol{Z}_1 & \\ & \ddots & \ddots & \ddots \\ & 0 & \boldsymbol{K}_n & \boldsymbol{B}_n - \boldsymbol{\Xi}_n \end{bmatrix} \begin{bmatrix} \boldsymbol{P}_{0, j-1} \\ \boldsymbol{P}_{1, j-1} \\ \vdots \\ \boldsymbol{P}_{n, j-1} \end{bmatrix} + \begin{bmatrix} \boldsymbol{R}_{0} \\ \boldsymbol{R}_{1} \\ \vdots \\ \boldsymbol{R}_{n} \end{bmatrix} \end{equation}$

Sovling an example system with initial value of $\rho=1.0$ at $x=40 \unicode{x212B}$ for all $n_w=0:9$ and boundary conditions with $\rho_L = \rho_R = 0$ using the potential profile of $n_w = 0$ .

In [7]:

# Non-adjustable constants
const kBT = 0.593 # kcal/mol
const kB = 1.38e-23
const T = 298.15
const r_ion = 1.75e-10
const r_water = 1.40e-10
const m_ion = 35.453
const m_water = 18.015
const NA = 6.02e23
const Λ3 = 4.87e-33
const μ_water = - 2.386 # kcal/mol

# Adjustable constants
const nx = 100
const dx = 1.e-10
const D = 1.69e-9
const steps = 1000
const nw = 10
const ρ_water = 130 # M/m^3

ΔE = [-12.14, -11.86, -11.42, -10.94, -10.47, -10.56, -10.35, -10.41, -10.35, -10.39] # kcal/mol
dt = 0.5 * dx^2 /(2 * D)

# Set grid and initial value(including boundary conditions) for dx.
# init_value = [[30, 1.0], [35, 0.8], [40, 0.6], [45, 0.4], [50, 0.2]]
init_value = []
for i = 1:nw
    push!(init_value, [40, 1.0])
end
grids = [] # [nx, steps] * nw
Lbound = 0.0
Rbound = 0.0
# Set potential profile fittings for every nw.
fits = [] # f * nw
for i = 1:nw
    push!(grids, zeros(Float64, nx, steps))
    grids[i][1, :] = Lbound
    grids[i][nx, :] = Rbound
    grids[i][Int(init_value[i][1]), 1] = init_value[i][2]
    push!(fits, fit)
end

λ = D * dt / dx^2
# Set η and τ array
ηs = [] # [nx - 2] * nw
τs = [] # [nx - 2] * nw

for i = 1:nw
    push!(ηs, [ - D * dt * fits[i][:grad](dx * x + 12.5e-10)[1] / (2 * kBT * dx) for x in 1:nx - 2])
    push!(τs, [ - D * dt * fits[i][:grad2](dx * x + 12.5e-10)[1] / kBT for x in 1:nx - 2])
end

# Set k_xx
k_evp =zeros(nw)
k_adp = zeros(nw)
k_wf = zeros(nw)

for i = 1:nw
    i - 1 < 4 ? ii = 4 : ii = i
    r_cluster = 1.2 * (r_ion^3 + r_water^3 * (ii - 1))^(1/3)
    μ = (m_ion + m_water * (i - 1)) * m_water / (m_ion + m_water * i) / NA
    k_adp_temp = π * (r_cluster + r_water)^2 * sqrt(8 * kB * T / (π * μ)) 
    if i < nw
        k_adp[i] = k_adp_temp * ρ_water * NA
        k_evp[i + 1] = k_adp_temp * i / Λ3 * exp((ΔE[i] - μ_water) / kBT)
    end
end

# Set up matrices A, B, Ξ, Z, K, R
matA = []
matB = []
matΞ = []
matZ = []
matK = []
matR = []
for i = 1:nw
    push!(matA, Tridiagonal([ - λ - ηs[i][x] for x in 2:nx - 2], 
            [2 + 2 * λ + τs[i][x] for x in 1:nx - 2], 
            [- λ + ηs[i][x] for x in 1:nx - 3]))
    push!(matB, Tridiagonal([ λ + ηs[i][x] for x in 2:nx - 2], 
            [2 - 2 * λ - τs[i][x] for x in 1:nx - 2], 
            [ λ - ηs[i][x] for x in 1:nx - 3]))
    push!(matΞ, Diagonal(fill((k_adp[i] + k_evp[i] + k_wf[i]) * dt, nx - 2)))
    push!(matZ, Diagonal(fill(k_evp[i] * dt, nx - 2))) # subdiagonal
    push!(matK, Diagonal(fill(k_adp[i] * dt, nx - 2))) # superdiagonal
    matRtmp = spzeros(nx - 2)
    matRtmp[1] = Lbound * 2 * (λ + ηs[i][1])
    matRtmp[nx - 2] = Rbound * 2 * (λ - ηs[i][end])
    push!(matR, matRtmp)
end

# Setup full matrices.
amatrix = hcat(matA[1] + matΞ[1], - matZ[2])
bmatrix = hcat(matB[1] - matΞ[1], matZ[2])
for i = 1:nw - 2
    amatrix = hcat(amatrix, zeros(nx - 2, nx - 2))
    bmatrix = hcat(bmatrix, zeros(nx - 2, nx - 2))
end
rmatrix = matR[1]
prepmatrix = grids[1][2:nx - 1, 1]
for i = 2:nw - 1
    rowA = hcat(-matK[i - 1], matA[i] + matΞ[i], - matZ[i + 1])
    rowB = hcat(matK[i - 1], matB[i] - matΞ[i], matZ[i + 1])
    for j = 1:i - 2
        rowA = hcat(zeros(nx - 2, nx - 2), rowA)
        rowB = hcat(zeros(nx - 2, nx - 2), rowB)
    end
    for j = i + 2: nw
        rowA = hcat(rowA, zeros(nx - 2, nx - 2))
        rowB = hcat(rowB, zeros(nx - 2, nx - 2))
    end
    amatrix = vcat(amatrix, rowA)
    bmatrix = vcat(bmatrix, rowB)
    rmatrix = vcat(rmatrix, matR[i])
    prepmatrix = vcat(prepmatrix, grids[i][2:nx - 1, 1])
end
rowA = hcat(-matK[nw - 1], matA[nw] + matΞ[nw])
rowB = hcat(matK[nw - 1], matB[nw] - matΞ[nw])
for i = 1:nw - 2
    rowA = hcat(zeros(nx - 2, nx - 2), rowA)
    rowB = hcat(zeros(nx - 2, nx - 2), rowB)
end
amatrix = vcat(amatrix, rowA)
bmatrix = vcat(bmatrix, rowB)
rmatrix = vcat(rmatrix, matR[nw])
prepmatrix = vcat(prepmatrix, grids[nw][2:nx - 1, 1])

# Solve tridiagonal system.
for i = 2:steps
    pmatrix = LAPACK.gesv!(copy(full(amatrix)), bmatrix * prepmatrix + rmatrix)
    # pmatrix = inv(full(amatrix)) * (bmatrix * prepmatrix + rmatrix)
    prepmatrix = pmatrix[1]
    for j = 1:nw
        grids[j][2:nx - 1, i] = pmatrix[1][1 + (nx - 2) * (j - 1):nx - 2 + (nx - 2) * (j - 1)]
    end
end

WARNING: redefining constant steps

In [8]:

anim = @animate for i=1:steps
    plot([grids[x][:, i] for x in 1:nw])
    plot!(ylims = (0, 1), yticks = 0:0.1:1, title="t = $(round(i * dt * 1.E12, 3)) ps")
    end every 10
gif(anim, "pics/coupled.gif", fps = 5)

INFO: Saved animation to /home/wanglj/PDE/pics/coupled.gif

Out[8]:

In [ ]: