Stochastic differential equation with callback in Julia - callback

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:
The first output value is just z0. I expected the reflection to be applied first, given that I set func_start = true.
The second value also being negative indicates two things:
The callback was not called prior to the first integrator call.
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)