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)