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.9650168 -1.0541967 1.0143851; -1.8106718 1.3827721 -0.2504116; … ; -0.11411189 1.083959 0.89708346; 2.4581158 1.0878487 -1.2851561], bias = Float32[0.039714336, 0.080408834, -0.47516906, -0.7222136, 0.08322632, -0.21592027, 0.049395334, 1.6934074, 1.0127257, -0.06481843 … -1.0453333, -3.2722178, 2.2262084, 0.7445199, -0.32184333, -0.5640919, -0.8249773, -0.24367915, -0.7101517, 2.546886]), layer_2 = (weight = Float32[14.760053 3.401427 … -14.018913 8.679321; 8.839743 0.32747328 … -7.2446938 4.5207105], bias = Float32[-19.96841, -10.200766]))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#f##3"{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"#2#3", 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, Nothing, @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{}, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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.9363545427459513, 0.11512007139314075, -0.05147461413909219], [0.8699971162259429, 0.2223118121883675, -0.09230892841431051], [0.8007616313355653, 0.31983038448708406, -0.12059201582265039], [0.728804058154263, 0.4065685580387039, -0.135372616192967], [0.6547602214055406, 0.482095548353587, -0.1368557697591278], [0.5796717337343648, 0.5464653090474518, -0.12613704278181548], [0.5047761164328203, 0.5999523094408157, -0.10472842587363579], [0.43131118298681415, 0.6428374598817495, -0.074148642868564], [0.36040573037314527, 0.6753223491242395, -0.03572807949738469], [0.2930220502817123, 0.6975670907356323, 0.009410858982655498]], 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#f##3"{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"#2#3", 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, Nothing, @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{}, Nothing, @NamedTuple{}}, SciMLBase.StandardODEProblem}(SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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"#2#3", 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, Nothing, @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"#2#3"(), (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[1.3640302 -0.36107618 0.37594157; 0.33330995 -1.4359815 1.4545712; … ; -0.19775648 0.24590233 -0.9404391; 0.71272886 0.8129535 -1.6413245], bias = Float32[-0.031678043, -0.0138918795, 0.13036866, 0.221475, -0.2524069, 0.31746367, 0.0761661, -0.1753235, 0.16692692, 0.55990374 … 0.38026392, -0.16160919, 0.4957599, -0.059035912, 0.50306064, 0.4307109, 0.19431432, 0.40638694, 0.3944812, 0.54980373]), layer_2 = (weight = Float32[-0.070465125 -0.05771675 … 0.018785004 -0.1550005; 0.117481075 -0.08074972 … -0.13462856 0.035107773], bias = Float32[-0.12248206, 0.025191948])), Base.Pairs{Symbol, Union{}, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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"#2#3", 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, Nothing, @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"#2#3"(), (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.9363545427459513, 0.11512007139314075, -0.05147461413909219], [0.8699971162259429, 0.2223118121883675, -0.09230892841431051], [0.8007616313355653, 0.31983038448708406, -0.12059201582265039], [0.728804058154263, 0.4065685580387039, -0.135372616192967], [0.6547602214055406, 0.482095548353587, -0.1368557697591278], [0.5796717337343648, 0.5464653090474518, -0.12613704278181548], [0.5047761164328203, 0.5999523094408157, -0.10472842587363579], [0.43131118298681415, 0.6428374598817495, -0.074148642868564], [0.36040573037314527, 0.6753223491242395, -0.03572807949738469], [0.2930220502817123, 0.6975670907356323, 0.009410858982655498]], [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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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"#2#3", 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, Nothing, @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"#2#3"(), (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.3623927287435007, 0.6745348284250752, -0.03692755716857593], (layer_1 = (weight = Float32[1.3640302 -0.36107618 0.37594157; 0.33330995 -1.4359815 1.4545712; … ; -0.19775648 0.24590233 -0.9404391; 0.71272886 0.8129535 -1.6413245], bias = Float32[-0.031678043, -0.0138918795, 0.13036866, 0.221475, -0.2524069, 0.31746367, 0.0761661, -0.1753235, 0.16692692, 0.55990374 … 0.38026392, -0.16160919, 0.4957599, -0.059035912, 0.50306064, 0.4307109, 0.19431432, 0.40638694, 0.3944812, 0.54980373]), layer_2 = (weight = Float32[-0.070465125 -0.05771675 … 0.018785004 -0.1550005; 0.117481075 -0.08074972 … -0.13462856 0.035107773], bias = Float32[-0.12248206, 0.025191948]))), SciMLBase.UDerivativeWrapper{false, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#f#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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#f##3"{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"#2#3", 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, Nothing, @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"#2#3", 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, Nothing, @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"#2#3"(), (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.897134545733878, (layer_1 = (weight = Float32[1.3640302 -0.36107618 0.37594157; 0.33330995 -1.4359815 1.4545712; … ; -0.19775648 0.24590233 -0.9404391; 0.71272886 0.8129535 -1.6413245], bias = Float32[-0.031678043, -0.0138918795, 0.13036866, 0.221475, -0.2524069, 0.31746367, 0.0761661, -0.1753235, 0.16692692, 0.55990374 … 0.38026392, -0.16160919, 0.4957599, -0.059035912, 0.50306064, 0.4307109, 0.19431432, 0.40638694, 0.3944812, 0.54980373]), layer_2 = (weight = Float32[-0.070465125 -0.05771675 … 0.018785004 -0.1550005; 0.117481075 -0.08074972 … -0.13462856 0.035107773], bias = Float32[-0.12248206, 0.025191948]))), 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(146, 0, 12, 84, 12, 0, 0, 0, 0, 0, 12, 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)))4.431774167574023Notice 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#5 (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[1.140235 -0.52473754 0.36952356; 0.6045647 -1.4949999 1.1031356; … ; -0.56821996 0.64372957 -0.91814685; 0.7688229 0.84544134 -1.6209465], bias = Float32[-0.4210742, -0.19214827, 0.22215322, 0.3962205, -0.58932805, 0.56127733, 0.040960822, 0.44354764, 0.1491895, 0.5518863 … 0.33237073, -0.14448142, 0.46688825, 0.065156505, 0.31207323, 1.0322604, -0.21757312, 0.3945852, 0.38936234, 0.69604516]), layer_2 = (weight = Float32[-0.04232437 -0.18652706 … 0.061254676 -0.1194165; 1.001475 4.0466375 … -1.8593764 -1.3029335], bias = Float32[0.046252392, -0.36969516]))