Silent release of UDEs with ModelingToolkit?

Specific Domains Modelling & Simulations
, , ,
1

I’ve seen several discussions about writing symbolic UDE models with ModelingToolkit, and I think it is fair to say it has been an anticipated feature for some years:

I was playing around with callable parameters and realized that you can in fact model and train UDEs using this functionality. I rewrote the example from the SciML book shown below.

using Pkg
Pkg.activate(@__DIR__)
using ModelingToolkit, DifferentialEquations, Lux, Optimization, OptimizationOptimisers, Plots, Random, ComponentArrays
using ModelingToolkit: t_nounits as t, D_nounits as D


u0 = Float32[0.8; 0.8]
tspan = (0.0f0,25.0f0)
ts = collect(0.0f0:0.1f0:25.0f0)


#Neural Network
ann = Chain(Dense(2,10,tanh), Dense(10,1))
p_nn, st = Lux.setup(Random.Xoshiro(0), ann)
nn_func(p_nn) = (x, y) -> first(first(ann([x, y], p_nn, st)))

p_nn = ComponentArray(p_nn)
p_ode = Float32[-2.0,1.1]
p3 = ComponentArray(nn=p_nn, ode=p_ode)

#UDE toy example
@variables x(t) y(t)
params = @parameters (nn::Function)(..) p1 = -2.0 p2 = 1.1 
eqs = [
    D(x) ~ nn(x,y)
    D(y) ~ p1*y + p2*x
]

@mtkcompile sys = System(eqs, t)
parameter_map(p) = [
  nn => nn_func(p.nn),
  p1 => p.ode[1],
  p2 => p.ode[2],
]

prob = ODEProblem(sys, [
  x => u0[1],
  y => u0[2],
  parameter_map(p3)...,
], tspan)
sol = solve(prob, Tsit5(), saveat=ts, abstol=1f-6, reltol=1f-6)

# Optimization
function predict(p)
  trial_prob = remake(prob; p=parameter_map(p))
  Array(solve(trial_prob, Tsit5(), saveat=ts, abstol=1f-6, reltol=1f-6))
end
loss(p, _) = sum(abs2, x - 1.0 for x in predict(p))
loss(p3, nothing)

iter = 0
cb = function (state, l)
  global iter += 1
  if iter % 50 == 0
    display(l)
    display(plot(solve(remake(prob; p=parameter_map(state.u)), Tsit5(), saveat=0.1), ylim=(0,6)))
  end
  return false
end

# Display the ODE with the current parameter values.
display(plot(sol, ylim=(0,6)))

optf = Optimization.OptimizationFunction(loss, Optimization.AutoForwardDiff())
optprob = Optimization.OptimizationProblem(optf, p3)
res = Optimization.solve(optprob, OptimizationOptimisers.Adam(0.01), callback=cb, maxiters=300)

Is this documented or even an intended feature? How is no one talking about this? Will there be support for reverse mode differentiation?

2

Just found Home · ModelingToolkitNeuralNets.jl

3

Yeah we put out things in the docs and such, but the old threads will always be there. Maybe it would be good to go back to them and mention that it does work. We use it extensively in Dyad.