NeuralODE Layer Strategy

New to Julia
,
1

To embed as much known information as possible, I’m trying to include a known equation in a NeuralODE. Using the simple pendulum example for dynamics and following the DiffEqFlux local minima documentation page with a basic neural ODE layer, I have the following code:

# Simple Pendulum Problem
using OrdinaryDiffEq, Plots, DiffEqFlux, Flux

#Constants
const g = 9.81
L = 1.0

#Define the problem
function simplependulum(du,u,p,t)
    θ = u[1]
    dθ = u[2]
    du[1] = dθ
    du[2] = -(g/L)*sin(θ)
end

#Initial Conditions
u0 = [0,π/2]
tspan = (0.0,6.0)
datasize = 30
t = range(tspan[1],tspan[2],length=datasize)

prob = ODEProblem(simplependulum, u0, tspan)
soln = solve(prob,Tsit5(),saveat=t)
ode_data = Array(soln)

diff = ODEProblem(simplependulum,u0,tspan)
dudt = Chain(
	Dense(2,20,tanh),
	Dense(20,2))

n_ode = NeuralODE(dudt,tspan,Tsit5(),saveat=t,reltol=1e-9,abstol=1e-9)

function predict_n_ode(p)
	Array(n_ode(u0,p))
end

function loss_n_ode(p)
pred = predict_n_ode(p)
loss = sum(abs2,ode_data[:,1:size(pred,2)] .- pred)
return loss, pred
end

cb = function (p,l,pred;doplot=false) #callback function to observe training
display(l)
display(sum(abs2,ode_data[:,1:size(pred,2)] .- pred))
# plot current prediction against data
pl = scatter(t[1:size(pred,2)],ode_data[1,1:size(pred,2)],label="data", title="\\theta")
scatter!(pl,t[1:size(pred,2)],pred[1,:],label="prediction")
display(plot(pl))
return false
end

result_node = DiffEqFlux.sciml_train(loss_n_ode, n_ode.p, ADAM(0.05), cb = cb, maxiters = 100)

My question is how to include the simplependulum function in the dudt layers. Like the “sciml_train” documentation page, I’ve tried to put the function directly as the top layer as:

dudt = Chain((x,p) -> simplependulum(x),
             Dense(2,50,tanh),
             Dense(50,2))

which doesn’t work; presumably because simplependulum takes more parameters than simply x. I tried to make an almost-equivalent pendulum function which only accepts x, but that doesn’t work either. I’ve also tried to follow this example and put solve directly in the layer, like:

dudt = Chain(
	(x,p) -> solve(diff,Tsit5(), saveat = 0:0.1:1.0,p = p)[1:2,:],
                vec,
	Dense(2,20,tanh),
	Dense(20,2))

But that also doesn’t work.

So - what’s the best way to approach putting a known function as a Chain layer? Also, what’s going on with the Chain input? I’ve tried to follow and replace relevant parts with the FastChain example, but it’s not immediately intuitive how the tmp1 variable is being used with FastDense (new to Julia…).

2

Chain only has a single input. FastChain carries parameters and x.

3

Thank you - that helps, but fully doesn’t get me there. If I follow the first example and make a function to accept only x and output du, I end up with the following:

# Simple Pendulum Problem
using OrdinaryDiffEq, Plots, DiffEqFlux, Flux

#Constants
const g = 9.81
L = 1.0

#Define the problem
function simplependulum(du,u,p,t)
    θ = u[1]
    dθ = u[2]
    du[1] = dθ
    du[2] = -(g/L)*sin(θ)
end

function pend(u)
    du = Vector{Float64}(undef,2) 
    θ = u[1]
    dθ = u[2]
    du[1] = dθ
    du[2] = -(g/L)*sin(θ)
end

#Initial Conditions
u0 = [0,π/2]
tspan = (0.0,6.0)
datasize = 30
t = range(tspan[1],tspan[2],length=datasize)

prob = ODEProblem(simplependulum, u0, tspan)
soln = solve(prob,Tsit5(),saveat=t)
ode_data = Array(soln)

diff = ODEProblem(simplependulum,u0,tspan)
dudt = FastChain(
	(x,p) -> pend(x),
	FastDense(2,20,tanh),
	FastDense(20,2))

n_ode = NeuralODE(dudt,tspan,Tsit5(),saveat=t,reltol=1e-9,abstol=1e-9)

function predict_n_ode(p)
	Array(n_ode(u0,p))
end

function loss_n_ode(p)
pred = predict_n_ode(p)
loss = sum(abs2,ode_data[:,1:size(pred,2)] .- pred)
return loss, pred
end

