Help with modifying SciML Missing Physics UDE example

Hello all,

I’ve made some minor modifications to the code from Automatically Discover Missing Physics by Embedding Machine Learning into Differential Equations · Overview of Julia's SciML

My goal is to expose unknown parameters and known parameters in the UDE so that unknown parameters can be trained and known parameters can be utilized.

The top level error is: “LoadError: type NamedTuple has no field mech_params”

  1. Am I taking an appropriate direction to implement my goal, or is there a better paradigm in the SciML framework?
  2. If this is a reasonable direction, in what directions should I make changes to the code?

Thank you in advance for your assistance.

A reduced working example follows. The stacktrace is fairly long, so I haven’t included it:

# SciML Tools
using OrdinaryDiffEq, ModelingToolkit, DataDrivenDiffEq, SciMLSensitivity, DataDrivenSparse
using Optimization, OptimizationOptimisers, OptimizationOptimJL

# Standard Libraries
using LinearAlgebra, Statistics

# External Libraries
using ComponentArrays, Lux, Zygote, Plots, StableRNGs
gr()

# Set a random seed for reproducible behaviour
rng = StableRNG(1111)

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.0, 5.0)
u0 = 5.0f0 * rand(rng, 2)
p_ = [1.3, 0.9, 0.8, 1.8]
prob = ODEProblem(lotka!, u0, tspan, p_)
solution = solve(prob, Vern7(), abstol = 1e-12, reltol = 1e-12, saveat = 0.25)

# Add noise in terms of the mean
X = Array(solution)
t = solution.t

x̄ = mean(X, dims = 2)
noise_magnitude = 5e-3
Xₙ = X .+ (noise_magnitude * x̄) .* randn(rng, eltype(X), size(X))

plot(solution, alpha = 0.75, color = :black, label = ["True Data" nothing])
scatter!(t, transpose(Xₙ), color = :red, label = ["Noisy Data" nothing])

rbf(x) = exp.(-(x .^ 2))

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

p_full = ComponentVector(dd_params=p, mech_params=[1.0,],known_params=[1.8,])

# Define the hybrid model

function ude_dynamics!(du, u, p, t)

    p_mech = p.mech_params
    p_true = p.known_params
    p_nn = p.dd_params

    û = U(u, p_nn, _st)[1] # Network prediction
    du[1] = p_mech[1] * u[1] + û[1]
    du[2] = -p_true[4] * u[2] + û[2]
end

prob_nn = ODEProblem(ude_dynamics!, Xₙ[:, 1], tspan, p_full)

function predict(θ, X = Xₙ[:, 1], T = t)
    _prob = remake(prob_nn, u0 = X, tspan = (T[1], T[end]), p = θ)
    Array(solve(_prob, Vern7(), saveat = T,
                abstol = 1e-6, reltol = 1e-6,
                sensealg=QuadratureAdjoint(autojacvec=ReverseDiffVJP(true))))
end

function loss(θ)
    X̂ = predict(θ)
    mean(abs2, Xₙ .- X̂)
end

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

optprob = Optimization.OptimizationProblem(optf, p_full)
res1 = Optimization.solve(optprob, ADAM(), maxiters = 5000)

optprob2 = Optimization.OptimizationProblem(optf, res1.u)
res2 = Optimization.solve(optprob2, Optim.LBFGS(), maxiters = 1000)

Did you try converting p to a ComponentArray first?

That, with a few other minor changes, fixed my problem. Thanks.