Where to find the code for ESRK1 and RSwM1 in the Julia library source code?

Question

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

Answer 1

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:

https://github.com/SciML/StochasticDiffEq.jl/blob/9d8eb5503f1d78cdb0de76691af2a89c20085486/src/tableaus.jl#L40

and the method (which can also work with other coefficients) here:

https://github.com/SciML/StochasticDiffEq.jl/blob/9d8eb5503f1d78cdb0de76691af2a89c20085486/src/perform_step/sri.jl#L58

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