What's the NeuralNetDiffEq.jl equivalent of this piece of code for solving ODE in Julia?

Viewed 71

I am trying to make documentation for NeuraNetDiffEq.jl. I find an ODE solution with Flux.jl from this amazing tutorial here https://mitmath.github.io/18S096SciML/lecture2/ml .

using Flux
using DifferentialEquations
using LinearAlgebra
using Plots
using Statistics

NNODE = Chain(x -> [x],
           Dense(1, 32, tanh),
           Dense(32, 1),
           first)
NNODE(1.0)
g(t) = t * NNODE(t) + 1f0

ϵ = sqrt(eps(Float32))
loss() = mean(abs2(((g(t + ϵ) - g(t)) / ϵ) - cos(2π * t)) for t in 0:1f-2:1f0)
opt = Flux.Descent(0.01)
data = Iterators.repeated((), 5000)
iter = 0
cb = function () # callback function to observe training
    global iter += 1
    if iter % 500 == 0
        display(loss())
    end
end
display(loss())
Flux.train!(loss, Flux.params(NNODE), data, opt; cb=cb)

t = 0:0.001:1.0
plot(t,g.(t),label="NN")
plot!(t,1.0 .+ sin.(2π .* t) / 2π, label="True")

I am having trouble understanding the parameters involved to invoke training process for NueralNetDiffEq.jl as in:

function DiffEqBase.solve(
    prob::DiffEqBase.AbstractODEProblem,
    alg::NeuralNetDiffEqAlgorithm,
    args...;
    dt,
    timeseries_errors = true,
    save_everystep=true,
    adaptive=false,
    abstol = 1f-6,
    verbose = false,
    maxiters = 100)

What would be a valid input for alg parameter? What would be the equivalent code in NeuralNetDiffEq.jl for the above ODE example?

0 Answers
Related