Nested derivatives with DiffEqFlux.jl

Specific Domains
,
1

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!

2

Ping me in a week and hopefully I have a solution

3

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.

4

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!

5

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.