DiffEqFlux.jl: Questions about Neural ODEs

Specific Domains Machine Learning
1

I’m relatively new to both Julia and deep learning. I want to learn the dynamical model of a system using a set of observed trajectories. I’ve heard neural ODEs may be of promise, but they seem quite new, and I haven’t been able to find many usage examples. I have some questions about neural ODES in general, and about a specific basic example I’ve been working on. Below is the example:

using Plots, DifferentialEquations, Flux, DiffEqFlux

function projectiletrajectory(du, u, p, t)
    cd1, cd2, cd3, a1, a2, cl1, srd, scm, ρ, μ, g, diameter, mass, area = p
    v_x, v_y, v_z, ω_x, ω_y, ω_z = u[4:9]
    # Speed and spin
    v = sqrt(v_x^2 + v_y^2 + v_z^2)
    ω = sqrt(ω_x^2 + ω_y^2 + ω_z^2)
    # Reynolds number
    re = ρ * v * diameter / μ
    # Spin coefficient
    sc = 0.5 * ω * diameter / v
    # Lift coefficient
    cl = cl1 * (sc^scm)
    # Drag coefficient
    cd = cd1 + cd2 * sc + cd3 * sin(π * (re-a1)/a2)
    # Differential equations
    du[1] = v_x
    du[2] = v_y
    du[3] = v_z
    du[4] = -0.5 * ρ * (area / mass) * v * (cd * v_x - cl * ( (ω_y / ω) * v_z - (ω_z / ω) * v_y )  )
    du[5] = - 0.5 * ρ * (area / mass) * v * (cd * v_y - cl * ( (ω_z / ω) * v_x - (ω_x / ω) * v_z ) )
    du[6] = -g - 0.5 * ρ * (area / mass) * v * (cd * v_z - cl * ( (ω_x / ω) * v_y - (ω_y / ω) * v_z ) )
    du[7] = - 2srd * ω_x *  v / diameter
    du[8] = - 2srd * ω_y *  v / diameter
    du[9] = - 2srd * ω_z *  v / diameter
end

# Initial conditions and parameters
n_dims = 3
u0 = [0, 0, 0, 72.86, 0, 14.03, 0, 0, -281.28]
trajectory_length = 124
dt = 20e-3
tspan = (0, trajectory_length * dt)
p = [0.2, 0.18, 0.06, 9e4, 2e5, 0.54, 4e-5π, 0.45, 1.225, 1.8e-5, 9.81, 42.67e-3, 0.04593, 1.43e-3]
# Solve ODE to generate data
prob = ODEProblem(projectiletrajectory, u0, tspan, p)
ode_data = Array(solve(prob, saveat=dt))
# Only keep x, y, z data
n_ode_data = ode_data[1:n_dims, :]
# MinMax kinda scaling, maybe?
n_ode_data = (n_ode_data .- minimum(n_ode_data)) / (maximum(n_ode_data) .- minimum(n_ode_data))
# Setup neural ODE
dudt = Chain(Dense(n_dims, 50, tanh), Dense(50, n_dims))
n_ode = NeuralODE(dudt, tspan, Tsit5(), saveat=dt)
# Loss function
v0 = n_ode_data[:, 1]
loss_n_ode() = sum(abs2, n_ode_data .- n_ode(v0))
# Training
Flux.train!(loss_n_ode, Flux.params(n_ode), Iterators.repeated((), 5000), ADAM(0.1))
# Predict
pred = n_ode(v0)
# Plot
t = range(tspan[1], tspan[2], length=trajectory_length+1)
l = @layout [a b c]
p1 = plot(xlabel="t (s)", ylabel="x (m)")
plot!(p1, t, n_ode_data[1, :], label="Ground truth")
plot!(p1, t, pred[1, :], label="Prediction")
p2 = plot(xlabel="t (s)", ylabel="y (m)")
plot!(p2, t, n_ode_data[2, :], label="Ground truth")
plot!(p2, t, pred[2, :], label="Prediction")
p3 = plot(xlabel="t (s)", ylabel="z (m)")
plot!(p3, t, n_ode_data[3, :], label="Ground truth")
plot!(p3, t, pred[3, :], label="Prediction")
plot(p1, p2, p3, layout = l)
plot!(legend=:outerbottom, legendcolumns=2)

This example learns the model for a projectile’s position trajectory in three-dimensional space. The data is generated using a nine-dimensional model that includes velocity and spin. How do I get this example working on my GPU? If I use n_ode = NeuralODE(gpu(dudt), tspan, Tsit5(), saveat=dt) as suggested, then everything breaks.

I also have some questions about neural ODEs in general:

  • Is it possible to train a model on multiple different observed trajectories, where each trajectory has a different length? If so, then how is this done?
  • Is it possible to incorporate initial conditions for multiple time-steps? For example, instead of just u_{0}, could I use u_{0}, u_{1}, u_{2}?

Basically, given some set of observed trajectories, I want to learn the underlying dynamical model. Using this dynamical model, I want to be forecast an entire trajectory using some fixed number of initial points. Is this possible using neural ODEs? I tried using LSTMs, but it didn’t go so well.

2

What do you mean by that? What the error? A priori the code looks fine, though I’d recommend you read the docs rather than a 4 year old blog post

Yes, just add them to the loss function.

Yes that’s multiple shooting, in the docs.

Well machine learning in general is bad for this kind of thing. If you ahve prior knowledge of some model, it’s helpful. That’s the whole point of SciML. See this talk:

1 Like