chrisrackauckas / universal_differential_equations

Repository for the Universal Differential Equations for Scientific Machine Learning paper, describing a computational basis for high performance SciML

Universal Differential Equations for Scientific Machine Learning (SciML)

Repository for the universal differential equations paper: arXiv:2001.04385 [cs.LG].

This repository is simply a demonstration. For more up-to-date versions of these codes, see the SciMLSensitivity.jl and DiffEqFlux.jl documentations. Their tutorials incorporate all of the ideas from this repo in well-maintained and tested software!

For more software, see the SciML organization and its Github organization

universal_differential_equations's Issues

LotkaVolterra manifest file outdated

Looks like the LotkaVolterra example folder has an outdated and invalid Manifest file with an old reference to a master of DataDrivenDiffEq

[2445eb08] DataDrivenDiffEq v0.5.5 `https://github.com/SciML/DataDrivenDiffEq.jl.git#master`

It also seems to have a lot of packages that aren't part of the experiment and there is no Project.toml

Instability detected in Fisher-KPP example

Hi,
when running the Fisher-KPP example (Fisher-KPP-CNN.jl), everything works fine until the network is trained with BFGS.
The code aborts due to a detected instability.

I guess it has something to do with ReverseDiffVJP?

Warning: Instability detected. Aborting
└ @ SciMLBase ~/.julia/packages/SciMLBase/9EjAY/src/integrator_interface.jl:351
┌ Warning: Instability detected. Aborting
└ @ SciMLBase ~/.julia/packages/SciMLBase/9EjAY/src/integrator_interface.jl:351
┌ Warning: Instability detected. Aborting
└ @ SciMLBase ~/.julia/packages/SciMLBase/9EjAY/src/integrator_interface.jl:351
┌ Warning: Instability detected. Aborting
└ @ SciMLBase ~/.julia/packages/SciMLBase/9EjAY/src/integrator_interface.jl:351
┌ Warning: Instability detected. Aborting
└ @ SciMLBase ~/.julia/packages/SciMLBase/9EjAY/src/integrator_interface.jl:351
┌ Warning: Instability detected. Aborting
└ @ SciMLBase ~/.julia/packages/SciMLBase/9EjAY/src/integrator_interface.jl:351
ERROR: LoadError: DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths 11 and 2")
Stacktrace:
  [1] _bcs1
    @ ./broadcast.jl:501 [inlined]
  [2] _bcs (repeats 2 times)
    @ ./broadcast.jl:495 [inlined]
  [3] broadcast_shape
    @ ./broadcast.jl:489 [inlined]
  [4] combine_axes
    @ ./broadcast.jl:484 [inlined]
  [5] instantiate
    @ ./broadcast.jl:266 [inlined]
  [6] materialize(bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{2}, Nothing, typeof(-), Tuple{Matrix{Float64}, Matrix{Float64}}})
    @ Base.Broadcast ./broadcast.jl:883
  [7] loss_rd(p::Vector{Float64})
    @ Main ~/Documents/Workshops/FUHRI2021/P4-Machine-Learning-And-PDEs/universal_differential_equations/FisherKPP/Fisher-KPP-CNN.jl:142
  [8] (::DiffEqFlux.var"#69#70"{typeof(loss_rd)})(x::Vector{Float64}, p::SciMLBase.NullParameters)
    @ DiffEqFlux ~/.julia/packages/DiffEqFlux/alPQ3/src/train.jl:3
  [9] (::OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing})(::Vector{Float64}, ::Vararg{Any, N} where N) (repeats 2 times)
    @ SciMLBase ~/.julia/packages/SciMLBase/9EjAY/src/problems/basic_problems.jl:107
 [10] WARNING: both Iterators and Flux export "flatten"; uses of it in module GalacticOptim must be qualified
WARNING: both Flux and Tracker export "gradient"; uses of it in module GalacticOptim must be qualified
(::GalacticOptim.var"#17#25"{OptimizationProblem{false, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}, Vector{Float64}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:cb, :maxiters), Tuple{var"#7#10", Int64}}}}, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}})(θ::Vector{Float64})
    @ GalacticOptim ~/.julia/packages/GalacticOptim/JnLwV/src/solve.jl:165
 [11] (::GalacticOptim.var"#18#26"{GalacticOptim.var"#17#25"{OptimizationProblem{false, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}, Vector{Float64}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:cb, :maxiters), Tuple{var"#7#10", Int64}}}}, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}}, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}})(G::Vector{Float64}, θ::Vector{Float64})
    @ GalacticOptim ~/.julia/packages/GalacticOptim/JnLwV/src/solve.jl:173
 [12] value_gradient!!(obj::TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}, x::Vector{Float64})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/geyh3/src/interface.jl:82
 [13] value_gradient!(obj::TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}, x::Vector{Float64})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/geyh3/src/interface.jl:69
 [14] value_gradient!(obj::Optim.ManifoldObjective{TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, x::Vector{Float64})
    @ Optim ~/.julia/packages/Optim/uwNqi/src/Manifolds.jl:50
 [15] (::LineSearches.var"#ϕdϕ#6"{Optim.ManifoldObjective{TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Vector{Float64}, Vector{Float64}, Vector{Float64}})(α::Float64)
    @ LineSearches ~/.julia/packages/LineSearches/Ki4c5/src/LineSearches.jl:84
 [16] secant2!(ϕdϕ::LineSearches.var"#ϕdϕ#6"{Optim.ManifoldObjective{TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Vector{Float64}, Vector{Float64}, Vector{Float64}}, alphas::Vector{Float64}, values::Vector{Float64}, slopes::Vector{Float64}, ia::Int64, ib::Int64, phi_lim::Float64, delta::Float64, sigma::Float64, display::Int64)
    @ LineSearches ~/.julia/packages/LineSearches/Ki4c5/src/hagerzhang.jl:368
 [17] (::LineSearches.HagerZhang{Float64, Base.RefValue{Bool}})(ϕ::Function, ϕdϕ::LineSearches.var"#ϕdϕ#6"{Optim.ManifoldObjective{TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Vector{Float64}, Vector{Float64}, Vector{Float64}}, c::Float64, phi_0::Float64, dphi_0::Float64)
    @ LineSearches ~/.julia/packages/LineSearches/Ki4c5/src/hagerzhang.jl:269
 [18] HagerZhang
    @ ~/.julia/packages/LineSearches/Ki4c5/src/hagerzhang.jl:101 [inlined]
 [19] perform_linesearch!(state::Optim.BFGSState{Vector{Float64}, Matrix{Float64}, Float64, Vector{Float64}}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}, d::Optim.ManifoldObjective{TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}})
    @ Optim ~/.julia/packages/Optim/uwNqi/src/utilities/perform_linesearch.jl:59
 [20] update_state!(d::TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}, state::Optim.BFGSState{Vector{Float64}, Matrix{Float64}, Float64, Vector{Float64}}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat})
    @ Optim ~/.julia/packages/Optim/uwNqi/src/multivariate/solvers/first_order/bfgs.jl:129
 [21] optimize(d::TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}, initial_x::Vector{Float64}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}, options::Optim.Options{Float64, GalacticOptim.var"#_cb#24"{var"#7#10", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}, Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}}}, state::Optim.BFGSState{Vector{Float64}, Matrix{Float64}, Float64, Vector{Float64}})
    @ Optim ~/.julia/packages/Optim/uwNqi/src/multivariate/optimize/optimize.jl:57
 [22] optimize(d::TwiceDifferentiable{Float64, Vector{Float64}, Matrix{Float64}, Vector{Float64}}, initial_x::Vector{Float64}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}, options::Optim.Options{Float64, GalacticOptim.var"#_cb#24"{var"#7#10", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}, Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}}})
    @ Optim ~/.julia/packages/Optim/uwNqi/src/multivariate/optimize/optimize.jl:35
 [23] __solve(prob::OptimizationProblem{false, OptimizationFunction{false, GalacticOptim.AutoZygote, OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true, GalacticOptim.AutoZygote, DiffEqFlux.var"#69#70"{typeof(loss_rd)}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#154#164", Nothing, Nothing, Nothing}, Vector{Float64}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:cb, :maxiters), Tuple{var"#7#10", Int64}}}}, opt::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}, data::Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}; maxiters::Int64, cb::Function, progress::Bool, kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ GalacticOptim ~/.julia/packages/GalacticOptim/JnLwV/src/solve.jl:182
 [24] #solve#468
    @ ~/.julia/packages/SciMLBase/9EjAY/src/solve.jl:3 [inlined]
 [25] sciml_train(::typeof(loss_rd), ::Vector{Float64}, ::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Flat}, ::GalacticOptim.AutoZygote; lower_bounds::Nothing, upper_bounds::Nothing, kwargs::Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:cb, :maxiters), Tuple{var"#7#10", Int64}}})
    @ DiffEqFlux ~/.julia/packages/DiffEqFlux/alPQ3/src/train.jl:6
 [26] top-level scope

