Difficulties with DifferentialEquations package

Specific Domains Modelling & Simulations
1

I’m testing the DifferentialEquations package on a model I developed ca. 1993 (just dug it of the archives…). I try to simulate 5 Argon molecules bouncing around in 2D, locked up between “Lennard-Jones” like potential walls, and influenced by Lennard-Jones potentials around their neighboring molecules. When I did the original simulation, MATLAB took 7 hours to solve the model on a DECstation (hence, only 5 molecules and in 2D…). The code now runs in 9 seconds using MATLAB on my desktop PC, and 4 seconds in Julia.

I’m simulating the evolution over 500 ps, with fixed step length of 0.01 ps, with my own RK4 implementation. I have restructured the code somewhat, to make it fit better with the DifferentialEquations package. This restructuring has made my RK4 solver a little bit slower… I have not attempted to optimize the code…

The problem is that DifferentialEquations doesn’t find the solution – nothing “moves”… The problem is challenging, I guess, in that for most of the time, the model is “almost” pure integrators, while it becomes very stiff during molecule “collisions”… But I ask DifferentialEquations to “saveat = 0.01e-12” = 0.01 ps.

Here is the solution that my own RK4 solver gives… positions during the first 25 ps:

image
image1436×969 184 KB

… and here is the solution during the entire 500 ps:

image
image1436×960 668 KB

Anyone has an idea? [I’ll include my code in a follow-up…]

2

Here is my model of the molecule evolution…

# Argon function
function Argon(u,p,t)
	np = Int(ceil(length(u)/4))
	#
	x = u[1:np]
	y = u[np+1:2*np]
	px = u[2*np+1:3*np]
	py = u[3*np+1:4*np]
	#
	nw = 4
	Xk = p[1]
	Yk = p[2]
	sig = p[3]
	epsi = p[4]
	m = p[5]
	kB = p[6]
	#
	dfimdx=zeros(np)
	dfimdy=zeros(np)
	tmp = 1:np
	for i in 1:np
		r = zeros(np)
		tm = collect(copy(tmp))
		deleteat!(tm,i)
		for j in 1:np
			r[j] =sqrt((x[i]-x[j])^2 +(y[i]- y[j])^2)
		end
		for j in 1:nw
			k = tm[j]
			dfimdx[i] = dfimdx[i] +
			24*epsi*((sig/r[k])^6 - 2*(sig/r[k])^12)*(x[i] - x[k])/r[k]^2
			dfimdy[i] = dfimdy[i] +
			24* epsi*((sig/r[k])^6 - 2*(sig/r[k])^12)*(y[i] - y[k])/r[k]^2
		end
	end
	#
	dfivdx = zeros(np)
	dfivdy = zeros(np)
	for i in 1:np
		R = zeros(nw)
		Xg= zeros(nw)
		Yg= zeros(nw)
		L2 = zeros(nw)
		for k in 1:nw
			L2[k] = (Xk[2,k]-Xk[1,k])^2 +(Yk[2,k]-Yk[1,k])^2
			gam =((x[i]-Xk[1,k])*(Xk[2,k]-Xk[1,k]) +
			(y[i]-Yk[1,k])*(Yk[2,k]-Yk[1,k]))/L2[k]
			Xg[k] = Xk[1,k] + (Xk[2,k]-Xk[1,k])*gam
			Yg[k] = Yk[1,k] + (Yk[2,k]-Yk[1,k])*gam
			R[k] = sqrt((x[i]-Xg[k])^2 + (y[i]-Yg[k])^2)
		end
		for k in 1:nw
			dfivdx[i] = dfivdx[i] +
			24*epsi*((sig/R[k])^6 - 2*(sig/R[k])^12)*
			(x[i]-Xg[k])/R[k]^2*(Yk[2,k]-Yk[1,k])^2/L2[k]
			dfivdy[i] = dfivdy[i] +
			24*epsi*((sig/R[k])^6 - 2*(sig/R[k])^12)*
			(y[i]-Yg[k])/R[k]^2*(Xk[2,k]-Xk[1,k])^2/L2[k]
		end
	end
	fpx = - (dfimdx + dfivdx)*kB
	fpy = - (dfimdy + dfivdy)*kB
	fx = px/m
	fy = py/m
	du = [fx...,fy...,fpx...,fpy...]
	return du
end
#

Here is my RK4 solver:

# Runge Kutta 4 solver
function RK4(func,u,tspan,p)
	#
	#
	t0 = tspan[1]
	t1 = tspan[2]
	h = tspan[3]
	N = Int(ceil((t1-t0)/h))
	nu = length(u)
	#
	T = zeros(N+1)
	U = zeros(N+1,nu)
	#
	T[1] = t0
	U[1,:] = u'
	for i in 1:N
		#
		# -----------------------------------------------------------------
		# Runge-Kutta-4 updating
		#
		# Computing vector fields and temporary quantities k1-k4
		#
		k1 = func(u,p,0)
		k2 = func(u+0.5*h*k1,p,0)
		k3 = func(u+0.5*h*k2,p,0)
	 	k4 = func(u+h*k3,p,0)
		#
		# Updating states
		u = u + h*(k1+2*k2+2*k3+k4)/6
		# ------------------------------------------------------------------
		# Storing time and states
		#
		T[i+1] = T[i]+h
		U[i+1,:] = u'
		#
		# Next time
	end;
	#
	return T,U
