# Nested derivatives with DiffEqFlux.jl

I’m working on a physics project where (very roughly) I need to integrate some Hamiltonian equations of motion and differentiate functions of the resulting trajectories with respect to a parameter \alpha. More precisely, I want to solve the following system of equations for x(t; \alpha) and p(t; \alpha):

\frac{dx}{dt}(t; \alpha) = \frac{\partial H}{\partial p}(x, p, t; \alpha)\\ \frac{dp}{dt}(t; \alpha) = -\frac{\partial H}{\partial x}(x, p, t; \alpha),

and then optimize some functions x(t; \alpha) and p(t; \alpha) with respect to \alpha. Here’s my attempt at this using HamiltonianProblem from DiffEqPhysics.jl (modified version of the DiffEqFlux.jl examples):

using DifferentialEquations
using DiffEqFlux
using Flux
using ForwardDiff

function H(p, x, params)  # complicated Hamiltonian
return params[1] * (1 + x[1]^2)^(1/3) * (1 + p[1]^2)^(1/4)  # randomly chosen function
end

# Solve the differential equation
function integrateeoms(params)
# Time interval to solve over
tspan = (0.0, 1.0)
# Initial state
p0, x0 = [0.05], [0.01]
# Construct ODE problem
prob = HamiltonianProblem(H, p0, x0, tspan, params)
# Integrate the ODE
Tracker.collect(diffeq_rd(params, prob, VelocityVerlet(), dt=0.1, saveat=0.1))
end


This runs fine if I pass in a plain array (eg, integrateeoms([1.0])), but fails if I try to track the parameter array (eg, integrateeoms(param([1.0]))). I get the following error followed by a huge stack trace:

ERROR: MethodError: *(::Tracker.TrackedReal{Float64}, ::ForwardDiff.Dual{ForwardDiff.Tag{DiffEqPhysics.PhysicsTag,Float64},Float64,1}) is ambiguous


Since H is too complicated to analytically differentiate, how can I go about using nested autodifferentiation here?

Internally, HamiltonianProblem uses ForwardDiff.jl to take partial derivatives of H. Is there a way to selectively turning off tracking for x and p after taking those derivatives to avoid the type problems?

Thanks!

Ping me in a week and hopefully I have a solution

Ok, awesome, thanks!

I should mention that my actual problem is more complicated and involves more nested derivatives. The system I’m modelling is nonconservative, so there are other terms in the equations of motion that depend on partial derivatives of H. In addition the initial conditions depend on finding a root of a function of the second derivative of H.

Ok, I got a simple example working with ForwardDiff.jl:

using DifferentialEquations
using ForwardDiff

function H(x, p, k, m)  # complicated Hamiltonian
return p^2 / (2 * m) + 1/2 * k * x^2
end

function dHdu(x, p, k, m)
# Differentiate the Hamiltonian
dHdx = ForwardDiff.derivative((x̄) -> H(x̄, p, k, m), x)
dHdp = ForwardDiff.derivative((p̄) -> H(x, p̄, k, m), p)
return dHdx, dHdp
end

function eoms!(dudt, u, params, t)
# println("(dudt, u, p, t) = ($dudt,$u, $p,$t)")
x, p = u
k, m = params

# Construct equations of motion
dHdx, dHdp = dHdu(x, p, k, m)
dudt[1] = dHdp
dudt[2] = -dHdx
end

function integrateeoms(p)
u0 = [0.1, 0.2]
tspan = (0.0, 1.0)
println("solving...")
prob = ODEProblem(eoms!, eltype(p).(u0), eltype(p).(tspan), p)
return solve(prob, Tsit5(), save_everystep=false)[end][1]
end


The last function can be differentiated with ForwardDiff.gradient(integrateeoms, [2.0, 0.01]).

Unfortunately, my equations of motion contain complex numbers in intermediate steps, which triggers a known StackOverflowError in base/complex.jl [1], [2], [3]. Any ideas about how to get this to work (with Flux, Zygote, defining a custom sqrt function, or whatever) would be much appreciated!

XRef: https://github.com/JuliaDiffEq/DiffEqFlux.jl/issues/53#issuecomment-506971079 It looks like the solution we are building towards will be in Zygote, and so I think this may be put on hold for a Zygote transition.