diff --git a/docs/make.jl b/docs/make.jl index cbbeb87..0889551 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -21,7 +21,8 @@ makedocs(; pages = [ "Home" => "index.md", "Tutorials" => ["NeuralNetworkBlock" => "nnblock.md" - "Friction Model" => "friction.md"], + "Friction Model" => "friction.md" + "Symbolic UDE Creation" => "symbolic_ude_tutorial.md"], "API" => "api.md" ] ) diff --git a/docs/src/symbolic_ude_tutorial.md b/docs/src/symbolic_ude_tutorial.md new file mode 100644 index 0000000..fd16016 --- /dev/null +++ b/docs/src/symbolic_ude_tutorial.md @@ -0,0 +1,114 @@ +# Symbolic UDE Creation + +This tutorial will demonstrate a simple interface for symbolic declaration neural networks that can be directly added to [ModelingToolkit.jl](https://github.com/SciML/ModelingToolkit.jl)-declared ODE models to create UDEs. The primarily functionality we show is the [`SymbolicNeuralNetwork`](@ref) 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`. + +```@example symbolic_ude +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) +``` + +Next, we simulate our model for a true parameter set (which we wish to recover). + +```@example symbolic_ude +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). + +```@example symbolic_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](https://github.com/LuxDL/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). + +```@example symbolic_ude +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) +) +``` + +Next, we can use [ModelingToolkitNeuralNets.jl](https://github.com/SciML/ModelingToolkitNeuralNets.jl) to turn our neural network to a Symbolic neural network representation (which can later be inserted into an ModelingToolkit model). + +```@example symbolic_ude +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] +``` + +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). + +```@example symbolic_ude +eqs_ude = [D(X) ~ sym_nn_func(Y) - d*X + D(Y) ~ X - d*Y] +@mtkcompile xy_model_ude = System(eqs_ude, t) +``` + +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. + +```@example symbolic_ude +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 +``` + +Next, we use [Optimization.jl](https://github.com/SciML/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. + +```@example symbolic_ude +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) +``` + +Finally, we can fit the UDE to our data. We will use the Adam optimizer. + +```@example symbolic_ude +import OptimizationOptimisers: Adam +@time opt_sol = solve(opt_prob, Adam(0.01); maxiters = 10000) +``` + +By plotting a simulation from our fitted UDE, we can confirm that it can reproduce the ground-truth model. + +```@example symbolic_ude +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)"]) +```