Fokker-planck equation solved with FEniCs

fokker-planck_fem.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d, Axes3D
from dolfin import *
from fenics import *
from matplotlib import cm


# model parameters
γ_GFR = 1.  # rate of GFR reaching its steady state
γ_ERM = 1.  # rate of ERM reaching its steady state
w_GFR = -2.56  # Basal inhibition of GFR
w_GFR_E2ER = -1.6  # GFR inhibition by E2ER
w_GFR_GFR = 6.8  # GFR self activation
w_GFR_ERM = 0.4  # GFR activation by ERM
w_ERM = -3.04  #  Basal inhibition of EPM 
w_ERM_E2ER = -1.6  # EPM inhibition by E2ER
w_ERM_ERM = 6.64  #  ERM self activation
w_ERM_GFR = 0.4  # ERM activation by GFR
w_E2ER = 5.0  # E2ER self activation
w_E2ER_E2 = 5. # E2ER activation by E2
E2 = -2
E2ER = 1./(1. + np.exp(-(w_E2ER + w_E2ER_E2 * E2))) 


T = 1000           # final time
num_steps = 1000    # number of time steps
dt = 1. #T / num_steps # time step size
output_steps = 1000

print("time step", dt)
# Create mesh and define function space
nx = ny = 51

mesh = RectangleMesh(Point(-0.2, -0.2), Point(1.2, 1.2), nx, ny)
n_ = FacetNormal(mesh)
V = FunctionSpace(mesh, 'P', 1)
W = VectorFunctionSpace(mesh, 'P', 2)

velocity = Expression(('1/(1+exp(-(w0 + w02 * x2 + w00 * x[0] + w01 * x[1]) ))  - x[0]', # GFR
                               '1/(1+exp(-(w1 + w12 * x2 + w11 * x[1] + w10 * x[0])))  - x[1]'), # ERM
                              x2 = E2ER, 
                              w0 = w_GFR,
                              w02 = w_GFR_E2ER,
                              w00 = w_GFR_GFR,
                              w01 = w_GFR_ERM, 
                              w1 = w_ERM, 
                              w12 = w_ERM_E2ER, 
                              w11 = w_ERM_ERM, 
                              w10 = w_ERM_GFR, degree=2)

vel = interpolate(velocity, W)

# plot velocity field
plot(vel)
plt.xlabel('GFR')
plt.ylabel('ERM')
plt.show()


# Define variational problem
u = TrialFunction(V) 
v = TestFunction(V)
u_ = Function(V)
u_n = Function(V)


# Define initial distribution
u0 = Expression(' exp(- (pow(x[0]-0.5, 2) + pow(x[1]-0.5, 2)) )',  degree=1)
#u0 = Expression('100.0',degree=0)

u_n = interpolate(u0, V)

k = Constant(dt)


# forward finite differences
#  diffusion (volume)
#  advection (volume)
#  diffusion (surface)
#  advection (surface)
F =   dot( (u - u_n)/k, v) * dx +\
    inner(grad(u), grad(v)) * dx \
     -   u * inner(vel, grad(v))  * dx \
     +  v * inner(grad(u), n_) * ds\
     -  u * v * inner(vel, n_) * ds \
     

a, L = lhs(F), rhs(F)

# Create VTK file for saving solution
vtkfile = File('fokker-planck.pvd')

# Time-stepping

t = 0


A = assemble(a) 
b = assemble(L)


for n in range(num_steps):
    
    t += dt

    if n % output_steps == 0:
        u_n.rename("probability", "")
        vtkfile << (u_n, t)
        im1 = plot(u_n)
        plt.colorbar(im1)
        plt.xlabel('GFR')
        plt.ylabel('ERM')
        plt.show()
        print(n)
        
    A = assemble(a)
    b = assemble(L)
    
   # Compute solution
    solve(A, u_.vector(), b) 

   # Update previous solution
    u_n.assign(u_)
    
    

im1 = plot(u_n)
plt.colorbar(im1)
plt.xlabel('GFR')
plt.ylabel('ERM')
plt.show()

u_n.rename("probability", "")

vtkfile << (u_n, t)