Lotka volterra UDE not recovered as per paper

scenario_2.jl outputs in following fit for UDE.. scenario_1.jl also does does not fit (I remember in early versions, this examples were able to recover the model accurately.. what changed ?)

line 207 error DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths 22 and 106")

line 207

L = β.(0:tspan[end],β0,D,N,κ,α).*S.*I./N

i get this error.
DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths 22 and 106")

Stacktrace:
[1] _bcs1 at .\broadcast.jl:490 [inlined]
[2] _bcs at .\broadcast.jl:484 [inlined]
[3] broadcast_shape at .\broadcast.jl:478 [inlined]
[4] combine_axes at .\broadcast.jl:473 [inlined]
[5] _axes at .\broadcast.jl:208 [inlined]
[6] axes at .\broadcast.jl:206 [inlined]
[7] combine_axes at C:\Users\matias.julia\packages\DiffEqBase\OmQhO\src\diffeqfastbc.jl:19 [inlined]
[8] _axes at .\broadcast.jl:208 [inlined]
[9] axes at .\broadcast.jl:206 [inlined]
[10] combine_axes at C:\Users\matias.julia\packages\DiffEqBase\OmQhO\src\diffeqfastbc.jl:19 [inlined]
[11] _axes at .\broadcast.jl:208 [inlined]
[12] axes at .\broadcast.jl:206 [inlined]
[13] combine_axes at C:\Users\matias.julia\packages\DiffEqBase\OmQhO\src\diffeqfastbc.jl:19 [inlined]
[14] instantiate at .\broadcast.jl:256 [inlined]
[15] materialize(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1},Nothing,typeof(/),Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1},Nothing,typeof(),Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1},Nothing,typeof(),Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1},Nothing,typeof(β),Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},Float64,SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true},SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true},Float64,Float64}},SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}}},SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}}},SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}}}) at .\broadcast.jl:820
[16] top-level scope at In[16]:7

type SparseIdentificationResult has no field basis

When I try to run the code, I get the following error:

universal_differential_equations/SEIR_exposure/seir_exposure.jl

Line 223 in dd83689

println(Ψ_direct.basis)

type SparseIdentificationResult has no field basis

Thanks!

SEIR Example

Translation of SEIR Example, based on Lotka Volterra 1:

Hiya, ok, here's the first...

• I can't get it to train only for dE, dI, dR the way the old example did, but seems like the way I did it should work (also, it seemed to be linearising the other equations in order to get there back when this did work on the old example, so...)
• It doesn't predict correctly, for either the UDE or NODE. UDE prediction is linear, same as the problem I was having before, NODE is non-linear but incorrect. (Hopefully it's something simple with the ADAM and BFGS setup, because I don't fully understand that part of the code?)
• I don't know the correct code to extrapolate after training (lines 155-158 and 248-251)
• I haven't done the SINDy part of the code yet since the approximation doesn't work.
• Why can't I drop the .jl file in here?
• The savefigs don't work in the Lotka Volterra examples, they can be changed to this format.

cd(@__DIR__)
using Pkg; Pkg.activate("."); Pkg.instantiate()

# Single experiment, move to ensemble further on
# Some good parameter values are stored as comments right now
# because this is really good practice
using OrdinaryDiffEq
using ModelingToolkit
using DataDrivenDiffEq
using LinearAlgebra 
using SciMLSensitivity
using Random
using Optimization, OptimizationFlux, OptimizationOptimJL #OptimizationFlux for ADAM and OptimizationOptimJL for BFGS
using Lux
using Statistics
using Plots
gr()
#using DiffEqSensitivity**, Optim**
#using DiffEqFlux**, Flux**

function corona!(du,u,p,t)
    S,E,I,R,N,D,C = u
    F, β0,α,κ,μ,σ,γ,d,λ = p
    dS = -β0*S*F/N - β(t,β0,D,N,κ,α)*S*I/N -μ*S # susceptible
    dE = β0*S*F/N + β(t,β0,D,N,κ,α)*S*I/N -(σ+μ)*E # exposed
    dI = σ*E - (γ+μ)*I # infected
    dR = γ*I - μ*R # removed (recovered + dead)
    dN = -μ*N # total population
    dD = d*γ*I - λ*D # severe, critical cases, and deaths
    dC = σ*E # +cumulative cases

    du[1] = dS; du[2] = dE; du[3] = dI; du[4] = dR
    du[5] = dN; du[6] = dD; du[7] = dC
end

β(t,β0,D,N,κ,α) = β0*(1-α)*(1-D/N)^κ
S0 = 14e6
u0 = [0.9*S0, 0.0, 0.0, 0.0, S0, 0.0, 0.0]
p_ = [10.0, 0.5944, 0.4239, 1117.3, 0.02, 1/3, 1/5,0.2, 1/11.2]
R0 = p_[2]/p_[7]*p_[6]/(p_[6]+p_[5])
tspan = (0.0, 21.0)
prob = ODEProblem(corona!, u0, tspan, p_)
solution = solve(prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 1)
t = solution.t
#[2:4] are Exposed, Infected, Removed 
X = Array(solution[2:4,:])'
plot(X)

#Extrapolate to a longer timespan
tspan2 = (0.0,60.0)
prob = ODEProblem(corona!, u0, tspan2, p_)
solution_extrapolate = solve(prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 1)
extrapolate = Array(solution_extrapolate[2:4,:])'
plot(extrapolate)

# Ideal data
tsdata = Array(solution)

# Add noise to the data
noisy_data = tsdata + Float32(1e-5)*randn(eltype(tsdata), size(tsdata))
# You can see that the noise looks random
plot(abs.(tsdata-noisy_data)')

### Neural ODE
#Predicts for unknown equations
rng = Random.default_rng()
Random.seed!(111)

#7 inputs for 7 equations, 5 outputs because we know 2 equations already
U = Lux.Chain(Lux.Dense(7, 64, tanh),Lux.Dense(64, 64, tanh), Lux.Dense(64, 64, tanh), Lux.Dense(64, 5))

# Get the initial parameters and state variables of the model
p, st = Lux.setup(rng, U)

function coronaNODE(du,u,p,t,p_)
    û = U(u, p, st)[1] # Network prediction
    S,E,I,R,N,D,C = u
    μ,σ = p_
    dS = û[1]
    dE = û[2]
    dI = û[3]
    dR = û[4]
    dN = -μ*N # total population
    dD = û[5]
    dC = σ*E # +cumulative cases
    du[1] = dS; du[2] = dE; du[3] = dI; du[4] = dR
    du[5] = dN; du[6] = dD; du[7] = dC
end

# Closure with the known parameters
NODE_dynamics!(du,u,p,t) = coronaNODE(du,u,p,t,p_)
# Define the problem
prob_node = ODEProblem(NODE_dynamics!, u0, tspan, p)

## Function to train the network
# Define a predictor
function predict(θ, X = noisy_data[:,1], T = t)
    Array(solve(prob_node, Vern7(), u0 = X, p=θ,
                saveat = T,
                abstol=1e-6, reltol=1e-6,
                sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP())))
