Symbolic UDE Creation

This tutorial will demonstrate a simple interface for symbolic declaration neural networks that can be directly added to ModelingToolkit.jl-declared ODE models to create UDEs. The primarily functionality we show is the SymbolicNeuralNetwork function, however, we will show how it can be incorporated into a full workflow. For our example we will use a simple self-activation loop model, however, it can be easily generalised to more model types.

Ground truth model and synthetic data generation

First we create the ground-truth model using ModelingToolkit. In it, Y activates X at the rate v * (Y^n) / (K^n + Y^n). Later on, we will attempt to learn this rate using a neural network. Both variables decay at constant rates that scales with the parameter d.

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
@variables X(t) Y(t)
@parameters v=1.0 K=1.0 n=1.0 d=1.0 # Sets unused default values for all parameters (but vaguely useful as potential optimization initial conditions).
eqs = [D(X) ~ v * (Y^n) / (K^n + Y^n) - d*X
       D(Y) ~ X - d*Y]
@mtkcompile xy_model = System(eqs, t)

\[ \begin{align} \frac{\mathrm{d} Y\left( t \right)}{\mathrm{d}t} &= X\left( t \right) - d Y\left( t \right) \\ \frac{\mathrm{d} X\left( t \right)}{\mathrm{d}t} &= \frac{\left( Y\left( t \right) \right)^{n} v}{\left( Y\left( t \right) \right)^{n} + K^{n}} - d X\left( t \right) \end{align} \]

Next, we simulate our model for a true parameter set (which we wish to recover).

using OrdinaryDiffEqTsit5, Plots
u0 = [X => 2.0, Y => 0.1]
ps_true = [v => 1.1, K => 2.0, n => 3.0, d => 0.5]
sim_cond = [u0; ps_true]
tend = 45.0
oprob_true = ODEProblem(xy_model, sim_cond, (0.0, tend))
sol_true = solve(oprob_true, Tsit5())
plot(sol_true; lw = 6, idxs = [X, Y])

Finally, we generate noisy measured samples from both X and Y (to which we will fit the UDE).

sample_t = range(0.0, tend; length = 20)
sample_X = [(0.8 + 0.4rand()) * X_sample for X_sample in sol_true(sample_t; idxs = X)]
sample_Y = [(0.8 + 0.4rand()) * Y_sample for Y_sample in sol_true(sample_t; idxs = Y)]
plot!(sample_t, sample_X, seriestype = :scatter,
    label = "X (data)", color = 1, ms = 6, alpha = 0.7)
plot!(sample_t, sample_Y, seriestype = :scatter,
    label = "Y (data)", color = 2, ms = 6, alpha = 0.7)

UDE declaration and training

First, we use Lux.jl to declare the neural network we wish to use for our UDE. For this case, we can use a fairly small network. We use softplus throughout the network we ensure that the fitted UDE function is positive (for our application this is the case, however, it might not always be true).

using Lux
nn_arch = Lux.Chain(
    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 1, Lux.softplus, use_bias = false)
)

Chain(
    layer_1 = Dense(1 => 3, softplus, use_bias=false),  # 3 parameters
    layer_2 = Dense(3 => 3, softplus, use_bias=false),  # 9 parameters
    layer_3 = Dense(3 => 1, softplus, use_bias=false),  # 3 parameters
)         # Total: 15 parameters,
          #        plus 0 states.

Next, we can use ModelingToolkitNeuralNets.jl to turn our neural network to a Symbolic neural network representation (which can later be inserted into an ModelingToolkit model).

using ModelingToolkitNeuralNets
sym_nn,
θ = SymbolicNeuralNetwork(; nn_p_name = :θ, chain = nn_arch, n_input = 1, n_output = 1)
sym_nn_func(x) = sym_nn([x], θ)[1]

sym_nn_func (generic function with 1 method)

Now we can create our UDE. We replace the (from now on unknown) function v * (Y^n) / (K^n + Y^n) with our symbolic neural network (which we let be a function of the variable Y only).

eqs_ude = [D(X) ~ sym_nn_func(Y) - d*X
           D(Y) ~ X - d*Y]
@mtkcompile xy_model_ude = System(eqs_ude, t)

\[ \begin{align} \frac{\mathrm{d} Y\left( t \right)}{\mathrm{d}t} &= X\left( t \right) - d Y\left( t \right) \\ \frac{\mathrm{d} X\left( t \right)}{\mathrm{d}t} &= \mathtt{NN}\_{1}\left( \left[ \begin{array}{c} Y\left( t \right) \\ \end{array} \right], \theta \right) - d X\left( t \right) \end{align} \]

We can now fit our UDE model (including the neural network and the parameter d) to the data. First, we define a loss function which compares the UDE's simulation to the data.

function loss(ps, (oprob_base, set_ps, sample_t, sample_X, sample_Y))
    p = set_ps(oprob_base, ps)
    new_oprob = remake(oprob_base; p)
    new_osol = solve(new_oprob, Tsit5(); saveat = sample_t, verbose = false, maxiters = 10000)
    SciMLBase.successful_retcode(new_osol) || return Inf # Simulation failed -> Inf loss.
    x_error = sum((x_sim - x_data)^2 for (x_sim, x_data) in zip(new_osol[X], sample_X))
    y_error = sum((y_sim - y_data)^2 for (y_sim, y_data) in zip(new_osol[Y], sample_Y))
    return x_error + y_error
end

loss (generic function with 1 method)

Next, we use Optimization.jl to create an OptimizationProblem. This uses a similar syntax to normal parameter inference workflows, however, we need to add the entire neural network parameterisation to the optimization parameter vector.

using Optimization
oprob_base = ODEProblem(xy_model_ude, u0, (0.0, tend))
set_ps = ModelingToolkit.setp_oop(oprob_base, [d, θ...])
loss_params = (oprob_base, set_ps, sample_t, sample_X, sample_Y)
ps_init = oprob_base.ps[[d, θ...]]
of = OptimizationFunction{true}(loss, AutoForwardDiff())
opt_prob = OptimizationProblem(of, ps_init, loss_params)

OptimizationProblem. In-place: true
u0: 16-element Vector{Float64}:
  1.0
 -0.04929668828845024
 -0.3266667425632477
 -1.6716315746307373
 -0.5999714136123657
  0.7242816686630249
 -0.458548903465271
 -0.8280583620071411
 -0.38509929180145264
  0.32322537899017334
 -0.32623517513275146
 -0.7673453092575073
  0.7302734851837158
 -0.16844713687896729
  0.4040088653564453
  0.28568029403686523

Finally, we can fit the UDE to our data. We will use the Adam optimizer.

import OptimizationOptimisers: Adam
@time opt_sol = solve(opt_prob, Adam(0.01); maxiters = 10000)

retcode: Default
u: 16-element Vector{Float64}:
  0.48580829783419616
  0.7894319969826042
 -0.44928532760559153
 -0.5496463229369764
 -1.2747190824583556
 -0.20318234480926148
 -0.28766868759800723
  3.2800019765463544
 -0.05604357327669147
  0.5363989270977549
  4.704203691790537
 -1.2887838499116326
  0.20727768725643742
 -1.2321266128534931
  0.2911785878058523
  0.8508099849487802

By plotting a simulation from our fitted UDE, we can confirm that it can reproduce the ground-truth model.

oprob_fitted = remake(oprob_base; p = set_ps(oprob_base, opt_sol.u))
sol_fitted = solve(oprob_fitted, Tsit5())
plot!(sol_true; lw = 4, la = 0.7, linestyle = :dash, idxs = [X, Y], color = [:blue :red],
    label = ["X (UDE)" "Y (UDE)"])