I tried to follow this tutorial: Experiments with Neural ODEs in Julia.
Data download link (monthly enso): Dropbox - monthly_enso.csv - Simplify your life
using Flux, DiffEqFlux, DifferentialEquations, Plots
using DataFrames, CSV
data = CSV.read(âmonthly_enso.csvâ,DataFrame)
first(data,5)
println(size(data))
tbegin = 0.0
tend = 456
t = range(tbegin,tend,length=tend)
u0 = [2.5; 0.5]
tspan = (tbegin,tend)
trange = range(tbegin,tend,length=tend)
dataset_ts = data
math_law(u) = sin.(2. * u) + cos.(2. * u)
dudt = Chain(u â math_law(u),Dense(2, 50, tanh),Dense(50, 2))
reltol = 1e-7 # tol = tolerances
abstol = 1e-9
n_ode = NeuralODE(dudt, tspan, Tsit5(), saveat=trange, reltol=reltol,abstol=abstol)
ps = Flux.params(n_ode.p)
function loss_n_ode()
pred = n_ode(u0)
loss = sum(abs2, dataset_ts .- pred)
end
n_epochs = 400
learning_rate = 0.01
data = Iterators.repeated((), n_epochs)
opt = ADAM(learning_rate)
cb = function () # callback function to observe training
loss = loss_n_ode()
println("Loss: ", loss)
end
println();
cb() # Display the ODE with the initial parameter values.
Flux.train!(loss_n_ode, ps, data, opt, cb=cb)
pl = plot(
trange,
dataset_ts[1,:],
linewidth=2, ls=:dash,
title=âNeural ODE for forecastingâ,
xaxis=âtâ,
label=âoriginal timeseries x(t)â,
legend=:right)
display(pl)
pl = plot!(
trange,
dataset_ts[2,:],
linewidth=2, ls=:dash,
label=âoriginal timeseries y(t)â)
display(pl)
pred = n_ode(u0)
pl = plot!(
trange,
pred[1,:],
linewidth=1,
label=âpredicted timeseries x(t)â)
display(pl)
pl = plot!(
trange,
pred[2,:],
linewidth=1,
label=âpredicted timeseries y(t)â)
display(pl)