end

# Simple L2 loss
function loss(θ)
    X̂ = predict(θ)
    sum(abs2, noisy_data .- X̂)
end


# Container to track the losses
losses = Float32[]

callback = function (p, l)
  push!(losses, l)
  if length(losses)%50==0
      println("Current loss after $(length(losses)) iterations: $(losses[end])")
  end
  return false
end

## Training

# First train with ADAM for better convergence -> move the parameters into a
# favourable starting positing for BFGS
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p)->loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, Lux.ComponentArray(p))
res1 = Optimization.solve(optprob, ADAM(0.01), callback=callback, maxiters = 200)
println("Training loss after $(length(losses)) iterations: $(losses[end])")
# Train with BFGS
optprob2 = Optimization.OptimizationProblem(optf, res1.minimizer)
res2 = Optimization.solve(optprob2, Optim.BFGS(initial_stepnorm=0.01), callback=callback, maxiters = 10000)
println("Final training loss after $(length(losses)) iterations: $(losses[end])")

# Plot the losses
pl_losses = plot(1:200, losses[1:200], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "ADAM", color = :blue)
plot!(201:length(losses), losses[201:end], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "BFGS", color = :red)
savefig(pl_losses, "pl_lossesNODE.png")

# Rename the best candidate
p_trained = res2.minimizer

## Analysis of the trained network
# Plot the data and the approximation

# Make the prediction to match solution.t
X̂ = predict(p_trained, noisy_data[:,1], t)
# Prediction trained on noisy data vs real solution
pl_trajectory = plot(t, transpose(X̂[2:4,:]), xlabel = "t", ylabel ="x(t), y(t)", color = :red, label = ["NODE Approximation" nothing])
scatter!(solution.t, transpose(noisy_data[2:4,:]), color = :black, label = ["Measurements" nothing])
savefig(pl_trajectory, "plots_trajectory_reconstructionNODE.png")

#Extrapolate the solution to match tspan2
ExtrapolateX̂ = predict(p_trained, noisy_data[:,1], solution_extrapolate.t)
extrapolate_trajectory = plot(solution_extrapolate.t, transpose(ExtrapolateX̂[2:4,:]), xlabel = "t", ylabel ="x(t), y(t)", color = :red, label = ["NODE Approximation" nothing])
scatter!(solution_extrapolate.t, transpose(solution_extrapolate[2:4,:]), color = :black, label = ["Measurements" nothing])
savefig(pl_trajectory, "ExtrapolateNODE.png")

### Universal ODE
##Prediction for missing parameters
rng = Random.default_rng()
Random.seed!(222)

#7 inputs for 7 equations, 1 output for 1 missing part of the equation
U = Lux.Chain(Lux.Dense(7, 64, tanh),Lux.Dense(64, 64, tanh), Lux.Dense(64, 64, tanh), Lux.Dense(64, 1))

# Get the initial parameters and state variables of the model
p, st = Lux.setup(rng, U)

function coronaUDE(du,u,p,t,p_true)
    û = U(u, p, st)[1] # Network prediction
    S,E,I,R,N,D,C = u
    F,β0,α,κ,μ,σ,γ,d,λ = p_
    dS = -β0*S*F/N - û[1] -μ*S # susceptible
    dE = β0*S*F/N + û[1] -(σ+μ)*E # exposed
    dI = σ*E - (γ+μ)*I # infected
    dR = γ*I - μ*R # removed (recovered + dead)
    dN = -μ*N # total population
    dD = d*γ*I - λ*D # severe, critical cases, and deaths
    dC = σ*E # +cumulative cases
    du[1] = dS; du[2] = dE; du[3] = dI; du[4] = dR
    du[5] = dN; du[6] = dD; du[7] = dC
end

# Closure with the known parameters
UDE_dynamics!(du,u,p,t) = coronaUDE(du,u,p,t,p_)
# Define the problem
prob_ude = ODEProblem(UDE_dynamics!, u0, tspan, p)

## Function to train the network
# Define a predictor
function predict(θ, X = noisy_data[:,1], T = t)
    Array(solve(prob_ude, Vern7(), u0 = X, p=θ,
                saveat = T,
                abstol=1e-6, reltol=1e-6,
                sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP())))
end

# Simple L2 loss
function loss(θ)
    X̂ = predict(θ)
    sum(abs2, noisy_data .- X̂)
end


# Container to track the losses
losses = Float32[]

callback = function (p, l)
  push!(losses, l)
  if length(losses)%50==0
      println("Current loss after $(length(losses)) iterations: $(losses[end])")
  end
  return false
end

## Training

# First train with ADAM for better convergence -> move the parameters into a
# favourable starting positing for BFGS
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p)->loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, Lux.ComponentArray(p))
res1UDE = Optimization.solve(optprob, ADAM(0.01), callback=callback, maxiters = 200)
println("Training loss after $(length(losses)) iterations: $(losses[end])")
# Train with BFGS
optprob2 = Optimization.OptimizationProblem(optf, res1.minimizer)
res2UDE = Optimization.solve(optprob2, Optim.BFGS(initial_stepnorm=0.01), callback=callback, maxiters = 10000)
println("Final training loss after $(length(losses)) iterations: $(losses[end])")

# Plot the losses
pl_losses = plot(1:200, losses[1:200], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "ADAM", color = :blue)
plot!(201:length(losses), losses[201:end], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "BFGS", color = :red)
savefig(pl_losses, "plot_lossesUDE.png")
# Rename the best candidate
p_trained = res2UDE.minimizer

## Analysis of the trained network
# Plot the data and the approximation
X̂ = predict(p_trained, noisy_data[:,1], t)
# Trained on noisy data vs real solution
pl_trajectory = plot(t, transpose(X̂[2:4,:]), xlabel = "t", ylabel ="x(t), y(t)", color = :red, label = ["UDE Approximation" nothing])
scatter!(solution.t, transpose(solution[2:4,:]), color = :black, label = ["Measurements" nothing])
savefig(pl_trajectory, "plot_trajectory_reconstructionUDE.png")

# Extrapolate out
ExtrapolateX̂ = predict(p_trained, noisy_data[:,1], solution_extrapolate.t)
extrapolate_trajectory = plot(solution_extrapolate.t, transpose(ExtrapolateX̂[2:4,:]), xlabel = "t", ylabel ="x(t), y(t)", color = :red, label = ["UDE Approximation" nothing])
scatter!(solution_extrapolate.t, transpose(solution_extrapolate[2:4,:]), color = :black, label = ["Measurements" nothing])
savefig(extrapolate_trajectory, "ExtrapolateUDE.png")

outdated / soon to be outdated example codes

The code for the sindy to equations parts of the examples are out of date and won't run.
In addition, running the SEIR model, the following warnings came up:
FastChain is being deprecated in favor of Lux.jl
sciml_train is being deprecated in favor of direct usage of Optimization.jl
Seems like these examples will soon be outdated, but they're really useful.

MethodError: objects of type SparseIdentificationResult are not callable

universal_differential_equations/SEIR_exposure/seir_exposure.jl

Line 235 in dd83689

z = Ψ([S/N,I,D/N]) # Exposure does not depend on exposed, removed, or cumulative!

I get this error when attempting to the function as it is
"MethodError: objects of type SparseIdentificationResult are not callable"

OptimizationFunction Definition

Hello,