cb = function (p,l,pred;doplot=false) #callback function to observe training
display(l)
display(sum(abs2,ode_data[:,1:size(pred,2)] .- pred))
# plot current prediction against data
pl = scatter(t[1:size(pred,2)],ode_data[1,1:size(pred,2)],label="data", title="\\theta")
scatter!(pl,t[1:size(pred,2)],pred[1,:],label="prediction")
display(plot(pl))
return false
end

result_node = DiffEqFlux.sciml_train(loss_n_ode, n_ode.p, ADAM(0.05), cb = cb, maxiters = 100)

which results in the error:

ERROR: MethodError: no method matching Vector{Float64}(::Matrix{Float64})
Closest candidates are:
  Array{T, N}(::AbstractArray{S, N}) where {T, N, S} at array.jl:540
  Vector{T}() where T at boot.jl:467
  Array{T, N}(::StaticArrays.SizedArray{S, T, N, M, TData} where {M, TData<:AbstractArray{T, M}}) where {T, S, N} at C:\Users\lindblot\.julia\packages\StaticArrays\vxjOO\src\SizedArray.jl:98
  ...
Stacktrace:
  [1] convert(#unused#::Type{Vector{Float64}}, a::Matrix{Float64})
    @ Base .\array.jl:532
  [2] setproperty!(x::OrdinaryDiffEq.ODEIntegrator{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, false, Vector{Float64}, Nothing, Float64, Vector{Float32}, Float64, Float64, Float64, Float64, Vector{Vector{Float64}}, ODESolution{Float64, 2, Vector{Vector{Float64}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{Vector{Float64}}}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, OrdinaryDiffEq.InterpolationData{ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Vector{Vector{Float64}}, Vector{Float64}, Vector{Vector{Vector{Float64}}}, OrdinaryDiffEq.Tsit5ConstantCache{Float64, Float64}}, DiffEqBase.DEStats}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, OrdinaryDiffEq.Tsit5ConstantCache{Float64, Float64}, OrdinaryDiffEq.DEOptions{Float64, Float64, Float64, Float64, PIController{Rational{Int64}}, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(LinearAlgebra.opnorm), Bool, CallbackSet{Tuple{}, Tuple{}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, Nothing, Nothing, Int64, Tuple{}, Tuple{}, Tuple{}}, Vector{Float64}, Float64, Nothing, OrdinaryDiffEq.DefaultInit}, f::Symbol, v::Matrix{Float64})
    @ Base .\Base.jl:34
  [3] initialize!(integrator::OrdinaryDiffEq.ODEIntegrator{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, false, Vector{Float64}, Nothing, Float64, Vector{Float32}, Float64, Float64, Float64, Float64, Vector{Vector{Float64}}, ODESolution{Float64, 2, Vector{Vector{Float64}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{Vector{Float64}}}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, OrdinaryDiffEq.InterpolationData{ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Vector{Vector{Float64}}, Vector{Float64}, Vector{Vector{Vector{Float64}}}, OrdinaryDiffEq.Tsit5ConstantCache{Float64, Float64}}, DiffEqBase.DEStats}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, OrdinaryDiffEq.Tsit5ConstantCache{Float64, Float64}, OrdinaryDiffEq.DEOptions{Float64, Float64, Float64, Float64, PIController{Rational{Int64}}, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(LinearAlgebra.opnorm), Bool, CallbackSet{Tuple{}, Tuple{}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, Nothing, Nothing, Int64, Tuple{}, Tuple{}, Tuple{}}, Vector{Float64}, Float64, Nothing, OrdinaryDiffEq.DefaultInit}, cache::OrdinaryDiffEq.Tsit5ConstantCache{Float64, Float64})
    @ OrdinaryDiffEq ~\.julia\packages\OrdinaryDiffEq\Zi9Zh\src\perform_step\low_order_rk_perform_step.jl:565
  [4] __init(prob::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, alg::Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{Val{true}}; saveat::Tuple{}, tstops::Tuple{}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Bool, callback::Nothing, dense::Bool, calck::Bool, dt::Float64, dtmin::Nothing, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Rational{Int64}, abstol::Float64, reltol::Float64, qmin::Rational{Int64}, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, beta1::Nothing, beta2::Nothing, qoldinit::Rational{Int64}, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(LinearAlgebra.opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEq.DefaultInit, kwargs::Base.Iterators.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:save_noise,), Tuple{Bool}}})
    @ OrdinaryDiffEq ~\.julia\packages\OrdinaryDiffEq\Zi9Zh\src\solve.jl:456
  [5] #__solve#471
    @ ~\.julia\packages\OrdinaryDiffEq\Zi9Zh\src\solve.jl:4 [inlined]
  [6] #solve_call#42
    @ ~\.julia\packages\DiffEqBase\WKucm\src\solve.jl:61 [inlined]
  [7] solve_up(prob::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, sensealg::Nothing, u0::Vector{Float64}, p::Vector{Float32}, args::Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}; kwargs::Base.Iterators.Pairs{Symbol, Real, NTuple{5, Symbol}, NamedTuple{(:save_noise, :save_start, :save_end, :reltol, :abstol), Tuple{Bool, Bool, Bool, Float64, Float64}}})
    @ DiffEqBase ~\.julia\packages\DiffEqBase\WKucm\src\solve.jl:87
  [8] #solve#43
    @ ~\.julia\packages\DiffEqBase\WKucm\src\solve.jl:73 [inlined]
  [9] _concrete_solve_adjoint(::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, ::Vector{Float64}, ::Vector{Float32}; save_start::Bool, save_end::Bool, saveat::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, save_idxs::Nothing, kwargs::Base.Iterators.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:reltol, :abstol), Tuple{Float64, Float64}}})
    @ DiffEqSensitivity ~\.julia\packages\DiffEqSensitivity\cLl5o\src\concrete_solve.jl:114
 [10] #_solve_adjoint#61
    @ ~\.julia\packages\DiffEqBase\WKucm\src\solve.jl:347 [inlined]
 [11] #rrule#59
    @ ~\.julia\packages\DiffEqBase\WKucm\src\solve.jl:310 [inlined]
 [12] chain_rrule_kw
    @ ~\.julia\packages\Zygote\TaBlo\src\compiler\chainrules.jl:115 [inlined]
 [13] macro expansion
    @ ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0 [inlined]
 [14] _pullback(::Zygote.Context, ::DiffEqBase.var"#solve_up##kw", ::NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, ::Vector{Float64}, ::Vector{Float32}, ::Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:9
 [15] _apply(::Function, ::Vararg{Any, N} where N)
    @ Core .\boot.jl:804
 [16] adjoint
    @ ~\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [17] _pullback
    @ ~\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57 [inlined]
 [18] _pullback
    @ ~\.julia\packages\DiffEqBase\WKucm\src\solve.jl:73 [inlined]
 [19] _pullback(::Zygote.Context, ::DiffEqBase.var"##solve#43", ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, ::Nothing, ::Nothing, ::Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}, ::typeof(solve), ::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [20] _apply(::Function, ::Vararg{Any, N} where N)
    @ Core .\boot.jl:804
 [21] adjoint
    @ ~\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [22] _pullback
    @ ~\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57 [inlined]
 [23] _pullback
    @ ~\.julia\packages\DiffEqBase\WKucm\src\solve.jl:68 [inlined]
 [24] _pullback(::Zygote.Context, ::CommonSolve.var"#solve##kw", ::NamedTuple{(:sensealg, :saveat, :reltol, :abstol), Tuple{InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}, ::typeof(solve), ::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [25] _apply(::Function, ::Vararg{Any, N} where N)
    @ Core .\boot.jl:804
 [26] adjoint
    @ ~\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [27] adjoint(::Zygote.Context, ::typeof(Core._apply_iterate), ::typeof(iterate), ::Function, ::Tuple{NamedTuple{(:sensealg, :saveat, :reltol, :abstol), Tuple{InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}, typeof(solve), ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, DiffEqFlux.var"#dudt_#101"{NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ::Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}})
    @ Zygote .\none:0
 [28] _pullback
    @ ~\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57 [inlined]
 [29] _pullback
    @ ~\.julia\packages\DiffEqFlux\N7blG\src\neural_de.jl:77 [inlined]
 [30] _pullback(::Zygote.Context, ::NeuralODE{FastChain{Tuple{var"#1#2", FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#82"{Vector{Float32}}}}}, Vector{Float32}, Nothing, Tuple{Float64, Float64}, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}}, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :reltol, :abstol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}}}, ::Vector{Float64}, ::Vector{Float32})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [31] _pullback
    @ .\REPL[19]:2 [inlined]
 [32] _pullback(ctx::Zygote.Context, f::typeof(predict_n_ode), args::Vector{Float32})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [33] _pullback
    @ .\REPL[20]:2 [inlined]
 [34] _pullback(ctx::Zygote.Context, f::typeof(loss_n_ode), args::Vector{Float32})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [35] _pullback
    @ ~\.julia\packages\DiffEqFlux\N7blG\src\train.jl:84 [inlined]
 [36] _pullback(::Zygote.Context, ::DiffEqFlux.var"#74#79"{typeof(loss_n_ode)}, ::Vector{Float32}, ::Nothing)
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [37] _apply
    @ .\boot.jl:804 [inlined]
 [38] adjoint
    @ ~\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [39] _pullback
    @ ~\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57 [inlined]
 [40] _pullback
    @ ~\.julia\packages\SciMLBase\UIp7W\src\problems\basic_problems.jl:107 [inlined]
 [41] _pullback(::Zygote.Context, ::OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#74#79"{typeof(loss_n_ode)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ::Vector{Float32}, ::Nothing)
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [42] _apply(::Function, ::Vararg{Any, N} where N)
    @ Core .\boot.jl:804
 [43] adjoint
    @ ~\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [44] _pullback(::Zygote.Context, ::typeof(Core._apply_iterate), ::typeof(iterate), ::Function, ::Tuple{Vector{Float32}, Nothing}, ::Tuple{})
    @ Zygote ~\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57
 [45] _pullback
    @ ~\.julia\packages\GalacticOptim\bEh06\src\function\zygote.jl:6 [inlined]
 [46] _pullback(ctx::Zygote.Context, f::GalacticOptim.var"#220#230"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#74#79"{typeof(loss_n_ode)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}, args::Vector{Float32})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [47] _apply(::Function, ::Vararg{Any, N} where N)
    @ Core .\boot.jl:804
 [48] adjoint
    @ ~\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [49] _pullback
    @ ~\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57 [inlined]
 [50] _pullback
    @ ~\.julia\packages\GalacticOptim\bEh06\src\function\zygote.jl:8 [inlined]
 [51] _pullback(ctx::Zygote.Context, f::GalacticOptim.var"#223#233"{Tuple{}, GalacticOptim.var"#220#230"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#74#79"{typeof(loss_n_ode)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, args::Vector{Float32})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [52] _pullback(f::Function, args::Vector{Float32})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface.jl:34
 [53] pullback(f::Function, args::Vector{Float32})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface.jl:40
 [54] gradient(f::Function, args::Vector{Float32})
    @ Zygote ~\.julia\packages\Zygote\TaBlo\src\compiler\interface.jl:75
 [55] (::GalacticOptim.var"#221#231"{GalacticOptim.var"#220#230"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#74#79"{typeof(loss_n_ode)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}})(::Vector{Float32}, ::Vector{Float32})
    @ GalacticOptim ~\.julia\packages\GalacticOptim\bEh06\src\function\zygote.jl:8
 [56] macro expansion
    @ ~\.julia\packages\GalacticOptim\bEh06\src\solve\flux.jl:43 [inlined]
 [57] macro expansion
    @ ~\.julia\packages\GalacticOptim\bEh06\src\solve\solve.jl:35 [inlined]
 [58] __solve(prob::OptimizationProblem{false, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#74#79"{typeof(loss_n_ode)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, GalacticOptim.var"#221#231"{GalacticOptim.var"#220#230"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#74#79"{typeof(loss_n_ode)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#224#234"{GalacticOptim.var"#220#230"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#74#79"{typeof(loss_n_ode)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#229#239", Nothing, Nothing, Nothing}, Vector{Float32}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Nothing, Base.Iterators.Pairs{Symbol, var"#3#5", Tuple{Symbol}, NamedTuple{(:cb,), Tuple{var"#3#5"}}}}, opt::ADAM, data::Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}; maxiters::Int64, cb::Function, progress::Bool, save_best::Bool, kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ GalacticOptim ~\.julia\packages\GalacticOptim\bEh06\src\solve\flux.jl:41
 [59] #solve#474
    @ ~\.julia\packages\SciMLBase\UIp7W\src\solve.jl:3 [inlined]
 [60] sciml_train(::typeof(loss_n_ode), ::Vector{Float32}, ::ADAM, ::Nothing; lower_bounds::Nothing, upper_bounds::Nothing, maxiters::Int64, kwargs::Base.Iterators.Pairs{Symbol, var"#3#5", Tuple{Symbol}, NamedTuple{(:cb,), Tuple{var"#3#5"}}})
    @ DiffEqFlux ~\.julia\packages\DiffEqFlux\N7blG\src\train.jl:89
 [61] top-level scope
    @ REPL[22]:1

What am I doing wrong there?

4

You weren’t returning the array you were creating in pend.

function pend(u)
    θ = u[1]
    dθ = u[2]
    [dθ,-(g/L)*sin(θ)]
end

is fine. Then you probably want to do multiple shooting here.

5

Thanks, Chris; that + multiple shooting does the trick.