# DiffEqFlux.jl: Questions about Neural ODEs

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.

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

https://docs.sciml.ai/DiffEqFlux/dev/

Yes, just add them to the loss function.

Yes that’s multiple shooting, in the docs.

https://docs.sciml.ai/DiffEqFlux/dev/examples/multiple_shooting/

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