How to define parameters of a model built by Lux.jl in ModelingToolkit.jl and perform optimization or computation?

I’ve just started using ModelingToolkit.jl recently, but I would like to describe and compute an ODE problem with a neural network model using ModelingToolkit.jl.

I would like to start by testing the feasibility of this with a simple demo. Here is my code:

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using Optimization, OptimizationOptimisers
using Lux
using OrdinaryDiffEq
using Random

model = Lux.Chain(Lux.Dense(3, 16, tanh), Lux.Dense(16, 1, leakyrelu), name=:model)
rng = MersenneTwister()
Random.seed!(rng, 42)
ps, st = Lux.setup(rng, model)
p, re = destructure(ps)
Lux.apply(model,[1,1,1],re(p),st)
@variables x(t)
@parameters pvec[1:Lux.parameterlength(model)]
temp_func(t) = (model([sin(t), cos(t), tan(t)], re(pvec), st)[1])[1]
eqs = [D(x) ~ temp_func(t) * (x)]
u0 = [x => 1.0]
pset = [
pvec => p,
]
sys = ODESystem(eqs, t, name=:sys)
simple_sys = structural_simplify(sys)
prob = ODEProblem(simple_sys, u0, (0.0, 10.0), pset)
sol = solve(prob, Tsit5())

However, I encountered an error when executing structural_simplify(sys)

ERROR: MethodError: no method matching concrete_symtype(::Symbolics.ArrayOp{Vector{Real}})

Closest candidates are:
  concrete_symtype(::Type{Integer})
   @ ModelingToolkit D:\Julia\Julia-1.10.0\packages\packages\ModelingToolkit\hPXD5\src\systems\index_cache.jl:300
  concrete_symtype(::Type{Real})
   @ ModelingToolkit D:\Julia\Julia-1.10.0\packages\packages\ModelingToolkit\hPXD5\src\systems\index_cache.jl:299
  concrete_symtype(::Type{A}) where {T, N, A<:Array{T, N}}
   @ ModelingToolkit D:\Julia\Julia-1.10.0\packages\packages\ModelingToolkit\hPXD5\src\systems\index_cache.jl:301
  ...

In addition, I also constructed the ode problem through another method, as shown below:

model = Lux.Chain(Lux.Dense(3, 16, tanh), Lux.Dense(16, 1, leakyrelu), name=:model)
rng = MersenneTwister()
Random.seed!(rng, 42)
ps, st = Lux.setup(rng, model)

@variables x(t)
@parameters layer1_w[1:16, 1:3], layer1_b[1:16, 1:1]
@parameters layer2_w[1:16, 1:16], layer2_b[1:16, 1:1]

temp_func(t) = (model([sin(t), cos(t), tan(t)], (layer_1=(weight=layer1_w, bias=layer1_b), layer_2=(weight=layer2_w, bias=layer2_b)), st)[1])[1]

eqs = [D(x) ~ temp_func(t)]
u0 = [x => 1.0]
pset = [
    layer1_w => ps[:layer_1][:weight],
    layer1_b => ps[:layer_1][:bias],
    layer2_w => ps[:layer_2][:weight],
    layer2_b => ps[:layer_2][:bias],
]
sys = ODESystem(eqs, t, name=:sys)
simple_sys = structural_simplify(sys)
prob = ODEProblem(simple_sys, u0, (0.0, 10.0), pset)
sol = solve(prob, Tsit5())

But this will lead to the following problem:

ERROR: MethodError: no method matching *(::Int64, ::SymbolicUtils.BasicSymbolic{Matrix{Real}})

Closest candidates are:
  *(::Any, ::Any, ::Any, ::Any...)
   @ Base operators.jl:587
  *(::Integer, ::ForwardDiff.Dual{Ty}) where Ty
   @ ForwardDiff D:\Julia\Julia-1.10.0\packages\packages\ForwardDiff\PcZ48\src\dual.jl:145
  *(::Integer, ::ReverseDiff.TrackedReal)
   @ ReverseDiff D:\Julia\Julia-1.10.0\packages\packages\ReverseDiff\UJhiD\src\derivatives\scalars.jl:18

I hope someone can help me solve my problem.

Just hold your horses. We only shared the library for putting neural network components into MTK yesterday

We’re still working out the last details.

1 Like

Thank you very much. Looking forward to the library.