How to autodifferentiate an ODEsolver?

[dkarrasch]

I am interested in autodifferentiating the map that takes an initial value to its "final" value of an ODE solution. Taking the Lorenz example from the docs, I tried the following:

using DifferentialEquations, ForwardDiff

g = @ode_def LorenzExample begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ=>10.0 ρ=>28.0 β=(8/3)

u0 = [1.0;0.0;0.0]
tspan = (0.0,1.0)
dg = ForwardDiff.jacobian(x -> g(0.,x),u0)

function flow(g,u0,tspan)
  prob = ODEProblem(g,u0,tspan)
  sol = solve(prob,dt=.1).u[end]
end

flow(g,u0,tspan)
dflow = ForwardDiff.jacobian(x -> flow(g,x,tspan),u0)

So, all functions are defined properly, ODE-solving works, vector field autodifferentiation works, but not autodifferentiation of flow. The error message that I get is:

MethodError: Cannot convert an object of type ForwardDiff.Dual{ForwardDiff.Tag{##12#13,0xb046287d533b082d},Float64,3} to an object of type Float64 This may have arisen from a call to the constructor Float64(...), since type constructors fall back to convert methods. in jacobian at ForwardDiff/src/jacobian.jl:20 in jacobian at ForwardDiff/src/jacobian.jl:21 in vector_mode_jacobian at ForwardDiff/src/jacobian.jl:132 in at base/console:1 in flow at autodiff_test.jl:15 in at base/ in #solve#1 at DifferentialEquations/src/default_solve.jl:5 in at base/ in #solve#2 at DifferentialEquations/src/default_solve.jl:14 in at base/ in at base/ in #solve#1830 at OrdinaryDiffEq/src/solve.jl:7 in solve! at OrdinaryDiffEq/src/solve.jl:304 in perform_step! at OrdinaryDiffEq/src/integrators/low_order_rk_integrators.jl:864

Is there an issue with my definition of flow and how I pick the solution state from the ODEsolution-structure, or is it with the solve-function? Or is the idea to define the equation of variations yourself, using the autoderivative of the vector field? Any help is much appreciated!

[ChrisRackauckas]

In order to differentiate through an ODE solver which has adaptive timestepping, you need to make sure that not just the dependent variables are dual numbers, but also the independent variables. This is because changes in x changes the points which are stepped to, so these need to be duals as well. So without testing it,

function flow(g,u0,tspan)
  prob = ODEProblem(g,u0,eltype(u0).(tspan))
  sol = solve(prob,dt=.1).u[end]
end

should be all that's needed (so it's dual when u0 is dual, and not dual when u0 is not).

In general it may be more efficient to solve the sensitivity equations than to do this though, but that's hard to know right now.

[dkarrasch]

Many thanks, Chris. It works without complaining. I'll do some testing and comparisons at some point.