Stochastic differential equation with callback in Julia

Question

I'm trying to solve a diffusion problem with reflecting boundaries, using various SDE integrators from DifferentialEquations.jl. I thought I could use the FunctionCallingCallback to handle the boundaries, by reflecting the solution about the domain boundaries after every integrator step.

This is my code

using DifferentialEquations

K0 = 1e-3
K1 = 5e-3 
alpha = 0.5

K(z)    = K0 + z*K1*exp(-z*alpha)
dKdz(z) = K1*exp(-alpha*z) - K1*alpha*z*exp(-alpha*z)

a(z,p,t) = dKdz(z)
b(z,p,t) = sqrt(2*K(z))

dt = 0.1
tspan = (0.0,1.0)

z0 = 1.0

prob  = SDEProblem(a,b,z0,tspan)

function reflect(z, p, t, integrator)
    bottom = 2.0
    if z < 0
        # Reflect from surface
        z = -z
    elseif z > bottom
        # Reflect from bottom
        z = 2*bottom - z
    end
    return z
end

cb = FunctionCallingCallback(reflect;
                 func_everystep = true,
                 func_start = true,
                 tdir=1)

sol = solve(prob, EM(), dt = dt, callback = cb)

Edit: After solving my initial problem thanks to the comment by Chris Rackauckas, I modified my reflect function. Now the code runs, but the solution contains negative values, which should have been prevented by reflection about 0 after every step.

Any ideas as to what's going wrong here would be greatly appreciated.

Note by the way that the FunctionCallingCallback example found here contains two different function signatures for the callback function, but I get the same problem with both. It's also not clear to me if the callback should modify the value of z in place, or return the new value.

Edit 2: Based on Chris Rackauckas' answer, and looking at this example, I've modified by reflect function thus:

function reflect(z, t, integrator)
    bottom = 2.0
    if integrator.u < 0
        # Reflect from surface
        integrator.u = -integrator.u
    elseif integrator.u > bottom
        # Reflect from bottom
        integrator.u = 2*bottom - integrator.u
    end
    # Not sure if the return statement is required
    return integrator.u
end

Running this with initial condition z0 = -0.1 produces the following output:

retcode: Success
Interpolation: 1st order linear
t: 11-element Array{Float64,1}:
 0.0                
 0.1                
 0.2                
 0.30000000000000004
 0.4                
 0.5                
 0.6                
 0.7                
 0.7999999999999999 
 0.8999999999999999 
 1.0                
u: 11-element Array{Float64,1}:
 -0.1                 
 -0.08855333388147717 
  0.09862543518953905 
  0.09412012313587219 
  0.11409372573454478 
  0.10316400521980074 
  0.06491042188420941 
  0.045042097789392624
  0.040565317051189105
  0.06787136817395374 
  0.055880083559589955

It seems to me that what's happening here is:

1. The first output value is just z0. I expected the reflection to be applied first, given that I set func_start = true.
2. The second value also being negative indicates two things:
   1. The callback was not called prior to the first integrator call.
   2. The callback was not called prior to storing the output after the first integrator call.

I would have expected all the values in the output to be positive (i.e., have the callback applied to them before storing the output). Am I doing something wrong, or should I simply adjust my expectations?

Answer 1

The FunctionCallingCallback is a function (u,t,integrator), so I'm not sure how your code didn't error for you. It should be:

using DifferentialEquations

K0 = 1e-3
K1 = 5e-3
alpha = 0.5

K(z)    = K0 + z*K1*exp(-z*alpha)
dKdz(z) = K1*exp(-alpha*z) - K1*alpha*z*exp(-alpha*z)

a(z,p,t) = dKdz(z)
b(z,p,t) = sqrt(2*K(z))

dt = 0.1
tspan = (0.0,1.0)

z0 = 1.0

prob  = SDEProblem(a,b,z0,tspan)

function reflect(z, t, integrator)
    bottom = 2.0
    if z < 0
        # Reflect from surface
        z = -z
    elseif z > bottom
        # Reflect from bottom
        z = 2*bottom - z
    end
    return z
end

cb = FunctionCallingCallback(reflect;
                 func_everystep = true,
                 func_start = true,
                 tdir=1)

sol = solve(prob, EM(), dt = dt, callback = cb)

Edit

You don't want the function calling callback. Just use the normal callback:

using DifferentialEquations

K0 = 1e-3
K1 = 5e-3
alpha = 0.5

K(z)    = K0 + z*K1*exp(-z*alpha)
dKdz(z) = K1*exp(-alpha*z) - K1*alpha*z*exp(-alpha*z)

a(z,p,t) = dKdz(z)
b(z,p,t) = sqrt(2*K(z))

dt = 0.1
tspan = (0.0,1.0)

z0 = 1.0

prob  = SDEProblem(a,b,z0,tspan)

condition(u,t,integrator) = true
function affect!(integrator)
    bottom = 2.0
    if integrator.u < 0
        # Reflect from surface
        integrator.u = -integrator.u
    elseif integrator.u > bottom
        # Reflect from bottom
        integrator.u = 2*bottom - integrator.u
    end
end

cb = DiscreteCallback(condition,affect!;save_positions=(false,false))

sol = solve(prob, EM(), dt = dt, callback = cb)