Advice on UDE modelling workflow (for future Catalyst implementations)

[TorkelE]

Ok, I think I have managed to get a full working UDE workflow working based on what Chris told me. The advantage of this one (as opposed to the one on the docs) is that it will be much easier to integrate with Catalyst.

I was wondering if someone who know about this stuff better could just confirm that what I am doing makes sense. If it does, I will look at making this working in conjecture with Catalyst (don't work directly as we make some assumption of types that means the neural networks do not work, but this should be changeable).

# Fetch packages.
using Lux, ComponentArrays, ModelingToolkit, ModelingToolkitNeuralNets, OrdinaryDiffEq, Plots, StableRNGs
import ModelingToolkitNeuralNets: lazyconvert
import ModelingToolkit.t_nounits as t
import ModelingToolkit.D_nounits as D
import ModelingToolkit: getu, setp_oop
using OptimizationOptimisers: Adam

# Prepare the true model (X produced as a function of itself, and degraded lineary).
@variables X(t)
@parameters d
prod(x) = (x^2)/(0.5 + x^2)
eq = D(X) ~ prod(X) - d*X
@mtkbuild osys_true = ODESystem([eq], t)

# Generate the synthetic meassurments (samples a simulation and adds some noise) (plot comes later on).
u0_true = [X => 0.5]
tend = 10.0
ps_true = [d => 0.5]
oprob_true = ODEProblem(osys_true, u0_true, tend, ps_true)
sol_true = solve(oprob_true)
sample_t = 0.0:0.1:tend
sample_x = sol_true(sample_t; idxs = X).u .* (0.9 .+ 0.2*rand(length(sample_t)))

# Create the Neural Network.
NN = Lux.Chain(
    Lux.Dense(1 => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => 1, use_bias = false)
)
ps_nn_init, st_nn = Lux.setup(StableRNG(1234), NN)
ca_nn = ComponentArray{Float64}(ps_nn_init)
@parameters p_nn[1:length(ca_nn)] = Vector(ca_nn)
@parameters T_nn::typeof(typeof(ca_nn))=typeof(ca_nn) [tunable = false]

# Create the UDE model.
nn_func(x) = LuxCore.stateless_apply(NN, [x], lazyconvert(T_nn, p_nn))[1]
eq_nn = D(X) ~ nn_func(X) - d*X
@mtkbuild osys_nn = ODESystem([eq_nn], t, [X], [d, p_nn, T_nn])

# Create the optimization function (l2 distance between new simulation and data).
function loss(ps, (prob_base, sample_t, sample_x, set_ps, get_us))
    p = set_ps(prob_base, ps)
    new_prob = remake(prob_base; p)
    new_sol = solve(new_prob; saveat = sample_t)
    SciMLBase.successful_retcode(new_sol) || return Inf
    return sum(abs2.(get_us(new_sol) .- sample_x))
end
of = OptimizationFunction{true}(loss, AutoForwardDiff())

# Prepares the optimisation problem.
ps_init = [1.0; ModelingToolkit.default_values(osys_nn)[osys_nn.p_nn]]
prob_base = ODEProblem(osys_nn, u0_true, tend, [d => ps_init[1]])
set_ps = setp_oop(osys_nn, [d; osys_nn.p_nn])
get_us = getu(osys_nn, X)
opt_prob = OptimizationProblem(of, ps_init, (prob_base, sample_t, sample_x, set_ps, get_us))

# Fits the UDE to the meassured data.
function callback(opt_state, loss)
    (opt_state.iter % 100 == 0) && (@info "step $(opt_state.iter), loss: $loss")
    return false
end
@time opt_sol = solve(opt_prob, Adam(5e-3); maxiters = 5000, callback)

# Simulates the fitted solution.
fitted_ps = set_ps(prob_base, opt_sol.u)
fitted_prob = remake(prob_base, p = fitted_ps)
fitted_sol = solve(fitted_prob)
init_sol = solve(prob_base)

# Evaluate the fit (plots true solution, data, and the pre/post fitted NN model simulation).
plot(sol_true; label = "True solution", lw = 5)
plot!(sample_t, sample_x; label = "Meassured values", seriestype = :scatter, color = 1, ms = 6, markeralpha = 0.6)
plot!(fitted_sol, lw = 4, la = 0.9, linestyle = :dash, label = "Trained solution", color = :blue)
plot!(init_sol, lw = 4, la = 0.9, linestyle = :dot, label = "Pre-training solution", color = :royalblue4)

# Plots the fitted/true functions.
X_grid = minimum(sample_x):0.1:maximum(sample_x)
prod_true = prod.(X_grid)
get_nn_val(x) = LuxCore.stateless_apply(NN, [x], convert(fitted_sol.ps[osys_nn.T_nn], fitted_sol.ps[osys_nn.p_nn]))[1]
prod_nn = get_nn_val.(X_grid)
plot(X_grid, prod_true; label = "True production rate", lw = 8, la = 0.8)
p2 = plot!(X_grid, prod_nn; label = "Fitted production rate", lw = 8, la = 0.8, xlimit = (X_grid[1], X_grid[end]))

(in this case we can fit well, but the addition of d mean we don't necessarily recover the true function, but that is besides the point here)

[TorkelE]

I also have a specific question. When I run

ps_nn_init, st_nn = Lux.setup(StableRNG(1234), NN)

I get a variable st_nn (corresponding to the initial states?). Here I do not really need it for anything. However, is there some cases where this one might be useful, or can I just generally discard it? (and if so, can I use initialparameters instead of setup?)

[SebastianM-C]

The funny thing about this package is that the main aspect is not actually implemented here. If you want to use the so called "low level interface", you don't really need to depend on this package. As you've seen above, the only thing that you needed was that lazyconvert. I have to review things again, but I'm not entirely sure if it's absolutely needed on more recent Symbolics versions. In any case, I wouldn't depend on this function for that one line, you can just copy that.

@ChrisRackauckas Since more people are asking about this, should we document this? It would be the extreme version of low dependency usage, you don't even use this package 😆

On the implementation side of the above, yes you should use initialparameters, we also do it

ModelingToolkitNeuralNets.jl/src/ModelingToolkitNeuralNets.jl

Line 28 in 7466069

init_params = Lux.initialparameters(rng, chain),

. It would never be needed because we specifically call LuxCore.stateless_apply, which as the name suggests, is without state (which is what st_nn) is.

ps_init = [1.0; ModelingToolkit.default_values(osys_nn)[osys_nn.p_nn]]

default_values is from SII

julia> @which default_values
SymbolicIndexingInterface

(same for setp_oop and getp), so I recommend depending on SII directly.

One thing that I've been wondering implementation wise is if we should also make the neural network a parameter, as that would make it easier to retrieve it and might also make the codegen a bit nicer. In any case, the relevant part for this discussion is that if we make some updates to how the NN is represented you won't get them (which can be good or bad depending on the situation 😅 ). Hopefully we won't need many updates to this package, so I don't think this is a significant concern.

[TorkelE]

Sounds good, thanks!

My dream scenario would be for neural networks to work something like callable parameters: https://docs.sciml.ai/ModelingToolkit/stable/tutorials/callable_params/

I.e. I first declare that a parameter represents a neural network like:

T = typeof(...) # Some type
@parameters (NN::T)(..)

Then my model:

@variables X[1:1]
eq = D(X) ~ NN(X) - X
@mtkbuild osys = ODESystem([eq], t)

Finally, I can now designate my neural network structure after the model, and provide it (and corresponding information as a parameter when I simualte the model:

nn_struct = Lux.Chain(
    Lux.Dense(1 => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => 1, use_bias = false)
)
nn_in = ... # Something.
ODEProblem(osys, [X => [1.0], 10.0, [NN => nn_in])

Now I can see several potential problems here, i.e. the neural networks have parameters, here these are not direct parameters of the model. However, if this can be solved, I think it would make for a very neat interface.

[ChrisRackauckas]

We could make a callable type in this package that does this.

[ChrisRackauckas]

You would need to have the parameters of the NN as separate in order to differentiate them, so NN(x,p) would be required.

[TorkelE]

Something like NN(x,p), sounds perfect.

[TorkelE]

I have also managed to get a prototype in Catalyst working based on this workflow, it is quite neat actually. Will play around with it personally, and then, depending on any updates here I will finalise it with docs later this year (ideally before summer though).

[ChrisRackauckas]

@SebastianM-C can you make a nice callable type that automates this a bit, and a tutorial on doing this directly? I think it could be useful to many folks.