end
#

Here is my set-up for the RK4 solver:

#
using Plots, LaTeXStrings, Statistics, DifferentialEquations; pyplot()
LS1 = :solid
LS2 = :dash
LS3 = :dot
LW1 = 2
LW2 = 1
#
#
# Setting up Parameters
m = 6.6359e-23		# mass per molecule [g]
kB = 1.3806e-23		# Boltzmann's constant - converts Kelvin to Joule
#
X01 = 0.			# wall position, Angstrom
X11 = 0				# --"--
X12 = 50			# --"--
X04 = 50			# --"--
X14 = X01
X02 = X11
X03 = X12
X13 = X04
Xk = [X01 X02 X03 X04;X11 X12 X13 X14]*1e-10
#
Y01 = 0.			# Angstrom
Y11 = 50			# --"--
Y12 = 50			# --"--
Y04 = 0				# --"--
Y14 = Y01
Y02 = Y11
Y03 = Y12
Y13 = Y04
Yk = [Y01 Y02 Y03 Y04;Y11 Y12 Y13 Y14]*1e-10
#
epsi = 119.8		# Lennard-Jones parameter, Kelvin
sig = 3.405			# Lennard Jones parameter, Angstrom
sig = sig*1e-10		# sig in m
m = m*1e-3			# mass in kg
#
p = [Xk,Yk,sig,epsi,m,kB]
#
# Initial values for positions
#
x = [5.,15,25,35,45] 	# Ångström
y = [5.,5,5,5,5]		# Ångstrom
x = x*1e-10
y = y*1e-10
#
# Initial values for momenta
px = m*[3.,3,-3,-3,3]	# kg*Angstrom/pico second
py = m*[3.,-3,3,-3,3] 	# --"--
px = px*1e2 			# conversion to kg*m/s
py = py*1e2 			# conversion to kg*m/s
#
u0 = [x...,y...,px...,py...]
# Time interval
#
t0 = 0. 			# picosecond
t1 = 500. 			# picosecond
h = 0.01			# picosecond
t1 = t1*1e-12 		# in second
h = h*1e-12			# in second
tspan1 = (t0,t1,h)
#
@time T,U = RK4(Argon,u0,tspan1,p)

N = length(T)
# Unpacking results
X = U[:,1:5]
Y = U[:,6:10]
PX = U[:,11:15]
PY = U[:,16:20]
# Convert position units from m to Angstrom
X=X*1e10
Y=Y*1e10

Plotting…

# Plotting results
plot(X[:,1],Y[:,1],lw=LW2,ls=LS1,lc=:black,label="#1")
plot!(X[:,2],Y[:,2],lw=LW2,ls=LS3,lc=:red,la=0.7,label="#2")
plot!(X[:,3],Y[:,3],lw=LW2,ls=LS3,lc=:gold,la=0.7,label="#3")
plot!(X[:,4],Y[:,4],lw=LW2,ls=LS3,lc=:cornflowerblue,la=0.7,label="#4")
plot!(X[:,5],Y[:,5],lw=LW2,ls=LS3,lc=:green,la=0.7,label="#5")
plot!(xlim=(0,50),ylim=(0,50))
plot!(xlabel=L"$x$ [Ångstrom]", ylabel=L"$y$ [Ångstrom]")
plot!(title=L"5 Argon molecules during $500$ ps")
#
Ns = Int(ceil(N/20))
plot(X[1:Ns,1],Y[1:Ns,1],lw=LW1,ls=LS1,lc=:black,label="#1")
plot!(X[1:Ns,2],Y[1:Ns,2],lw=LW1,ls=LS3,lc=:red,label="#2")
plot!(X[1:Ns,3],Y[1:Ns,3],lw=LW1,ls=LS3,lc=:gold,label="#3")
plot!(X[1:Ns,4],Y[1:Ns,4],lw=LW1,ls=LS3,lc=:cornflowerblue,label="#4")
plot!(X[1:Ns,5],Y[1:Ns,5],lw=LW1,ls=LS3,lc=:green,label="#5")
plot!(xlim=(0,50),ylim=(0,50))
plot!(xlabel=L"$x$ [Ångstrom]", ylabel=L"$y$ [Ångstrom]")
plot!(title=L"5 Argon molecules during $25$ ps")
scatter!((X[1,1],Y[1,1]),mc=:black,ms=6,label="")
scatter!((X[1,2],Y[1,2]),mc=:red,ms=6,label="")
scatter!((X[1,3],Y[1,3]),mc=:gold,ms=6,label="")
scatter!((X[1,4],Y[1,4]),mc=:cornflowerblue,ms=6,label="")
scatter!((X[1,5],Y[1,5]),mc=:green,ms=6,label="")
3

OK… and then I tried the DifferentialEquations approach:

# Using DifferentialEquations package
# Setting up problem
tspan = (t0,t1)
prob = ODEProblem(Argon,u0,tspan,p)
#
sol = solve(prob, saveat=0.01e-12) #reltol=1e-6)
#
plot(sol[1,:]*1e10,sol[6,:]*1e10)

I must have done something wrong, because nothing “moves” when I use the DifferentialEquations package… Suggestions are appreciated!

4

Using fixed time steps worked:

sol = solve(prob, RK4(), adaptive=false, dt=h) #reltol=1e-6)
#
plot(sol[1,:]*1e10,sol[6,:]*1e10)

sol = solve(prob, Tsit5(), adaptive=false, dt=h)
plot(sol[1,:]*1e10,sol[6,:]*1e10)

That’s when I realized that the issue is that the default tolerances, like abstol=1e-6, are just not appropriate for your model. They need to be tweaked if you want to use adaptivity:

sol = solve(prob, Vern9(), saveat=h, reltol=1e-6,abstol=1e-14)
plot(sol[1,:]*1e10,sol[6,:]*1e10)

That for example works just fine.

I have three things to mention. First, for speed, I would recommend setting up an inplace Argon(du,u,p,t). Secondly, you might want to check out GitHub - SciML/NBodySimulator.jl: A differentiable simulator for scientific machine learning (SciML) with N-body problems, including astrophysical and molecular dynamics . Third, usually you want to use different units so that way your values are closer to one.

2 Likes
5

Thanks for excellent suggestions! I’ll take a look at NBodySimulator. [Way back, some of our Ph.D. students worked with molecule simulations as part of thermodynamics studies, and I played with relevant example to illustrate Lagrangian mechanics… My model is unnecessarily complicated because of the wall model I used to lock in the molecules. In molecular simulations, the standard trick to keep the mass constant when a molecule moves out of the box, is to inject a new molecule in position modula the box length, I think.]

1 Like
6

NBodySimulator.jl looks interesting – I’ll look into it later.

I have done some more testing…

sol = solve(prob,Tsit5(),adaptive=false, dt=h)

works, and gives the correct solution.
However, if I change the function signature from:

function Argon(u,p,t)

to:

function Argon(du,u,p,t)

and ask Julia to solve the model by

sol = solve(prob,Tsit5(),adaptive=false, dt=h)

[no other changes than changing the function call)] the old problem re-surfaces: no change is computed, and everything is constant in the positions.

I also tried to use adaptive step length with the Vern9() solver, and Julia cast an error message.

When using the DifferentialEquations framework, finding the solution takes some 30+ seconds on my laptop. If I use my own RK4 solver, the solution is found in ca. 20 seconds. If I use my original code which did not involve function calls, but only nested for loops, finding the solution on my laptop took ca. 9 seconds.

I realize that the DifferentialEquations framework has some overhead (e.g., computing interpolators, etc.), but I’m somewhat surprised that my code with my RK4 solver is slower than the for loop implementation. I’ll have to check more on that, to understand performance issues.

7

In the mean time, I’ll try to scale the problem to use Angstrom or nanometers [NBodySimulator appears to use nm unit for molecule simulations], and picoseconds.

8

No, you have to change the function call of course. It should be a mutating function that updates the values in du instead of building any new arrays.

The out of place code isn’t broadcasting anything because it was made for things like StaticArrays and not Arrays. The inplace code was made for arrays.

9

In my original code, I did:

function Argon(u,p,t)
...
   du = [...]
end

Therefore, I guess I am creating array du every time I call Arg(u,p,t).

In the new version, I do:

function Argon(du,u,p,t)
...
   du = [...]
end

Is this sufficient to make the code mutating? I compute du at the last line in the function. Thus, I assume that I fill in values in the array du in the function call, and do not create a new array inside of the function??

10

No. du .= [...] would be. But you’d still have a ton of temporary allocations and missing a lot of possible performance.

11

Just to add a bit to what Chris was implying; your version of argon(u,p,t) looks to create (i.e. allocate) a lot of new arrays every time it is called. Using the in-place version argon(du,u,p,t) will allow you to update the passed-in derivative array instead of reallocating it. Similarly, you might want to pass in cache arrays within p instead of the many places where you are currently calling zeros. You might also want to use views for accessing slices of u instead of making copies of those slices, which is what your initial method was doing.

12

Thanks for detailed tips. I’ll start checking out view(), etc.

One question: I used [vec1...,vec2...,vec3...,vec4...] to append 4 vectors instead of append!(vec1,append!(vec2,append!(vec3,vec4))). Is there a penalty to this?

13

Yes, splatting has quite a penalty. But appending does as well. There is definitely a way to reformulate the algorithm to not require any of these allocations.

14

If you haven’t already looked through it, you might also look at the FAQ question: My ODE is solving really slow… what do I do?

This has some info on how you can check that your function is not allocating memory each time it is called, and info on how to look for type instabilities in your code (which can also lead to dramatic slow downs). I haven’t look at your code closely enough to see if there are any immediately obvious instabilities. (I usually find type instabilities are the most common performance degradation in the first version of my own codes – it’s the biggest Julia coding gotcha I think.)

1 Like