I was investigating the LotkaVolterra problem scripts and i saw that the optimization function is defined as follows;

optf = Optimization.OptimizationFunction((x,p)->loss(x), adtype)

i am wondering why loss function is given to optimization-function definition as : (x,p)->loss(x) in here while it is only depending on ''p'' ? Since, it isn't working when it is written as loss(p).

"MethodError: objects of type SparseIdentificationResult are not callable"

universal_differential_equations/SEIR_exposure/seir_exposure.jl

Line 235 in dd83689

z = Ψ([S/N,I,D/N]) # Exposure does not depend on exposed, removed, or cumulative!

I get this error when attempting to the function as it is
"MethodError: objects of type SparseIdentificationResult are not callable"

No matching function wrapper found

I am relatively new to Julia, so please excuse potentially basic questions. I tried running the Lotka Volterra scripts using Julia 1.8. Once I run the line "Optimization.solve(...)" I get a very long error message concluding in "No matching function wrapper was found!". I tried it in a new environment,... Do you have any recommendations how I could solve this?

Hudson Bay

I couldn't find an example of multiple shooting in SciMLSensitivity, sorry - it's coming up with the error "UndefVarError: multiple_shoot not defined" when I try to run ADAM.

I also tried FENEP but it seems like there's some issue with Tracked Arrays, here: SciML/SciMLSensitivity.jl#609
And also with Flux, because the original doesn't work with it. I get the error: UndefVarError: TrackedArray not defined
So i'll leave this one.

using OrdinaryDiffEq
using ModelingToolkit
using DataDrivenDiffEq
using LinearAlgebra 
using SciMLSensitivity
using Random
using Optimization, OptimizationOptimisers, OptimizationOptimJL #OptimizationOptimisers for ADAM and OptimizationOptimJL for BFGS
using Lux
using Statistics
using JLD2, FileIO
using DelimitedFiles
using Plots
gr()

Random.seed!(5443)

#### NOTE
# Since the recent release of DataDrivenDiffEq v0.6.0 where a complete overhaul of the optimizers took
# place, SR3 has been used. Right now, STLSQ performs better and has been changed.
# Additionally, the behaviour of the optimization has changed slightly. This has been adjusted
# by decreasing the tolerance of the gradient.


svname = "HudsonBay"
## Data Preprocessing
# The data has been taken from https://jmahaffy.sdsu.edu/courses/f00/math122/labs/labj/q3v1.htm
# Originally published in E. P. Odum (1953), Fundamentals of Ecology, Philadelphia, W. B. Saunders
hudson_bay_data = readdlm("hudson_bay_data.dat", '\t', Float32, '\n')
# Measurements of prey and predator
Xₙ = Matrix(transpose(hudson_bay_data[:, 2:3]))
t = hudson_bay_data[:, 1] .- hudson_bay_data[1, 1]
# Normalize the data; since the data domain is strictly positive
# we just need to divide by the maximum
xscale = maximum(Xₙ, dims =2)
Xₙ .= 1f0 ./ xscale .* Xₙ
# Time from 0 -> n
tspan = (t[1], t[end])

# Plot the data
scatter(t, transpose(Xₙ), xlabel = "t", ylabel = "x(t), y(t)")
plot!(t, transpose(Xₙ), xlabel = "t", ylabel = "x(t), y(t)")

## Direct Identification via SINDy + Collocation

# Create the problem using a gaussian kernel for collocation
full_problem = ContinuousDataDrivenProblem(Xₙ, t, DataDrivenDiffEq.GaussianKernel())
# Look at the collocation
plot(full_problem.t, full_problem.X')
plot(full_problem.t, full_problem.DX')

# Create a Basis
@variables u[1:2]

# Generate the basis functions, multivariate polynomials up to deg 5
# and sine
b = [polynomial_basis(u, 5); sin.(u)]
basis = Basis(b, u)

# Create the thresholds which should be used in the search process
λ = Float32.(exp10.(-7:0.1:5))
# Create an optimizer for the SINDy problem
opt = STLSQ(λ)

# Best result so far
full_res = solve(full_problem, basis, opt, maxiter = 10000, progress = true, denoise = true, normalize = true)

println(full_res)
println(result(full_res))

## Define the network
# Gaussian RBF as activation
rbf(x) = exp.(-(x.^2))

# Define the network 2->5->5->5->2
U = Lux.Chain(
    Lux.Dense(2,5,rbf), Lux.Dense(5,5, rbf), Lux.Dense(5,5, tanh), Lux.Dense(5,2)
)

# Get the initial parameters, first two is linear birth / decay of prey and predator
rng = Random.default_rng()
p, st = Lux.setup(rng, U)

# Define the hybrid model
function ude_dynamics!(du,u, p, t)
    û = U(u, p[3:end]) # Network prediction
    # We assume a linear birth rate for the prey
    du[1] = p[1]*u[1] + û[1]
    # We assume a linear decay rate for the predator
    du[2] = -p[2]*u[2] + û[2]
end

# Define the problem
prob_nn = ODEProblem(ude_dynamics!,Xₙ[:, 1], tspan, p)

## Function to train the network
# Define a predictor
function predict(θ, X = Xₙ[:,1], T = t)
    Array(solve(prob_nn, Vern7(), u0 = X, p=θ,
                tspan = (T[1], T[end]), saveat = T,
                abstol=1e-6, reltol=1e-6,
                sensealg = ForwardDiffSensitivity()
                ))
end


# Define parameters for Multiple Shooting
group_size = 5
continuity_term = 200.0f0

function loss(data, pred)
	return sum(abs2, data - pred)
end

function shooting_loss(p)
    return multiple_shoot(p, Xₙ, t, prob_nn, loss, Vern7(),
                          group_size; continuity_term)
end

function loss(θ)
    X̂ = predict(θ)
    sum(abs2, Xₙ - X̂) / size(Xₙ, 2) + convert(eltype(θ), 1e-3)*sum(abs2, θ[3:end]) ./ length(θ[3:end])
end

# Container to track the losses
losses = Float32[]

# Callback to show the loss during training
callback(θ,args...) = begin
	l = loss(θ) # Equivalent L2 loss
    push!(losses, l)
    if length(losses)%5==0
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    false
end

## Training -> First shooting / batching to get a rough estimate

# First train with ADAM for better convergence -> move the parameters into a
# favourable starting positing for BFGS
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p)->shooting_loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, Lux.ComponentArray(p))
res1 = Optimization.solve(optprob, ADAM(0.1f0), cb=callback, maxiters = 100)
println("Training loss after $(length(losses)) iterations: $(losses[end])")
# Train with BFGS to achieve partial fit of the data
optprob2 = Optimization.OptimizationProblem(optf, res1.minimizer)
res2 = Optimization.solve(optprob2, Optim.BFGS(initial_stepnorm=0.01f0), cb=callback, maxiters = 500)
println("Training loss after $(length(losses)) iterations: $(losses[end])")
# Full L2-Loss for full prediction

optf2 = Optimization.OptimizationFunction((x,p)->loss(x), adtype)
optprob2 = Optimization.OptimizationProblem(optf2, res2.minimizer)
res3 = Optimization.solve(optprob2, Optim.BFGS(initial_stepnorm=0.01f0), cb=callback, maxiters = 10000)
println("Final training loss after $(length(losses)) iterations: $(losses[end])")

pl_losses = plot(1:101, losses[1:101], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "ADAM (Shooting)", color = :blue)
plot!(102:302, losses[102:302], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "BFGS (Shooting)", color = :red)
plot!(302:length(losses), losses[302:end], color = :black, label = "BFGS (L2)")
savefig(pl_losses, "plot_losses.png"))

LV Scenario 2

Incorrect predictions, again....

