MTK system type instability with physics-informed neural networks for parameter estimation for ODEs

[oameye]

Describe the bug 🐞

When trying out the "Parameter Estimation with Physics-Informed Neural Networks for ODEs" example for a ModelingToolkit system some type instability happens such that the initial condition type is not known?
Minimal Reproducible Example 👇

using ModelingToolkit, NeuralPDE, OrdinaryDiffEq, Lux, Random
using OptimizationOptimJL, LineSearches
using ModelingToolkit: t_nounits as t, D_nounits as D

vars_sym = @variables aᵣ(t) aᵢ(t)
param_new = @parameters Δ U G κ

eqs = [
    D(aᵣ) ~ Δ*aᵢ - G*aᵢ - κ*aᵣ/2 - U*(aᵣ^2+aᵢ^2)*aᵢ/2,
    D(aᵢ) ~ - Δ*aᵣ - G*aᵣ - κ*aᵢ/2 + U*(aᵣ^2+aᵢ^2)*aᵣ/2
    ]
@mtkbuild sys_man = ODESystem(eqs, t)

tspan = (0.0, 6_000.0)
u0 = [0.1, 0.2]
true_p = [1.0, 0.002, 0.0, 0.02]
param_val = (U, κ, Δ, G) .=> true_p
prob = ODEProblem{false}(sys_man, u0, tspan, param_val)

sol_data = solve(prob, DP5(), saveat=0.1);

t_ = sol_data.t
u_ = reduce(hcat, sol_data.u)


n = 15
chain = Chain(Dense(1, n, σ), Dense(n, n, σ), Dense(n, n, σ), Dense(n, 2))
ps, st = Lux.setup(Random.default_rng(), chain) |> f64

opt = LBFGS()
additional_loss(phi, θ) = sum(abs2, phi(t_, θ) .- u_) / size(u_, 2)
alg = NNODE(chain, opt, ps; param_estim = true, additional_loss)

sol = solve(prob, alg, verbose = true, abstol = 1e-8, maxiters = 5000, saveat = t_)

Error & Stacktrace ⚠️

ERROR: MethodError: no method matching zero(::Type{Any})

Closest candidates are:
  zero(::Type{Union{Missing, T}}) where T
   @ Base missing.jl:105
  zero(::Type{Union{}}, Any...)
   @ Base number.jl:310
  zero(::Type{Dates.DateTime})
   @ Dates ~/.julia/juliaup/julia-1.10.9+0.x64.linux.gnu/share/julia/stdlib/v1.10/Dates/src/types.jl:438
  ...

