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)