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.