using OrdinaryDiffEq
using ModelingToolkit
using DataDrivenDiffEq
using LinearAlgebra 
using Optimization, OptimizationFlux, OptimizationOptimJL #OptimizationOptimisers for ADAM and OptimizationOptimJL for BFGS
using Lux
using SciMLSensitivity
using Plots
gr()
using JLD2, FileIO
using Statistics
# Set a random seed for reproduceable behaviour
using Random
Random.seed!(2345)

#### NOTE
# Since the recent release of DataDrivenDiffEq v0.6.0 where a complete overhaul of the optimizers took
# place, SR3 has been used. Right now, STLSQ performs better and has been changed.

#### NOTE
# Since the recent release of DataDrivenDiffEq v0.6.0 where a complete overhaul of the optimizers took
# place, SR3 has been used. Right now, STLSQ performs better and has been changed.


# Create a name for saving ( basically a prefix )
svname = "Scenario_2_"

## Data generation
function lotka!(du, u, p, t)
    α, β, γ, δ = p
    du[1] = α*u[1] - β*u[2]*u[1]
    du[2] = γ*u[1]*u[2]  - δ*u[2]
end

# Define the experimental parameter
tspan = (0.0f0,6.0f0)
u0 = Float32[0.44249296,4.6280594]
p_ = Float32[1.3, 0.9, 0.8, 1.8]
prob = ODEProblem(lotka!, u0,tspan, p_)
solution = solve(prob, Vern7(), abstol=1e-6, reltol=1e-6, saveat = 0.1)

scatter(solution, alpha = 0.25)
plot!(solution, alpha = 0.5)

# Ideal data
X = Array(solution)
t = solution.t
# Add noise in terms of the mean
x̄ = mean(X, dims = 2)
noise_magnitude = Float32(1e-2)
Xₙ = X .+ (noise_magnitude*x̄) .* randn(eltype(X), size(X))

# Subsample the data in y
# We assume we have only 5 measurements in y, evenly distributed
ty = collect(t[1]:Float32(6/5):t[end])
# Create datasets for the different measurements
round(Int64, mean(diff(ty))/mean(diff(t)))
XS = zeros(eltype(X), length(ty)-1, floor(Int64, mean(diff(ty))/mean(diff(t)))+1) # All x data
TS = zeros(eltype(t), length(ty)-1, floor(Int64, mean(diff(ty))/mean(diff(t)))+1) # Time data
YS = zeros(eltype(X), length(ty)-1, 2) # Just two measurements in y

for i in 1:length(ty)-1
    idxs = ty[i].<= t .<= ty[i+1]
    XS[i, :] = Xₙ[1, idxs]
    TS[i, :] = t[idxs]
    YS[i, :] = [Xₙ[2, t .== ty[i]]'; Xₙ[2, t .== ty[i+1]]]
end

XS

scatter!(t, transpose(Xₙ))
## Define the network
# Gaussian RBF as activation
rbf(x) = exp.(-(x.^2))

# Define the network 2->5->5->5->2
U = Lux.Chain(
    Lux.Dense(2,5,rbf), Lux.Dense(5,5, rbf), Lux.Dense(5,5, rbf), Lux.Dense(5,2)
)

rng = Random.default_rng()
p1, st = Lux.setup(rng, U)

#for birth, decay parameters -> initializing random values.
parameter_array = Float64[0.5]
p = (layer_1 = p1, layer_2 = parameter_array)
p = Lux.ComponentArray(p)

# Define the hybrid model
function ude_dynamics!(du,u, p, t, p_true)
    û = U(u, p.layer_1, st)[1] # Network prediction
    du[1] = p_true[1]*u[1] + û[1]
    # We assume a linear decay rate for the predator
    du[2] = -p.layer_2[1]*u[2] + û[2]
end

p_true = 1.3
# Closure with the known parameter
nn_dynamics!(du,u,p,t) = ude_dynamics!(du,u,p,t,p_true)
# Define the problem

prob_nn = ODEProblem(nn_dynamics!,Xₙ[:, 1], tspan, p)

## Function to train the network
# Define a predictor
function predict(θ, X = Xₙ[:,1], T = t)
    Array(solve(prob_nn, Vern7(), u0 = X, p=θ,
                saveat = T,
                abstol=1e-6, reltol=1e-6,
                sensealg = ForwardDiffSensitivity()
                ))
end

# Multiple shooting like loss
function loss(θ)
    # Start with a regularization on the network
    l = convert(eltype(θ), 1e-3)*sum(abs2, θ[2:end]) ./ length(θ[2:end])
    for i in 1:size(XS,1)
        X̂ = predict(θ, [XS[i,1], YS[i,1]], TS[i, :])
        # Full prediction in x
        l += sum(abs2, XS[i,:] .- X̂[1,:])
        # Add the boundary condition in y
        l += abs(YS[i, 2] .- X̂[2, end])
    end
    return l
end

# Container to track the losses
losses = Float32[]

# Callback to show the loss during training
callback(θ,l) = begin
    push!(losses, l)
    if length(losses)%50==0
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    false
end

## Training

# First train with ADAM for better convergence -> move the parameters into a
# favourable starting positing for BFGS
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p)->loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, Lux.ComponentArray(p))
res1 = Optimization.solve(optprob, ADAM(0.01f0), callback=callback, maxiters = 300)
println("Training loss after $(length(losses)) iterations: $(losses[end])")
# Train with BFGS to achieve partial fit of the data
optprob2 = remake(optprob,u0 = res1.u)
res2 = Optimization.solve(optprob2, Optim.BFGS(initial_stepnorm=0.01f0), callback=callback, maxiters = 10000, g_tol = 1e-10)
println("Training loss after $(length(losses)) iterations: $(losses[end])")

# Plot the losses
pl_losses = plot(1:300, losses[1:300], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "ADAM", color = :blue)
plot!(301:length(losses), losses[301:end], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "BFGS", color = :red)
savefig(pl_losses, "plot_losses.png")

# Rename the best candidate
p_trained = res2.minimizer

## Analysis of the trained network
# Plot the data and the approximation
ts = first(solution.t):mean(diff(solution.t))/2:last(solution.t)
X̂ = predict(p_trained, Xₙ[:, 1], ts)
# Trained on noisy data vs real solution
pl_trajectory = plot(ts, transpose(X̂), ylabel = "t", xlabel ="x(t), y(t)", color = :red, label = ["UDE Approximation" nothing])
scatter!(t, X[1,:], color = :black, label = "Measurements")
ymeasurements = unique!(vcat(YS...))
tmeasurements = unique!(vcat([[ts[1], ts[end]] for ts in eachrow(TS)]...))
scatter!(tmeasurements, ymeasurements, color = :black, label = nothing, legend = :topleft)
savefig(pl_trajectory, "plot_trajectory_reconstruction.png")

# Ideal unknown interactions of the predictor
Ȳ = [-p_[2]*(X̂[1,:].*X̂[2,:])';p_[3]*(X̂[1,:].*X̂[2,:])']
# Neural network guess
Ŷ = U(X̂, p_trained.layer_1, st)[1]

pl_reconstruction = plot(ts, transpose(Ŷ), xlabel = "t", ylabel ="U(x,y)", color = :red, label = ["UDE Approximation" nothing])
plot!(ts, transpose(Ȳ), color = :black, label = ["True Interaction" nothing], legend = :topleft)
savefig(pl_reconstruction, "plot_missingterm_reconstruction.png")

# Plot the error
pl_reconstruction_error = plot(ts, norm.(eachcol(Ȳ-Ŷ)), yaxis = :log, xlabel = "t", ylabel = "L2-Error", color = :red, label = nothing)
pl_missing = plot(pl_reconstruction, pl_reconstruction_error, layout = (2,1))
savefig(pl_missing, "plots_missingterm_reconstruction_and_error.pdf")
pl_overall = plot(pl_trajectory, pl_missing)
savefig(pl_overall, "plots_reconstruction.png")

