I'm trying to implement the SDE solver called ESRK1 and the adaptive stepsize algorithm called RSwM1 from Rackauckas & Nie (2017). I'm writing a python implementation, mainly to confirm to myself that I've understood the algorithm correctly. However, I'm running into a problem already at the implementation of ESRK1: When I test my implementation with shorter and shorter timesteps on a simple SDE describing geometric Brownian motion, the solution does not converge as dt becomes smaller, indicating that I have a mistake in my code.
I believe this algorithm is implemented as part of the library DifferentialEquations.jl in Julia, so I thought perhaps I could find some help by looking at the Julia code. However, I have had some trouble locating the relevant code. If someone could point me to the implementation of ESRK1 and RSwM1 in the relevant Julia librar(y/ies) (or indeed any other readable and correct implementation) of these algorithms, I would be most grateful.
I searched for ESRK and RSwM in the github repo of StochasticDiffEq.jl, but I didn't find anything I could really recognise as the method from the paper I'm reading:
https://github.com/search?q=repo%3ASciML%2FStochasticDiffEq.jl+rswm&type=code
Update: I found the code for ESRK1, as shown in my answer below, but I'm still unable to find the code for RSwM1.
For completeness, here is my own not-yet-correct implementation of ESRK1 in python:
def ESRK1(U, t, dt, f, g, dW, dZ):
# Implementation of ESRK1, following Rackauckas & Nie (2017)
# Eq. (2), (3) and (4) and Table 1
# Stochastic integrals, taken from Eqs. (25) - (30) in Rackauckas & Nie (2017)
I1 = dW
I11 = (I1**2 - dt) / 2
I111 = (I1**3 - 3*dt*I1) / 6
I10 = (I1 + dZ/np.sqrt(3))*dt / 2
# Coefficients, taken from Table 1 in Rackauckas & Nie (2017)
# All coefficients not included below are zero
c0_2 = 3/4
c1_2, c1_3, c1_4 = 1/4, 1, 1/4
A0_21 = 3/4
B0_21 = 3/2
A1_21 = 1/4
A1_31 = 1
A1_43 = 1/4
B1_21 = 1/2
B1_31 = -1
B1_41, B1_42, B1_43 = -5, 3, 1/2
alpha1, alpha2 = 1/2, 2/3
alpha_tilde1, alpha_tilde2 = 1/2, 1/2
beta1_1, beta1_2, beta1_3 = -1, 4/3, 2/3
beta2_1, beta2_2, beta2_3 = -1, 4/3, -1/3
beta3_1, beta3_2, beta3_3 = 2, -4/3, -2/3
beta4_1, beta4_2, beta4_3, beta4_4 = -2, 5/3, -2/3, 1
# Stages in the Runge-Kutta approximation
# Eqs. (3) and (4) and Table 1 in Rackauckas & Nie (2017)
# First stages
H0_1 = U # H^(0)_1
H1_1 = U
# Second stages
H0_2 = U + A0_21 * f(t, H0_1)*dt + B0_21 * g(t, H1_1)*I10/dt
H1_2 = U + A1_21 * f(t, H0_1)*dt + B1_21 * g(t, H1_1)*np.sqrt(dt)
# Third stages
H0_3 = U
H1_3 = U + A1_31 * f(t, H0_1) * dt + B1_31 * g(t, H1_1) * np.sqrt(dt)
# Fourth stages
H0_4 = U
H1_4 = U + A1_43 * f(t, H0_3) * dt + (B1_41 * g(t, H1_1) + B1_42 * g(t+c1_2*dt, H1_2) + B1_43 * g(t+c1_3*dt, H1_3)) * np.sqrt(dt)
# Construct next position
# Eq. (2) and Table 1 in Rackauckas & Nie (2017)
U_ = U + (alpha1*f(t, H0_1) + alpha2*f(t+c0_2*dt, H0_2))*dt \
+ (beta1_1*I1 + beta2_1*I11/np.sqrt(dt) + beta3_1*I10/dt ) * g(t, H1_1) \
+ (beta1_2*I1 + beta2_2*I11/np.sqrt(dt) + beta3_2*I10/dt ) * g(t + c1_2*dt, H1_2) \
+ (beta1_3*I1 + beta2_3*I11/np.sqrt(dt) + beta3_3*I10/dt ) * g(t + c1_3*dt, H1_3) \
+ (beta4_4*I111/dt ) * g(t + c1_4*dt, H1_4)
# Calculate error estimate
# Eq. (9) and Table 1 in Rackauckas & Nie (2017)
E = -dt*(f(t, H0_1) + f(t + c0_2*dt, H0_2))/6 \
+ (beta3_1*I10/dt + beta4_1*I111/dt)*g(t, H1_1) \
+ (beta3_2*I10/dt + beta4_2*I111/dt)*g(t + c1_2*dt, H1_2) \
+ (beta3_3*I10/dt + beta4_3*I111/dt)*g(t + c1_3*dt, H1_3) \
+ (beta4_4*I111/dt)*g(t + c1_4*dt, H1_4)
# Return next position and error
return U_, E
Rackauckas & Nie (2017): https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5844583/pdf/nihms920388.pdf
So, I guess I found the first half of the answer to my question: The code for the ESRK1 method appears to be found in two places in StochasticDiffEq.jl:
The coefficients from the paper (although they are now called SRIW1 instead of ESRK1) here:
and the method (which can also work with other coefficients) here:
It's not super readable, at least not to a non-Julia programmer like me, as it's making use of some advanced features, but I think I'll be able to work it out with a bit of patience.
Update:
By translating the Julia code into Python I was able to write an implementation that at least seems to work in the sense that I observe an order 1.5 convergence in the strong sense when I'm testing it as if it was a fixed-step integrator. In case it may be of use to anyone else, here is the line-for-line translation.
# c₀ = [0; 3//4; 0; 0]
c0 = np.array([0, 3/4, 0, 0])
# c₁ = [0; 1//4; 1; 1//4]
c1 = np.array([0, 1/4, 1, 1/4])
# A₀ = [0 0 0 0
# 3//4 0 0 0
# 0 0 0 0
# 0 0 0 0]
A0 = np.array([
[0, 0, 0, 0],
[3/4, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0],
])
# A₁ = [0 0 0 0
# 1//4 0 0 0
# 1 0 0 0
# 0 0 1//4 0]
A1 = np.array([
[0, 0, 0, 0],
[1/4, 0, 0, 0],
[1, 0, 0, 0],
[0, 0, 1/4, 0],
])
# B₀ = [0 0 0 0
# 3//2 0 0 0
# 0 0 0 0
# 0 0 0 0]
B0 = np.array([
[0, 0, 0, 0],
[3/2, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0],
])
# B₁ = [0 0 0 0
# 1//2 0 0 0
# -1 0 0 0
# -5 3 1//2 0]
B1 = np.array([
[0, 0, 0, 0],
[1/2, 0, 0, 0],
[-1, 0, 0, 0],
[-5, 3, 1/2, 0],
])
# α = [1//3; 2//3; 0;0]
alpha = np.array([ 1/3, 2/3, 0, 0])
# β₁ = [-1; 4//3; 2//3; 0]
beta1 = np.array([-1, 4/3, 2/3, 0])
# β₂ = -[1; -4//3; 1//3; 0]
beta2 = np.array([-1, 4/3, -1/3, 0])
# β₃ = [2; -4//3; -2//3; 0]
beta3 = np.array([ 2, -4/3, -2/3, 0])
# β₄ = [-2; 5//3; -2//3; 1]
beta4 = np.array([-2, 5/3, -2/3, 1])
def ESRK1(U, t, dt, f, g, dW=None, dZ=None):
# sqrt3 = sqrt(3one(eltype(W.dW)))
sqrt3 = np.sqrt(3)
# chi1 = (W.dW.^2 - abs(dt))/2integrator.sqdt #I_(1,1)/sqrt(h)
chi1 = (dW**2 - np.abs(dt)) / (2*np.sqrt(dt))
# chi2 = (W.dW + W.dZ/sqrt3)/2 #I_(1,0)/h
chi2 = (dW + dZ/sqrt3)/2
# chi3 = (W.dW.^3 - 3W.dW*dt)/6dt #I_(1,1,1)/h
chi3 = (dW**3 - 3*dW*dt) / (6*dt)
# for i=1:stages
# fill!(H0[i],zero(eltype(integrator.u)))
# fill!(H1[i],zero(eltype(integrator.u)))
# end
stages = 4
H0 = np.zeros(stages)
H1 = np.zeros(stages)
# for i = 1:stages
# fill!(A0temp,zero(eltype(integrator.u)))
# fill!(B0temp,zero(eltype(integrator.u)))
# fill!(A1temp,zero(eltype(integrator.u)))
# fill!(B1temp,zero(eltype(integrator.u)))
for i in range(stages):
A0temp = 0.0
B0temp = 0.0
A1temp = 0.0
B1temp = 0.0
# for j = 1:i-1
# integrator.f(ftemp,H0[j],p,t + c₀[j]*dt)
# integrator.g(gtemp,H1[j],p,t + c₁[j]*dt)
# @.. A0temp = A0temp + A₀[j,i]*ftemp
# @.. B0temp = B0temp + B₀[j,i]*gtemp
# @.. A1temp = A1temp + A₁[j,i]*ftemp
# @.. B1temp = B1temp + B₁[j,i]*gtemp
# end
for j in range(i):
ftemp = f(H0[j], t+c0[j]*dt)
gtemp = g(H1[j], t+c1[j]*dt)
A0temp = A0temp + A0[i,j]*ftemp
B0temp = B0temp + B0[i,j]*gtemp
A1temp = A1temp + A1[i,j]*ftemp
B1temp = B1temp + B1[i,j]*gtemp
# @.. H0[i] = uprev + A0temp*dt + B0temp*chi2
# @.. H1[i] = uprev + A1temp*dt + B1temp*integrator.sqdt
# end
H0[i] = U[0] + A0temp*dt + B0temp*chi2
H1[i] = U[0] + A1temp*dt + B1temp*np.sqrt(dt)
# fill!(atemp,zero(eltype(integrator.u)))
# fill!(btemp,zero(eltype(integrator.u)))
# fill!(E₂,zero(eltype(integrator.u)))
# fill!(E₁temp,zero(eltype(integrator.u)))
atemp = 0.0
btemp = 0.0
E2 = 0.0
E1temp = 0.0
# for i = 1:stages
# integrator.f(ftemp,H0[i],p,t+c₀[i]*dt)
# integrator.g(gtemp,H1[i],p,t+c₁[i]*dt)
# @.. atemp = atemp + α[i]*ftemp
# @.. btemp = btemp + (β₁[i]*W.dW + β₂[i]*chi1)*gtemp
# @.. E₂ = E₂ + (β₃[i]*chi2 + β₄[i]*chi3)*gtemp
# if i <= error_terms
# @.. E₁temp += ftemp
# end
for i in range(stages):
ftemp = f(H0[i], t+c0[i]*dt)
gtemp = g(H1[i], t+c1[i]*dt)
atemp = atemp + alpha[i]*ftemp
btemp = btemp + (beta1[i]*dW + beta2[i]*chi1)*gtemp
E2 = E2 + (beta3[i]*chi2 + beta4[i]*chi3)*gtemp
# @.. u = uprev + (dt*atemp + btemp) + E₂
return U + (dt*atemp + btemp) + E2