Stochastic differential equation sensitivity analysis with specified noise

Question

I am trying to calculate the gradient of a functional of a stochastic differential equation (SDE) solution given a specific realization of the noise. I can successfully calculate these gradients if I leave the noise unspecified, as shown in DiffEqFlux.jl: Using Other Differential Equations. I can also successfully obtain the solution to my SDE for a specific noise realization, like shown in DifferentialEquations.jl: NoiseWrapper Example. When I try and put the two together, though, the code returns an error.

Here is a minimal working example adapted from the two separate examples referenced above:

using StochasticDiffEq, DiffEqBase, DiffEqNoiseProcess, DiffEqSensitivity, Zygote

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end
function lotka_volterra_noise(du,u,p,t)
  du[1] = 0.1u[1]
  du[2] = 0.1u[2]
end
dt = 1//2^(4)
u0 = [1.0,1.0]
p = [2.2, 1.0, 2.0, 0.4]
prob1 = SDEProblem(lotka_volterra,lotka_volterra_noise,u0,(0.0,10.0),p)
sol1 = solve(prob1,EM(),dt=dt,save_noise=true)

W2 = NoiseWrapper(sol1.W)
prob2 = SDEProblem(lotka_volterra,lotka_volterra_noise,u0,(0.0,10.0),p,noise=W2)
sol2 = solve(prob2,EM(),dt=dt)

function predict_sde1(p)
  Array(concrete_solve(remake(prob1,p=p),EM(),dt=dt,sensealg=ForwardDiffSensitivity(),saveat=0.1))
end
loss_sde1(p) = sum(abs2,x-1 for x in predict_sde1(p))

loss_sde1(p)

# This gradient is successfully calculated
Zygote.gradient(loss_sde1,p)

function predict_sde2(p)
  W2 = NoiseWrapper(sol1.W)
  Array(concrete_solve(remake(prob2,p=p,noise=W2),EM(),dt=dt,sensealg=ForwardDiffSensitivity(),saveat=0.1))
end
loss_sde2(p) = sum(abs2,x-1 for x in predict_sde2(p))

# This loss is successfully calculated
loss_sde2(p)

# This gradient calculation raises and error
Zygote.gradient(loss_sde2,p)

The error I get at the end of running this code is

TypeError: in setfield!, expected Float64, got ForwardDiff.Dual{Nothing,Float64,4}

Stacktrace:
 [1] setproperty! at ./Base.jl:21 [inlined]
...

followed by an interminable conclusion to the stacktrace (I can post it if you think it would be helpful, but since it's longer than the rest of this question I'd rather not clutter things up off the bat).

Is calculating gradients for SDE problems with specified noise realizations not currently supported, or am I just not making the appropriate function calls? I could easily believe the latter, since it was a bit of a struggle just to get to the point where the working parts of the above code worked, but I couldn't find any clue as to what I had incorrectly supplied after stepping through this code with the Juno debugger.

Chris Rackauckas · Accepted Answer · 2020-03-11T04:50:30.980

As a StackOverflow solution, you can use ForwardDiffSensitivity(convert_tspan=false) to work around this. Working code:

using StochasticDiffEq, DiffEqBase, DiffEqNoiseProcess, DiffEqSensitivity, Zygote

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end
function lotka_volterra_noise(du,u,p,t)
  du[1] = 0.1u[1]
  du[2] = 0.1u[2]
end
dt = 1//2^(4)
u0 = [1.0,1.0]
p = [2.2, 1.0, 2.0, 0.4]
prob1 = SDEProblem(lotka_volterra,lotka_volterra_noise,u0,(0.0,10.0),p)
sol1 = solve(prob1,EM(),dt=dt,save_noise=true)

W2 = NoiseWrapper(sol1.W)
prob2 = SDEProblem(lotka_volterra,lotka_volterra_noise,u0,(0.0,10.0),p,noise=W2)
sol2 = solve(prob2,EM(),dt=dt)

function predict_sde1(p)
  Array(concrete_solve(remake(prob1,p=p),EM(),dt=dt,sensealg=ForwardDiffSensitivity(convert_tspan=false),saveat=0.1))
end
loss_sde1(p) = sum(abs2,x-1 for x in predict_sde1(p))

loss_sde1(p)

# This gradient is successfully calculated
Zygote.gradient(loss_sde1,p)

function predict_sde2(p)
  Array(concrete_solve(prob2,EM(),prob2.u0,p,dt=dt,sensealg=ForwardDiffSensitivity(convert_tspan=false),saveat=0.1))
end
loss_sde2(p) = sum(abs2,x-1 for x in predict_sde2(p))

# This loss is successfully calculated
loss_sde2(p)

# This gradient calculation raises and error
Zygote.gradient(loss_sde2,p)

As a developer... this isn't a nice solution and our default should be better here. I'll work on this. You can track the development here https://github.com/JuliaDiffEq/DiffEqSensitivity.jl/issues/204. It'll probably get solved in an hour or so.

Edit: The fix is released and your original code works.

I noticed you changed `predict_sde2` to not use `remake`, and to not wrap the noise again. I had been re-wrapping the noise each time because I was observing some side effects where re-using the same wrapped noise was giving different solutions. Is there a simple reason why I was running into this problem and why your change is important, or should I open another question? — Jonathan A. Gross, Mar 13 '20 at 18:09
I don't think it's important: it should already have W2 in there, and internally it's calling a version of remake. But remaking yourself should work too. — Chris Rackauckas, Mar 14 '20 at 02:09

Stochastic differential equation sensitivity analysis with specified noise

1 Answers1