@RajDandekar

Has this paper been published on a journal or a conference?

This brilliant work gives me a lot of insights! I want to cite it in my paper now. ♥
Is it published? A peer-reviewed article sounds much better than a preprint.

recompile after activating

Do I really need to recompile every package (OrdinaryDiffEq, ModelingToolki ...) after Pkg.activate("."); Pkg.instantiate() ? It seems to be doing so ...

`MethodError: no method matching SInDy` in DelayLotkaVolterra/VolterraExp.jl

I'm running the DelayLotkaVolterra/VolterraExp.jl unmodified and I'm getting the following error:

MethodError: no method matching SInDy

Thanks!

training method for multiple neural-nets in universal ODEs

what is the way of training multiple neural-nets in universal ODEs . Do we train them together (if yes , how ?) or do we train them subsequently ?

NN_1 = FastChain(FastDense(2,32,tanh), FastDense(32,32,tanh),FastDense(32,1)) 
p1 = initial_params(NN_1)
NN_2 = FastChain(FastDense(2,32,tanh), FastDense(32,32,tanh),FastDense(32,1)) 
p2 = initial_params(NN_1)

function dudt_(u, p,t)
    x, y, z = u
    z1 = NN_1([x,y],p1)
    z2 = NN_2([y,z],p2)
    [p_[1]*x + z[1],
    -p_[3]*y + z[2],
     p_[4]*z - z[1] - z[2]
      ]
end # where p_[1], p_[3],p_[4] are known parameters passed through p

prob_nn = ODEProblem(dudt_,u0, tspan, p)
sol_nn = solve(prob_nn)

function predict()  # etc..etc
function loss()

How do I update/ train NN_1 and NN_2 together using say sci_ml if thats how I should proceed ?

TrackedArray Error when trying to reproduce the script

Hi,

Thank you so much for open sourcing the paper. I was trying to reproduce the results before trying to build on top of it, and ran into some issues.

This is on a fresh install of juliapro after installing all packages with instantiate.

