Enforcing Physical Constraints via Universal Differential-Algebraic Equations
As shown in DiffEqDocs, differential-algebraic equations (DAEs) can be used to impose physical constraints. One way to define a DAE is through an ODE with a singular mass matrix. For example, if we make Mu' = f(u) where the last row of M is all zeros, then we have a constraint defined by the right-hand side. Using NeuralODEMM, we can use this to define a neural ODE where the sum of all 3 terms must add to one. An example of this is as follows:
using DiffEqFlux
using Lux, ComponentArrays, Optimization, OptimizationOptimJL, OrdinaryDiffEq, Plots
using Random
rng = Random.default_rng()
function f!(du, u, p, t)
y₁, y₂, y₃ = u
k₁, k₂, k₃ = p
du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
du[2] = k₁ * y₁ - k₃ * y₂ * y₃ - k₂ * y₂^2
du[3] = y₁ + y₂ + y₃ - 1
return nothing
end
u₀ = [1.0, 0, 0]
M = [1.0 0 0
0 1.0 0
0 0 0]
tspan = (0.0, 1.0)
p = [0.04, 3e7, 1e4]
stiff_func = ODEFunction(f!; mass_matrix = M)
prob_stiff = ODEProblem(stiff_func, u₀, tspan, p)
sol_stiff = solve(prob_stiff, Rodas5(); saveat = 0.1)
nn_dudt2 = Lux.Chain(Lux.Dense(3, 64, tanh), Lux.Dense(64, 2))
pinit, st = Lux.setup(rng, nn_dudt2)
model_stiff_ndae = NeuralODEMM(nn_dudt2, (u, p, t) -> [u[1] + u[2] + u[3] - 1],
tspan, M, Rodas5(; autodiff = false); saveat = 0.1)
function predict_stiff_ndae(p)
return model_stiff_ndae(u₀, p, st)[1]
end
function loss_stiff_ndae(p)
pred = predict_stiff_ndae(p)
loss = sum(abs2, Array(sol_stiff) .- pred)
return loss
end
# callback = function (state, l, pred) #callback function to observe training
# display(l)
# return false
# end
l1 = first(loss_stiff_ndae(ComponentArray(pinit)))
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_stiff_ndae(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))
result_stiff = Optimization.solve(optprob, OptimizationOptimJL.BFGS(); maxiters = 100)retcode: Failure
u: ComponentVector{Float32}(layer_1 = (weight = Float32[0.8172428 1.470233 0.91816556; 0.19951487 0.52637887 1.434899; … ; 1.5944408 -0.68565613 -0.35477436; -0.46823388 -0.4332052 0.78421825], bias = Float32[-0.13956481, 0.41927928, -0.120698996, 0.38777503, 0.3389905, 0.55854356, -0.48526105, 0.3892994, 0.25512162, -0.27512997 … -0.24473669, 0.4502602, -0.19704269, 0.43308616, 0.087899916, -0.35008088, 0.56332165, 0.30172768, 0.24721399, -0.45868942]), layer_2 = (weight = Float32[-0.30201703 0.097124614 … -0.038215827 -0.11498351; -0.06786923 -0.082503214 … -0.252229 0.0034923286], bias = Float32[-0.022657923, -0.12753373]))Step-by-Step Description
Load Packages
using DiffEqFlux
using Lux, ComponentArrays, Optimization, OptimizationOptimJL, OrdinaryDiffEq, Plots
using Random
rng = Random.default_rng()Random.TaskLocalRNG()Differential Equation
First, we define our differential equations as a highly stiff problem, which makes the fitting difficult.
function f!(du, u, p, t)
y₁, y₂, y₃ = u
k₁, k₂, k₃ = p
du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
du[2] = k₁ * y₁ - k₃ * y₂ * y₃ - k₂ * y₂^2
du[3] = y₁ + y₂ + y₃ - 1
return nothing
endf! (generic function with 1 method)Parameters
u₀ = [1.0, 0, 0]
M = [1.0 0 0
0 1.0 0
0 0 0]
tspan = (0.0, 1.0)
p = [0.04, 3e7, 1e4]3-element Vector{Float64}:
0.04
3.0e7
10000.0u₀= Initial ConditionsM= Semi-explicit Mass Matrix (last row is the constraint equation and are therefore all zeros)tspan= Time span over which to evaluatep= parametersk1,k2andk3of the differential equation above
ODE Function, Problem and Solution
We define and solve our ODE problem to generate the “labeled” data which will be used to train our Neural Network.
stiff_func = ODEFunction(f!; mass_matrix = M)
prob_stiff = ODEProblem(stiff_func, u₀, tspan, p)
sol_stiff = solve(prob_stiff, Rodas5(); saveat = 0.1)retcode: Success
Interpolation: 1st order linear
t: 11-element Vector{Float64}:
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
u: 11-element Vector{Vector{Float64}}:
[1.0, 0.0, 0.0]
[0.9960777474341874, 3.580437232874095e-5, 0.003886448193483852]
[0.9923059457218126, 3.5123030150794895e-5, 0.007658931248036548]
[0.9886739385487181, 3.44771604650164e-5, 0.011291584290816925]
[0.9851721109941468, 3.3863965535471e-5, 0.01479402504031772]
[0.981791774709964, 3.328089042278384e-5, 0.01817494439961413]
[0.9785250342445555, 3.2725768110357636e-5, 0.02144223998733427]
[0.9753647131268983, 3.21965297854234e-5, 0.024603090343316574]
[0.972304297901924, 3.1691238996228466e-5, 0.027664010859079317]
[0.9693377993879673, 3.12082968346845e-5, 0.030630992315197364]
[0.9664597388050115, 3.07462661101518e-5, 0.03350951492887844]Because this is a DAE, we need to make sure to use a compatible solver. Rodas5 works well for this example.
Neural Network Layers
Next, we create our layers using Lux.Chain. We use this instead of Flux.Chain because it is more suited to SciML applications (similarly for Lux.Dense). The input to our network will be the initial conditions fed in as u₀.
nn_dudt2 = Lux.Chain(Lux.Dense(3, 64, tanh), Lux.Dense(64, 2))
pinit, st = Lux.setup(rng, nn_dudt2)
model_stiff_ndae = NeuralODEMM(nn_dudt2, (u, p, t) -> [u[1] + u[2] + u[3] - 1],
tspan, M, Rodas5(; autodiff = false); saveat = 0.1)
model_stiff_ndae(u₀, ComponentArray(pinit), st)(SciMLBase.ODESolution{Float64, 2, Vector{Vector{Float64}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{Vector{Float64}}}, Nothing, SciMLBase.ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}, OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}, OrdinaryDiffEqCore.InterpolationData{SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{Vector{Float64}}, Vector{Float64}, Vector{Vector{Vector{Float64}}}, Nothing, OrdinaryDiffEqRosenbrock.RosenbrockCombinedConstantCache{SciMLBase.TimeDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}}, SciMLBase.UDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}}, OrdinaryDiffEqRosenbrock.RodasTableau{Float64, Float64}, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Nothing, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}}, BitVector}, SciMLBase.DEStats, Nothing, Nothing, Nothing, Nothing}([[1.0, 0.0, 0.0], [0.9913291259676597, -0.07231725131240258, 0.08098812534474335], [0.9679688544853329, -0.1371617683433477, 0.16919291385801488], [0.9297986765680304, -0.19244041358178746, 0.2626417370137574], [0.8769913418766243, -0.23725997247994177, 0.3602686306033172], [0.8097214162975963, -0.2719575197612888, 0.46223610346369226], [0.7281505991259758, -0.29806804744315535, 0.5699174483171795], [0.6329801557473118, -0.3184905710539163, 0.6855104153066047], [0.5263857730352284, -0.3376948322417077, 0.8113090592064797], [0.41242782612969275, -0.3611695108507423, 0.9487416847210491], [0.295735951893218, -0.39320246488264643, 1.0974665129894288]], nothing, nothing, [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], [[[1.0, 0.0, 0.0]]], nothing, SciMLBase.ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}(SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}(NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}(Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(3 => 64, tanh), layer_2 = Dense(64 => 2)), nothing), Main.var"#1#2"(), (0.0, 1.0), [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0], (OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}(nothing, OrdinaryDiffEqCore.DEFAULT_PRECS, OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, AutoFiniteDiff()),), Base.Pairs(:saveat => 0.1)), StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}(Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(3 => 64, tanh), layer_2 = Dense(64 => 2)), nothing), nothing, (layer_1 = NamedTuple(), layer_2 = NamedTuple()), nothing, Val{true}())), [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0], nothing, DiffEqFlux.basic_tgrad, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), [1.0, 0.0, 0.0], (0.0, 1.0), (layer_1 = (weight = Float32[0.65176344 -0.8900583 -0.00614206; -0.33992806 -0.9590705 -1.4015273; … ; -1.576024 -1.6410742 1.3571451; 0.4546056 1.5752465 1.3588256], bias = Float32[0.27661148, -0.423115, 0.34880325, 0.40609792, 0.271, -0.029302876, -0.049681492, 0.34828326, 0.28735182, -0.09361144 … -0.12409447, -0.2889236, -0.26697645, 0.08191165, 0.34306607, -0.36691323, 0.23065206, 0.1589453, -0.32673448, -0.24885118]), layer_2 = (weight = Float32[-0.07352558 0.081135966 … -0.16866808 -0.036217097; -0.095352724 -0.17643526 … 0.07492595 0.03190571], bias = Float32[0.10281293, -0.0933979])), Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}(), SciMLBase.StandardODEProblem()), OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}(nothing, OrdinaryDiffEqCore.DEFAULT_PRECS, OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, AutoFiniteDiff()), OrdinaryDiffEqCore.InterpolationData{SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{Vector{Float64}}, Vector{Float64}, Vector{Vector{Vector{Float64}}}, Nothing, OrdinaryDiffEqRosenbrock.RosenbrockCombinedConstantCache{SciMLBase.TimeDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}}, SciMLBase.UDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}}, OrdinaryDiffEqRosenbrock.RodasTableau{Float64, Float64}, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Nothing, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}}, BitVector}(SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}(NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}(Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(3 => 64, tanh), layer_2 = Dense(64 => 2)), nothing), Main.var"#1#2"(), (0.0, 1.0), [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0], (OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}(nothing, OrdinaryDiffEqCore.DEFAULT_PRECS, OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, AutoFiniteDiff()),), Base.Pairs(:saveat => 0.1)), StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}(Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(3 => 64, tanh), layer_2 = Dense(64 => 2)), nothing), nothing, (layer_1 = NamedTuple(), layer_2 = NamedTuple()), nothing, Val{true}())), [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0], nothing, DiffEqFlux.basic_tgrad, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), [[1.0, 0.0, 0.0], [0.9913291259676597, -0.07231725131240258, 0.08098812534474335], [0.9679688544853329, -0.1371617683433477, 0.16919291385801488], [0.9297986765680304, -0.19244041358178746, 0.2626417370137574], [0.8769913418766243, -0.23725997247994177, 0.3602686306033172], [0.8097214162975963, -0.2719575197612888, 0.46223610346369226], [0.7281505991259758, -0.29806804744315535, 0.5699174483171795], [0.6329801557473118, -0.3184905710539163, 0.6855104153066047], [0.5263857730352284, -0.3376948322417077, 0.8113090592064797], [0.41242782612969275, -0.3611695108507423, 0.9487416847210491], [0.295735951893218, -0.39320246488264643, 1.0974665129894288]], [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], [[[1.0, 0.0, 0.0]]], nothing, false, OrdinaryDiffEqRosenbrock.RosenbrockCombinedConstantCache{SciMLBase.TimeDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}}, SciMLBase.UDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}}, OrdinaryDiffEqRosenbrock.RodasTableau{Float64, Float64}, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Nothing, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}}(SciMLBase.TimeDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}}(SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}(NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}(Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(3 => 64, tanh), layer_2 = Dense(64 => 2)), nothing), Main.var"#1#2"(), (0.0, 1.0), [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0], (OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}(nothing, OrdinaryDiffEqCore.DEFAULT_PRECS, OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, AutoFiniteDiff()),), Base.Pairs(:saveat => 0.1)), StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}(Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(3 => 64, tanh), layer_2 = Dense(64 => 2)), nothing), nothing, (layer_1 = NamedTuple(), layer_2 = NamedTuple()), nothing, Val{true}())), [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0], nothing, DiffEqFlux.basic_tgrad, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), [0.36510661341025774, -0.37301928820251823, 1.0079126747922604], (layer_1 = (weight = Float32[0.65176344 -0.8900583 -0.00614206; -0.33992806 -0.9590705 -1.4015273; … ; -1.576024 -1.6410742 1.3571451; 0.4546056 1.5752465 1.3588256], bias = Float32[0.27661148, -0.423115, 0.34880325, 0.40609792, 0.271, -0.029302876, -0.049681492, 0.34828326, 0.28735182, -0.09361144 … -0.12409447, -0.2889236, -0.26697645, 0.08191165, 0.34306607, -0.36691323, 0.23065206, 0.1589453, -0.32673448, -0.24885118]), layer_2 = (weight = Float32[-0.07352558 0.081135966 … -0.16866808 -0.036217097; -0.095352724 -0.17643526 … 0.07492595 0.03190571], bias = Float32[0.10281293, -0.0933979]))), SciMLBase.UDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:256, Axis(weight = ViewAxis(1:192, ShapedAxis((64, 3))), bias = ViewAxis(193:256, Shaped1DAxis((64,))))), layer_2 = ViewAxis(257:386, Axis(weight = ViewAxis(1:128, ShapedAxis((2, 64))), bias = ViewAxis(129:130, Shaped1DAxis((2,))))))}}}}(SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Matrix{Float64}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(DiffEqFlux.var"#f#37"{NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}, StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}(NeuralODEMM{Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Main.var"#1#2", Tuple{Float64, Float64}, Matrix{Float64}, Tuple{OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}}, Base.Pairs{Symbol, Float64, Tuple{Symbol}, @NamedTuple{saveat::Float64}}}(Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(3 => 64, tanh), layer_2 = Dense(64 => 2)), nothing), Main.var"#1#2"(), (0.0, 1.0), [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0], (OrdinaryDiffEqRosenbrock.Rodas5{0, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Int64}, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:forward}(), true, nothing, typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!)}(nothing, OrdinaryDiffEqCore.DEFAULT_PRECS, OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, AutoFiniteDiff()),), Base.Pairs(:saveat => 0.1)), StatefulLuxLayer{Val{true}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}(Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(3 => 64, tanh), layer_2 = Dense(64 => 2)), nothing), nothing, (layer_1 = NamedTuple(), layer_2 = NamedTuple()), nothing, Val{true}())), [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0], nothing, DiffEqFlux.basic_tgrad, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), 0.9406376780513307, (layer_1 = (weight = Float32[0.65176344 -0.8900583 -0.00614206; -0.33992806 -0.9590705 -1.4015273; … ; -1.576024 -1.6410742 1.3571451; 0.4546056 1.5752465 1.3588256], bias = Float32[0.27661148, -0.423115, 0.34880325, 0.40609792, 0.271, -0.029302876, -0.049681492, 0.34828326, 0.28735182, -0.09361144 … -0.12409447, -0.2889236, -0.26697645, 0.08191165, 0.34306607, -0.36691323, 0.23065206, 0.1589453, -0.32673448, -0.24885118]), layer_2 = (weight = Float32[-0.07352558 0.081135966 … -0.16866808 -0.036217097; -0.095352724 -0.17643526 … 0.07492595 0.03190571], bias = Float32[0.10281293, -0.0933979]))), OrdinaryDiffEqRosenbrock.RodasTableau{Float64, Float64}([0.0 0.0 … 0.0 0.0; 2.0 0.0 … 0.0 0.0; … ; -14.09640773051259 6.925207756232704 … 0.0 0.0; -14.09640773051259 6.925207756232704 … 1.0 0.0], [0.0 0.0 … 0.0 0.0; -10.31323885133993 0.0 … 0.0 0.0; … ; 34.20013733472935 -14.1553540271769 … -6.551835421242162 0.0; 42.57076742291101 -13.80770672017997 … -6.685968952921985 -5.810979938412932], 0.19, [0.0, 0.38, 0.3878509998321533, 0.483971893787384, 0.457047700881958, 1.0, 1.0, 1.0], [0.19, -0.18230792253337147, -0.3192318321868749, 0.3449828624725343, -0.37741756439208984, 0.0, 0.0, 0.0], [27.354592673333357 -6.925207756232857 … -1.5306074446748028 -1.3929872940716344; 44.19024239501722 1.3677947663381929e-13 … 2.5793540257308067 2.2435122582734066; -44.0988150021747 -5.755396159656812e-13 … -5.792426076169686 -5.32503859794143]), [0.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0], LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}(Matrix{Float64}(undef, 0, 0), Int64[], 0), nothing, AutoFiniteDiff(), 3), Bool[1, 1, 0], false), false, 0, SciMLBase.DEStats(158, 0, 13, 91, 13, 0, 0, 0, 0, 0, 13, 0, 0.0), nothing, SciMLBase.ReturnCode.Success, nothing, nothing, nothing), (layer_1 = NamedTuple(), layer_2 = NamedTuple()))Because this is a stiff problem, we have manually imposed that sum constraint via (u,p,t) -> [u[1] + u[2] + u[3] - 1], making the fitting easier.
Prediction Function
For simplicity, we define a wrapper function that only takes in the model's parameters to make predictions.
function predict_stiff_ndae(p)
return model_stiff_ndae(u₀, p, st)[1]
endpredict_stiff_ndae (generic function with 1 method)Train Parameters
Training our network requires a loss function, an optimizer, and a callback function to display the progress.
Loss
We first make our predictions based on the current parameters, then calculate the loss from these predictions. In this case, we use least squares as our loss.
function loss_stiff_ndae(p)
pred = predict_stiff_ndae(p)
loss = sum(abs2, sol_stiff .- pred)
return loss
end
l1 = first(loss_stiff_ndae(ComponentArray(pinit)))5.703452316853184Notice that we are feeding the parameters of model_stiff_ndae to the loss_stiff_ndae function. model_stiff_node.p are the weights of our NN and is of size 386 (4 * 64 + 65 * 2) including the biases.
Optimizer
The optimizer is BFGS(see below).
Callback
The callback function displays the loss during training.
callback = function (state, l) #callback function to observe training
display(l)
return false
end#3 (generic function with 1 method)Train
Finally, training with Optimization.solve by passing: loss function, model parameters, optimizer, callback and maximum iteration.
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_stiff_ndae(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))
result_stiff = Optimization.solve(optprob, OptimizationOptimJL.BFGS(); maxiters = 100)retcode: Failure
u: ComponentVector{Float32}(layer_1 = (weight = Float32[0.5961264 -0.8692461 0.004942318; -0.6341386 -0.89056134 -1.3389535; … ; -1.5910673 -1.6358973 1.340079; 0.8897456 1.5287077 1.2640114], bias = Float32[0.22729672, -0.5993444, 0.32258293, 0.52820474, -0.10291811, -0.022541817, -0.019249178, 0.2504297, 0.49585792, 0.06621203 … -0.2555972, -0.31265694, -0.2217152, 0.08399092, 0.3394847, -0.34998178, 0.22907646, 0.12043193, -0.3546264, 0.08848713]), layer_2 = (weight = Float32[-0.09618316 -0.03032317 … 0.034115046 0.38566038; 0.16177022 -0.03577527 … 0.42494175 -0.20766708], bias = Float32[0.12448733, -0.21237192]))