Stacktrace:
  [1] zero(::Type{Any})
    @ Base ./missing.jl:106
  [2] __solve(cache::OptimizationCache{…})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/e3bUa/src/OptimizationOptimJL.jl:200
  [3] solve!(cache::OptimizationCache{…})
    @ SciMLBase ~/.julia/packages/SciMLBase/Ub9ub/src/solve.jl:187
  [4] solve(::OptimizationProblem{…}, ::LBFGS{…}; kwargs::@Kwargs{…})
    @ SciMLBase ~/.julia/packages/SciMLBase/Ub9ub/src/solve.jl:95
  [5] __solve(::ODEProblem{…}, ::NNODE{…}; dt::Nothing, timeseries_errors::Bool, save_everystep::Bool, adaptive::Bool, abstol::Float64, reltol::Float32, verbose::Bool, saveat::Vector{…}, maxiters::Int64, tstops::Nothing)
    @ NeuralPDE ~/.julia/packages/NeuralPDE/nkWKK/src/ode_solve.jl:384
  [6] __solve
    @ ~/.julia/packages/NeuralPDE/nkWKK/src/ode_solve.jl:293 [inlined]
  [7] solve_call(_prob::ODEProblem{…}, args::NNODE{…}; merge_callbacks::Bool, kwargshandle::Nothing, kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/y5GOy/src/solve.jl:635
  [8] solve_up(prob::ODEProblem{…}, sensealg::Nothing, u0::Vector{…}, p::MTKParameters{…}, args::NNODE{…}; kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/y5GOy/src/solve.jl:1126
  [9] solve_up
    @ ~/.julia/packages/DiffEqBase/y5GOy/src/solve.jl:1120 [inlined]
 [10] solve(prob::ODEProblem{…}, args::NNODE{…}; sensealg::Nothing, u0::Nothing, p::Nothing, wrap::Val{…}, kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/y5GOy/src/solve.jl:1057
Some type information was truncated. Use `show(err)` to see complete types.

Additional Information

When instead defining the ode manually, it works.

function ode(u,p,t)
    du = similar(u)
    du[1] = p[1]*u[2] - p[3]*u[2] - p[4]*u[1]/2 - p[2]*(u[1]^2+u[2]^2)*u[2]/2
    du[2] = -p[1]*u[1] - p[3]*u[1] - p[4]*u[2]/2 + p[2]*(u[1]^2+u[2]^2)*u[1]/2
    return du
end

prob = ODEProblem{false}(ode, u0, tspan, true_p)

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

• Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)

Status `~/Documents/GitHub/Quantum_Reconstruction/Project.toml`
  [634d3b9d] DrWatson v2.18.0
  [587475ba] Flux v0.16.3
  [e13b9ff6] HarmonicBalance v0.14.3
  [d3d80556] LineSearches v7.3.0
  [b2108857] Lux v1.10.0
  [961ee093] ModelingToolkit v9.66.0
  [315f7962] NeuralPDE v5.18.1
  [7f7a1694] Optimization v4.1.1
  [36348300] OptimizationOptimJL v0.4.1
  [42dfb2eb] OptimizationOptimisers v0.3.7
  [1dea7af3] OrdinaryDiffEq v6.92.0
  [f0f68f2c] PlotlyJS v0.18.15
  [91a5bcdd] Plots v1.40.11
  [35bcea6d] QuantumCumulants v0.3.7
  [90137ffa] StaticArrays v1.9.13
  [0c5d862f] Symbolics v6.31.0
  [9a3f8284] Random
  [10745b16] Statistics v1.10.0

• Output of versioninfo()

Julia Version 1.10.9
Commit 5595d20a287 (2025-03-10 12:51 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 12 × AMD Ryzen 5 5600X 6-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, znver3)
Threads: 10 default, 0 interactive, 5 GC (on 12 virtual cores)
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 10

[ChrisRackauckas]

The parameters don't seem right here. MTK generates equations with an MTKParameters object. θ would be the concatenation of the ODE parameters and the NN parameters, which I think it's messing up on. We'd need to teach this to use just the tunables.

When trying out the "Parameter Estimation with Physics-Informed Neural Networks for ODEs" example for a ModelingToolkit system some type instability happens such that the initial condition type is not known?

Out of curiosity, why are you doing it the PINN way? While it can be interesting, it's pretty much guaranteed to be less robust and 10000x slower than using the differentiable simulation approach. We should make it work, but I don't see a real use case for this kind of thing.

[oameye]

Out of curiosity, why are you doing it the PINN way? While it can be interesting, it's pretty much guaranteed to be less robust and 10000x slower than using the differentiable simulation approach. We should make it work, but I don't see a real use case for this kind of thing.

The eventually goal is to fit many trajectories (resonators ringdowns) from different initial conditions all converging to the same attractor/steady-state/fixed point. We want to learn the vector field of the device. The experimental data is quite noisy and we want to compare how PINNs performed compared to using DiffEqParamEstim.jl. With PINNs the vector field du can deviate from the model allowing us to also capture some deviations from the model which are likely present in the experiment.

[ChrisRackauckas]

With PINNs the vector field du can deviate from the model allowing us to also capture some deviations from the model which are likely present in the experiment.

And are you testing UDEs in comparison? That seems more UDEs than DiffEqParamEstim

[oameye]

And are you testing UDEs in comparison? That seems more UDEs than DiffEqParamEstim

Yes I am doing a comparison with the DiffEqFlux UDEs examples. With PINNs it seemed easier to train against trajectories of many different ICs. But yeah as you already highlighted UDEs are trained much faster.