julia> include("FullVolterraExp.jl") Updating registry at~/.juliapro/JuliaPro_v1.2.0-2/registries/General┌ Warning: Base.IOError("write: broken pipe (EPIPE)", -32) └ @ Pkg.Operations /Users/vagrant/worker/juliapro-release-osx1011-0_6/build/tmp_julia/Julia-1.2.app/Contents/Resources/julia/share/julia/stdlib/v1.2/Pkg/src/Operations.jl:485 Updating git-repohttps://pkg.juliacomputing.com/registry/General`
ERROR: LoadError: Not implemented: convert TrackedArray{…,Array{Float32,1}} to TrackedArray{…,Array{Float64,1}}
Stacktrace:
[1] convert(::Type{TrackedArray{…,Array{Float64,1}}}, ::TrackedArray{…,Array{Float32,1}}) at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/Tracker/cpxco/src/lib/array.jl:38
[2] setproperty!(::OrdinaryDiffEq.ODEIntegrator{Tsit5,false,TrackedArray{…,Array{Float64,1}},Float32,TrackedArray{…,Array{Float32,1}},Float32,Float32,Float32,Array{TrackedArray{…,Array{Float64,1}},1},ODESolution{Tracker.TrackedReal{Float64},2,Array{TrackedArray{…,Array{Float64,1}},1},Nothing,Nothing,Array{Float32,1},Array{Array{TrackedArray{…,Array{Float64,1}},1},1},ODEProblem{TrackedArray{…,Array{Float64,1}},Tuple{Float32,Float32},false,TrackedArray{…,Array{Float32,1}},ODEFunction{false,typeof(dudt_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODEFunction{false,typeof(dudt_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Array{TrackedArray{…,Array{Float64,1}},1},Array{Float32,1},Array{Array{TrackedArray{…,Array{Float64,1}},1},1},OrdinaryDiffEq.Tsit5ConstantCache{Float64,Float32}},DiffEqBase.DEStats},ODEFunction{false,typeof(dudt_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},OrdinaryDiffEq.Tsit5ConstantCache{Float64,Float32},OrdinaryDiffEq.DEOptions{Tracker.TrackedReal{Float64},Tracker.TrackedReal{Float64},Float32,Float32,typeof(DiffEqBase.ODE_DEFAULT_NORM),typeof(opnorm),CallbackSet{Tuple{},Tuple{}},typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN),typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE),typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK),DataStructures.BinaryHeap{Float32,DataStructures.LessThan},DataStructures.BinaryHeap{Float32,DataStructures.LessThan},Nothing,Nothing,Int64,Array{Float32,1},Array{Float32,1},Array{Float32,1}},TrackedArray{…,Array{Float64,1}},Tracker.TrackedReal{Float64},Nothing}, ::Symbol, ::TrackedArray{…,Array{Float32,1}}) at ./Base.jl:21
[3] initialize!(::OrdinaryDiffEq.ODEIntegrator{Tsit5,false,TrackedArray{…,Array{Float64,1}},Float32,TrackedArray{…,Array{Float32,1}},Float32,Float32,Float32,Array{TrackedArray{…,Array{Float64,1}},1},ODESolution{Tracker.TrackedReal{Float64},2,Array{TrackedArray{…,Array{Float64,1}},1},Nothing,Nothing,Array{Float32,1},Array{Array{TrackedArray{…,Array{Float64,1}},1},1},ODEProblem{TrackedArray{…,Array{Float64,1}},Tuple{Float32,Float32},false,TrackedArray{…,Array{Float32,1}},ODEFunction{false,typeof(dudt_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODEFunction{false,typeof(dudt_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Array{TrackedArray{…,Array{Float64,1}},1},Array{Float32,1},Array{Array{TrackedArray{…,Array{Float64,1}},1},1},OrdinaryDiffEq.Tsit5ConstantCache{Float64,Float32}},DiffEqBase.DEStats},ODEFunction{false,typeof(dudt_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},OrdinaryDiffEq.Tsit5ConstantCache{Float64,Float32},OrdinaryDiffEq.DEOptions{Tracker.TrackedReal{Float64},Tracker.TrackedReal{Float64},Float32,Float32,typeof(DiffEqBase.ODE_DEFAULT_NORM),typeof(opnorm),CallbackSet{Tuple{},Tuple{}},typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN),typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE),typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK),DataStructures.BinaryHeap{Float32,DataStructures.LessThan},DataStructures.BinaryHeap{Float32,DataStructures.LessThan},Nothing,Nothing,Int64,Array{Float32,1},Array{Float32,1},Array{Float32,1}},TrackedArray{…,Array{Float64,1}},Tracker.TrackedReal{Float64},Nothing}, ::OrdinaryDiffEq.Tsit5ConstantCache{Float64,Float32}) at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/OrdinaryDiffEq/DJItT/src/perform_step/low_order_rk_perform_step.jl:565
[4] #__init#332(::Array{Float32,1}, ::Array{Float32,1}, ::Array{Float32,1}, ::Nothing, ::Bool, ::Bool, ::Bool, ::Bool, ::Nothing, ::Bool, ::Bool, ::Float32, ::Float32, ::Float32, ::Bool, ::Bool, ::Rational{Int64}, ::Nothing, ::Nothing, ::Rational{Int64}, ::Int64, ::Int64, ::Int64, ::Rational{Int64}, ::Bool, ::Int64, ::Nothing, ::Nothing, ::Int64, ::typeof(DiffEqBase.ODE_DEFAULT_NORM), ::typeof(opnorm), ::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), ::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Int64, ::String, ::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), ::Nothing, ::Bool, ::Bool, ::Bool, ::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}}, ::typeof(DiffEqBase._init), ::ODEProblem{TrackedArray{…,Array{Float64,1}},Tuple{Float32,Float32},false,TrackedArray{…,Array{Float32,1}},ODEFunction{false,typeof(dudt),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}},DiffEqBase.StandardODEProblem}, ::Tsit5, ::Array{TrackedArray{…,Array{Float64,1}},1}, ::Array{Float32,1}, ::Array{Any,1}, ::Type{Val{true}}) at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/OrdinaryDiffEq/DJItT/src/solve.jl:357
[5] _init(::ODEProblem{TrackedArray{…,Array{Float64,1}},Tuple{Float32,Float32},false,TrackedArray{…,Array{Float32,1}},ODEFunction{false,typeof(dudt),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}},DiffEqBase.StandardODEProblem}, ::Tsit5, ::Array{TrackedArray{…,Array{Float64,1}},1}, ::Array{Float32,1}, ::Array{Any,1}, ::Type{Val{true}}) at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/OrdinaryDiffEq/DJItT/src/solve.jl:66 (repeats 4 times)
[6] #__solve#331 at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/OrdinaryDiffEq/DJItT/src/solve.jl:4 [inlined]
[7] solve at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/OrdinaryDiffEq/DJItT/src/solve.jl:4 [inlined]
[8] #solve_call#442(::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}}, ::typeof(DiffEqBase.solve_call), ::ODEProblem{TrackedArray{…,Array{Float64,1}},Tuple{Float32,Float32},false,TrackedArray{…,Array{Float32,1}},ODEFunction{false,typeof(dudt),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}},DiffEqBase.StandardODEProblem}, ::Tsit5) at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/DiffEqBase/jMOwn/src/solve.jl:40
[9] solve_call at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/DiffEqBase/jMOwn/src/solve.jl:37 [inlined]
[10] #solve#443 at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/DiffEqBase/jMOwn/src/solve.jl:57 [inlined]
[11] solve at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/DiffEqBase/jMOwn/src/solve.jl:45 [inlined]
[12] #diffeq_rd#7 at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/DiffEqFlux/o3yCl/src/Flux/layers.jl:19 [inlined]
[13] diffeq_rd(::TrackedArray{…,Array{Float32,1}}, ::ODEProblem{TrackedArray{…,Array{Float64,1}},Tuple{Float32,Float32},false,TrackedArray{…,Array{Float32,1}},ODEFunction{false,typeof(dudt),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}},DiffEqBase.StandardODEProblem}, ::Tsit5) at /Users/vaggi/.juliapro/JuliaPro_v1.2.0-2/packages/DiffEqFlux/o3yCl/src/Flux/layers.jl:7
[14] top-level scope at /Users/vaggi/Code_Libraries/Julia/universal_differential_equations/DelayLotkaVolterra/FullVolterraExp.jl:53
[15] include at ./boot.jl:328 [inlined]
[16] include_relative(::Module, ::String) at ./loading.jl:1094
[17] include(::Module, ::String) at ./Base.jl:31
[18] include(::String) at ./client.jl:431
[19] top-level scope at none:0
in expression starting at /Users/vaggi/Code_Libraries/Julia/universal_differential_equations/DelayLotkaVolterra/FullVolterraExp.jl:53`

I tried googling the error, and found this:
FluxML/Flux.jl#136

I hope I am not posting this in the wrong place - let me know if this belongs on discourse as well.

Non Newtonian Fluids: FENE-P example error

Here is the code which I have for now. For simplicity, I am currently using a single strain rate and not summing over multiple strain rates.

using DiffEqFlux, Flux
using DifferentialEquations
using Plots, Statistics
using Sundials,  Optimization, OptimizationFlux

function FENEP!(out,du,u,p,t,γd)
  # DAE definition for FENE-p
  θ₁₁,θ₂₂,θ₁₂, τ₁₁,τ₂₂,τ₁₂ = u
  λ,η,L = [2.0, 2.0, 4.0]
  a = L^2 /(L^2 -3)
  fλ = (L^2 + (λ/η/a)*(τ₂₂+τ₁₁))/(L^2 - 3)
  out[1] =  τ₁₁ + du[1] - 2*λ*γd(t)*τ₁₂/fλ
  out[2] =  τ₂₂ + du[2]
  out[3] =  τ₁₂ + du[3] - λ*γd(t)*τ₂₂/fλ - η/fλ * γd(t)

  out[4] = θ₁₁ - λ*τ₁₁/fλ
  out[5] = θ₂₂ - λ*τ₂₂/fλ
  out[6] = θ₁₂ - λ*τ₁₂/fλ
end

function find_σ_exact(tsave,γd)
  # finds the exact solution of the FENE-p equations given a strain rate fn and
  # time points to save at
  λ = 2.0
  L = 2.0
  η = 4.0
  p = [λ,η,L]
  u₀ = zeros(6)
  du₀ = [0.0, 0.0, η*γd(0.0)*(L^2-3)/L^2, 0.0,0.0,0.0]
  tspan = (Float64(tsave[1]),Float64(tsave[end]))
  differential_vars = [true,true,true,false,false,false]
  h(out, du,u,p,t) = FENEP!(out, du,u,p,t,γd)

  prob = DAEProblem(h,du₀,u₀,tspan,differential_vars=differential_vars)

  sol = solve(prob,IDA(),saveat=Float64.(tsave))
  return [Float32(σ[6]) for σ in sol.u]
end

p0_vec = Float64[]

f0_n = FastChain(FastDense(2,4,tanh), FastDense(4,1))
p0 = initial_params(f0_n)
append!(p0_vec, p0)

f1_n = FastChain(FastDense(2,4,tanh), FastDense(4,1))
p0 = initial_params(f1_n)
append!(p0_vec, p0)

function dudt_(du, u, p, t, γd)
    NN1 = abs(f1_n([u[1],γd(t) ], p[1:17])[1])
    du[1] = NN1
end

# Define the problem


α = p0_vec


#=
function predict_adjoint(θ)
  x = Array(solve(prob_pred,Tsit5(),p=θ,saveat=tsave,
                  sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true))))
end
=#

function loss_adjoint(θ, γd)
    tspan = (0.0f0,6.2831f0)
    tsave = range(tspan[1],tspan[2],length=100)
    σ_exact = find_σ_exact(tsave,γd)
    σ0 =0.0f0
    u0 = [σ0]

    dudt2_(du,u,p,t) = dudt_(du,u,p,t,γd)
    prob_pred = ODEProblem{true}(dudt2_,u0,tspan)
    x = Array(solve(prob_pred,Tsit5(),p=θ,saveat=tsave,
                    sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true))))

    P_RD = vcat(x, γd.(tsave)')
    σ_out = [f0_n(P_RD[:,i], θ[18:34])[1] for i = 1:size(P_RD,2)]
    return sum( (σ_out .- σ_exact).^2 )
end

#loss_adjoint2(θ, γd) =  sum(loss_adjoint(θ,t -> 12.0f0*cos.(ω.*t))() for ω in 1.0f0:0.2f0:2.0f0)

#γd = t -> 12.0f0*cos.(1.6.*t)

iter = 0
function callback(θ,l)
  global iter
  iter += 1
  if iter%10 == 0
    println(l)
  end
  return false
end

ω = 1.0f0
γd = t -> 12.0f0*cos.(ω.*t)

The error which I get on the last line is this:

ERROR: MethodError: no method matching length(::SciMLBase.NullParameters)
Closest candidates are:
  length(::Union{Base.KeySet, Base.ValueIterator}) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/base/abstractdict.jl:58
  length(::Union{DataStructures.OrderedRobinDict, DataStructures.RobinDict}) at ~/.julia/packages/DataStructures/59MD0/src/ordered_robin_dict.jl:86
  length(::Union{DataStructures.SortedDict, DataStructures.SortedMultiDict, DataStructures.SortedSet}) at ~/.julia/packages/DataStructures/59MD0/src/container_loops.jl:322
=#

@ChrisRackauckas Is this something to do with how the DAE Problem is defined? Can you also add @ccrnn here so that we can debug this together.

issue with regression with SInDy

universal_differential_equations/SEIR_exposure/seir_exposure.jl

Line 228 in dd83689

Ψ = SInDy(X̂[:, 2:end], L̂[2:end], basis, thresholds, opt = opt, maxiter = 10000, normalize = true, denoise = true) # Succeed

Ψ = SInDy(X̂[:, 2:end], L̂[2:end], basis, thresholds, opt = opt, maxiter = 10000, normalize = true, denoise = true) # Succeed

the base i get is different from the one I am suppose to obtain, not sure why.
if I do
print_equations(Ψ)
println(parameters(Ψ))

1 dimensional basis in ["u₁", "u₂", "u₃"]
f_1 = p₁ * u₂ + u₁ * u₂ * p₂ + u₁ ^ 2 * u₂ * p₃
with
p1 , p2, p3= 0.021474172006328535, 0.023860980751708793, 0.026513078226048998

And I am suppose to get :
(u₂ * 0.3065241445838003 + u₁ *0.0011560597253354426)

Application to a different system of ODEs

Hi,

I was trying to apply this method to a different system od ODEs, namely the 2D Hopf normal form defined by:

function Hopf2D(du, u, p, t)
μ, ω, A1, A2 = p
du[1] = μ*u[1] - ω*u[2] + A1*u[1]*(u[1]^(2) + u[2]^(2))
du[2] = μ*u[2] + ω*u[1] + A2*u[2]*(u[1]^(2) + u[2]^(2))
end

where I approximate the third term of the right-hand side of the equations with a UA L(u, p).

Even though it seems that I get good convergence (at least visually) between the (unavailable) ideal data L̄ and L̂ = L(Xₙ, res2.minimizer), the SiNDy algorithm does not output the correct linear combination of u_1 and u_2 (namely A1u_1*(u_1^(2) + u_2^(2) and A2u_2*(u_1^(2) + u_2^(2)).

I was wondering whether the part of the code

opt = SR3()
λ = exp10.(-7:0.1:3)
g(x) = x[1] < 1 ? Inf : norm(x, 2)

is in some way "specific" for the Lotka-Volterra problem.

Or are there any other parts of the code that cannot be applied to some similar systems of ODEs? Am I missing something?

Here's the trainined UA L̂, together with the correct L̄:

The output of the SiNDy is:

2 dimensional basis in ["u₁", "u₂"]
du₁ = u₁ ^ 3 * p₁
du₂ = p₂ * u₂

UDE giving Linear Approximation / NeuralODEMM problems

I also had a similar issue to this for the first part of the SEIR example a few weeks back, though the second part seemed to work.

I copied the example here:
github.com
[ChrisRackauckas/universal_differential_equations/blob/master/LotkaVolterra/scenario_1.jl]

And put in my own function. Everything else is the same except for I didn’t add noise.

U = Lux.Chain(
    Lux.Dense(5,3,rbf), Lux.Dense(3,3, rbf), Lux.Dense(3,3, rbf), Lux.Dense(3,5)
)
# Get the initial parameters and state variables of the model
p, st = Lux.setup(rng, U)


function ude_dynamics!(du, u, p, t, p_true)
    û = U(u, p, st)[1] # Network prediction
    p1, p2, p3, p4, p5, p6, p7, p8, p9 = p
    a, b, c, d, e = u
    da = û[1]
    db = c - p8 - (p3 + p4*p9 + (-p5*e) / a)*p9*a
    dc = p6*((p3 + p4*p9 + (-p5*e) / a)*p9*a + d + p8) - c
    dd = p7*(p2*a + p1*e) - d
    de = (p3 + p4*p9 + (-p5*e) / a)*p9*a + d + p8 - c - e
    du[1] = da; du[2] = db; du[3] = dc; du[4] = dd; du[5] = de
end

#Initial system: 
da = (p3 + p4*p9 + (-p5*e) / a)*p9*a + p8 - c
tspan = (0.0,90.0)
u0 = Float32[86.5,21.62,21.62,86.5,86.5]
p_ = Float32[0.6,0.4,0.635,5,0.01,0.2,0.3,20,0.025]

However, it’s giving me a linear approximation for da, though BFGS finishes way before max iterations.

I have tried with multiple solvers with the above:
Vern7 and KenCarp4 works but gives a linear approximation for da
Tsit5 aborts as unstable when BFGS runs (this solved the original ODE)
Tried a few others, Rosenbrock23 and Rodas4 but these don’t run at all.

I tried running it with 36 in the chain, it finished ADAM then hung all night on BFGS.

The original system was a DAE which I used ModelingToolkit structural_simplify(dae_index_lowering) on then got the full equations, which eliminated the algebraic constraints. I tried using a mass matrix with no constraints and then use NeuralODEMM in case this was the issue. The example I was working from for this is here: Readme · DiffEqFlux.jl (juliahub.com)
However, it's so different from the Lotka Volterra model that I had difficulty translating between the two.

#For original system:
M = [1.0 0 0 0 0
     0 1.0 0 0 0
     0 0 1.0 0 0
     0 0 0 1.0 0
     0 0 0 0 1.0]
#######
# Closure with the known parameter
nn_dynamics!(du,u,p,t) = ude_dynamics!(du,u,p,t,p_)
# Define the problem
prob_nn = NeuralODEMM(nn_dynamics!,X[:, 1], tspan, M, p)

## Function to train the network
# Define a predictor
function predict(θ, X = X[:,1], T = t)
    Array(solve(prob_nn, Rodas5(autodiff=false), u0 = X, p=θ,
                saveat = T,
                abstol=1e-8, reltol=1e-8,
                sensealg = ForwardDiffSensitivity()
                ))
end

I got a really long error, the first part was:
ERROR: MethodError: no method matching init(::NeuralODEMM{typeof(nn_dynamics!), Vector{Float32}, Vector{Bool}, Optimisers.Restructure{typeof(nn_dynamics!), Tuple{}}, Tuple{Float64, Float64}, Matrix{Float64}, Tuple{NamedTuple{(:layer_1, :layer_2, :layer_3, :layer_4), NTuple{4, NamedTuple{(:weight, :bias), Tuple{Matrix{Float32}, Matrix{Float32}}}}}}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ::Rodas5{0, false, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing}; u0=Float32[86.5, 21.62, 21.62, 86.5, 86.5], p=(layer_1 = (weight = Float32[-0.54 ...

I then tried to go back to the original DAE and make it work with the NeuralODEMM example here.
Readme · DiffEqFlux.jl (juliahub.com)
The example works fine, but when I type mine in, it tells me “a” or “p1” etc are not recognised. I put it in u[1], p[1] format, and it tells me u is not recognised. I tried all kinds of brackets in case the problem is that I have more than one constraint. I tried deleting 2 of the constraints. I'm not sure what is different from the example.

function f(du,u,p,t)
    a, b, c, d, e, f, g, h = u
    p1, p2, p3, p4, p5, p6, p7, p8, p9 = p
    du[1] = (p8*f) + p9 - c 
    du[2] =  c + d - (p8*f) - g  
    du[3] =  p6*(g + (p8*f)) - c
    du[4] = (p1*e + p2*a)*p7 - d
    du[5] = g + p8*f - c - e
    du[6] = a*(p3 + p4*p8 - p5*(e/a)) - f
    du[7] = d + p9 - g
    du[8] = a - f - h
end

#...
ndae = NeuralODEMM(dudt2, (u,p,t) -> [a*p3 + p4*p8 - p5*(e/a)) - f],[d + p9 - g],[a - f - h], tspan, M, Rodas5(autodiff=false),saveat=0.1)