diff --git a/Project.toml b/Project.toml index 599404d..09c07f0 100644 --- a/Project.toml +++ b/Project.toml @@ -20,7 +20,7 @@ JET = "0.8" Lux = "0.5.32" LuxCore = "0.1.14" ModelingToolkit = "9.9" -ModelingToolkitStandardLibrary = "2.6" +ModelingToolkitStandardLibrary = "2.7" Optimization = "3.24" OptimizationOptimisers = "0.2.1" OrdinaryDiffEq = "6.74" diff --git a/docs/Project.toml b/docs/Project.toml index 79f849e..92889db 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,6 +1,27 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" +ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739" +Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" +OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" ModelingToolkitNeuralNets = "f162e290-f571-43a6-83d9-22ecc16da15f" [compat] Documenter = "1.3" +Lux = "0.5.32" +ModelingToolkit = "9.9" +ModelingToolkitStandardLibrary = "2.7" +Optimization = "3.24" +OptimizationOptimisers = "0.2.1" +OrdinaryDiffEq = "6.74" +Plots = "1" +SciMLStructures = "1.1.0" +StableRNGs = "1" +SymbolicIndexingInterface = "0.3.15" +ModelingToolkitNeuralNets = "1" diff --git a/docs/make.jl b/docs/make.jl index 54367ee..8575c56 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,22 +1,29 @@ using ModelingToolkitNeuralNets using Documenter +ENV["GKSwstype"] = "100" +ENV["JULIA_DEBUG"] = "Documenter" cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml", force = true) cp("./docs/Project.toml", "./docs/src/assets/Project.toml", force = true) -DocMeta.setdocmeta!(ModelingToolkitNeuralNets, :DocTestSetup, :(using ModelingToolkitNeuralNets); recursive = true) +DocMeta.setdocmeta!(ModelingToolkitNeuralNets, :DocTestSetup, + :(using ModelingToolkitNeuralNets); recursive = true) makedocs(; modules = [ModelingToolkitNeuralNets], authors = "Sebastian Micluța-Câmpeanu and contributors", sitename = "ModelingToolkitNeuralNets.jl", - format = Documenter.HTML(; - canonical = "https://SciML.github.io/ModelingToolkitNeuralNets.jl", - edit_link = "main", - assets = String[] - ), + format = Documenter.HTML(assets = ["assets/favicon.ico"], + canonical = "https://docs.sciml.ai/ModelingToolkitNeuralNets.jl/stable/"), + clean = true, + doctest = false, + linkcheck = true, pages = [ - "Home" => "index.md" + "Home" => "index.md", + "Tutorials" => [ + "Friction Model" => "friction.md" + ], + "API" => "api.md" ] ) diff --git a/docs/src/api.md b/docs/src/api.md new file mode 100644 index 0000000..28b49d5 --- /dev/null +++ b/docs/src/api.md @@ -0,0 +1,5 @@ +# API + +```@docs +NeuralNetworkBlock +``` diff --git a/docs/src/assets/favicon.ico b/docs/src/assets/favicon.ico new file mode 100644 index 0000000..3c6bd47 Binary files /dev/null and b/docs/src/assets/favicon.ico differ diff --git a/docs/src/assets/logo.png b/docs/src/assets/logo.png new file mode 100644 index 0000000..6f4c3e2 Binary files /dev/null and b/docs/src/assets/logo.png differ diff --git a/docs/src/friction.md b/docs/src/friction.md new file mode 100644 index 0000000..842d6b7 --- /dev/null +++ b/docs/src/friction.md @@ -0,0 +1,179 @@ +# Modeling Non Linear Friction Model using UDEs + +Friction between moving bodies is not trivial to model. There have been idealised linear models which are not always useful in complicated systems. There have been many theories and non linear models which we can use, but they are not perfect. The aim of this tutorial to use Universal Differential Equations to showcase how we can embed a neural network to learn an unknown non linear friction model. + +## Julia Environment + +First, lets import the required packages. + +```@example friction +using ModelingToolkitNeuralNets +using ModelingToolkit +import ModelingToolkit.t_nounits as t +import ModelingToolkit.D_nounits as Dt +using ModelingToolkitStandardLibrary.Blocks +using OrdinaryDiffEq +using Optimization +using OptimizationOptimisers: Adam +using SciMLStructures +using SciMLStructures: Tunable +using SymbolicIndexingInterface +using StableRNGs +using Lux +using Plots +using Test #hide +``` + +## Problem Setup + +Let's use the friction model presented in https://www.mathworks.com/help/simscape/ref/translationalfriction.html for generating data. + +```@example friction +Fbrk = 100.0 +vbrk = 10.0 +Fc = 80.0 +vst = vbrk / 10 +vcol = vbrk * sqrt(2) +function friction(v) + sqrt(2 * MathConstants.e) * (Fbrk - Fc) * exp(-(v / vst)^2) * (v / vst) + + Fc * tanh(v / vcol) +end +``` + +Next, we define the model - an object sliding in 1D plane with a constant force `Fu` acting on it and friction force opposing the motion. + +```@example friction +function friction_true() + @variables y(t) = 0.0 + @constants Fu = 120.0 + eqs = [ + Dt(y) ~ Fu - friction(y) + ] + return ODESystem(eqs, t, name = :friction_true) +end +``` + +Now that we have defined the model, we will simulate it from 0 to 0.1 seconds. + +```@example friction +model_true = structural_simplify(friction_true()) +prob_true = ODEProblem(model_true, [], (0, 0.1), []) +sol_ref = solve(prob_true, Rodas4(); saveat = 0.001) +``` + +Let's plot it. + +```@example friction +scatter(sol_ref, label = "velocity") +``` + +That was the velocity. Let's also plot the friction force acting on the object throughout the simulation. + +```@example friction +scatter(sol_ref.t, friction.(first.(sol_ref.u)), label = "friction force") +``` + +## Model Setup + +Now, we will try to learn the same friction model using a neural network. We will use [`NeuralNetworkBlock`](@ref) to define neural network as a component. The input of the neural network is the velocity and the output is the friction force. We connect the neural network with the model using `RealInputArray` and `RealOutputArray` blocks. + +```@example friction +function friction_ude(Fu) + @variables y(t) = 0.0 + @constants Fu = Fu + @named nn_in = RealInputArray(nin = 1) + @named nn_out = RealOutputArray(nout = 1) + eqs = [Dt(y) ~ Fu - nn_in.u[1] + y ~ nn_out.u[1]] + return ODESystem(eqs, t, name = :friction, systems = [nn_in, nn_out]) +end + +Fu = 120.0 +model = friction_ude(Fu) + +chain = 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 = NeuralNetworkBlock(1, 1; chain = chain, rng = StableRNG(1111)) + +eqs = [connect(model.nn_in, nn.output) + connect(model.nn_out, nn.input)] + +ude_sys = complete(ODESystem(eqs, t, systems = [model, nn], name = :ude_sys)) +sys = structural_simplify(ude_sys) +``` + +## Optimization Setup + +We now setup the loss function and the optimization loop. + +```@example friction +function loss(x, (prob, sol_ref, get_vars, get_refs)) + new_p = SciMLStructures.replace(Tunable(), prob.p, x) + new_prob = remake(prob, p = new_p, u0 = eltype(x).(prob.u0)) + ts = sol_ref.t + new_sol = solve(new_prob, Rodas4(), saveat = ts, abstol = 1e-8, reltol = 1e-8) + loss = zero(eltype(x)) + for i in eachindex(new_sol.u) + loss += sum(abs2.(get_vars(new_sol, i) .- get_refs(sol_ref, i))) + end + if SciMLBase.successful_retcode(new_sol) + loss + else + Inf + end +end + +of = OptimizationFunction{true}(loss, AutoForwardDiff()) + +prob = ODEProblem(sys, [], (0, 0.1), []) +get_vars = getu(sys, [sys.friction.y]) +get_refs = getu(model_true, [model_true.y]) +x0 = reduce(vcat, getindex.((default_values(sys),), tunable_parameters(sys))) + +cb = (opt_state, loss) -> begin + @info "step $(opt_state.iter), loss: $loss" + return false +end + +op = OptimizationProblem(of, x0, (prob, sol_ref, get_vars, get_refs)) +res = solve(op, Adam(5e-3); maxiters = 10000, callback = cb) +``` + +## Visualization of results + +We now have a trained neural network! We can check whether running the simulation of the model embedded with the neural network matches the data or not. + +```@example friction +res_p = SciMLStructures.replace(Tunable(), prob.p, res) +res_prob = remake(prob, p = res_p) +res_sol = solve(res_prob, Rodas4(), saveat = sol_ref.t) +@test first.(sol_ref.u)≈first.(res_sol.u) rtol=1e-3 #hide +@test friction.(first.(sol_ref.u))≈(Fu .- first.(res_sol(res_sol.t, Val{1}).u)) rtol=1e-1 #hide +``` + +Also, it would be interesting to check the simulation before the training to get an idea of the starting point of the network. + +```@example friction +initial_sol = solve(prob, Rodas4(), saveat = sol_ref.t) +``` + +Now we plot it. + +```@example friction +scatter(sol_ref, idxs = [model_true.y], label = "ground truth velocity") +plot!(res_sol, idxs = [sys.friction.y], label = "velocity after training") +plot!(initial_sol, idxs = [sys.friction.y], label = "velocity before training") +``` + +It matches the data well! Let's also check the predictions for the friction force and whether the network learnt the friction model or not. + +```@example friction +scatter(sol_ref.t, friction.(first.(sol_ref.u)), label = "ground truth friction") +plot!(res_sol.t, Fu .- first.(res_sol(res_sol.t, Val{1}).u), + label = "friction from neural network") +``` + +It learns the friction model well! diff --git a/docs/src/index.md b/docs/src/index.md index f2945cc..a9ed1cf 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -4,7 +4,7 @@ CurrentModule = ModelingToolkitNeuralNets # ModelingToolkitNeuralNets -Documentation for [ModelingToolkitNeuralNets](https://github.com/SebastianM-C/ModelingToolkitNeuralNets.jl). +Documentation for [ModelingToolkitNeuralNets](https://github.com/SciML/ModelingToolkitNeuralNets.jl). ```@index ``` diff --git a/src/ModelingToolkitNeuralNets.jl b/src/ModelingToolkitNeuralNets.jl index b9346ba..fd55d09 100644 --- a/src/ModelingToolkitNeuralNets.jl +++ b/src/ModelingToolkitNeuralNets.jl @@ -1,7 +1,7 @@ module ModelingToolkitNeuralNets using ModelingToolkit: @parameters, @named, ODESystem, t_nounits -using ModelingToolkitStandardLibrary.Blocks: RealInput, RealOutput +using ModelingToolkitStandardLibrary.Blocks: RealInputArray, RealOutputArray using Symbolics: Symbolics, @register_array_symbolic, @wrapped using LuxCore: stateless_apply using Lux: Lux @@ -30,12 +30,12 @@ function NeuralNetworkBlock(n_input = 1, ca = ComponentArray{eltype}(init_params) @parameters p[1:length(ca)] = Vector(ca) - @parameters T::typeof(typeof(p))=typeof(p) [tunable = false] + @parameters T::typeof(typeof(ca))=typeof(ca) [tunable = false] - @named input = RealInput(nin = n_input) - @named output = RealOutput(nout = n_output) + @named input = RealInputArray(nin = n_input) + @named output = RealOutputArray(nout = n_output) - out = stateless_apply(chain, input.u, lazyconvert(typeof(ca), p)) + out = stateless_apply(chain, input.u, lazyconvert(T, p)) eqs = [output.u ~ out] diff --git a/test/lotka_volterra.jl b/test/lotka_volterra.jl index 3879c64..6b6f7a9 100644 --- a/test/lotka_volterra.jl +++ b/test/lotka_volterra.jl @@ -16,8 +16,8 @@ function lotka_ude() @variables t x(t)=3.1 y(t)=1.5 @parameters α=1.3 [tunable = false] δ=1.8 [tunable = false] Dt = ModelingToolkit.D_nounits - @named nn_in = RealInput(nin = 2) - @named nn_out = RealOutput(nout = 2) + @named nn_in = RealInputArray(nin = 2) + @named nn_out = RealOutputArray(nout = 2) eqs = [ Dt(x) ~ α * x + nn_in.u[1], @@ -50,7 +50,8 @@ eqs = [connect(model.nn_in, nn.output) connect(model.nn_out, nn.input)] ude_sys = complete(ODESystem( - eqs, ModelingToolkit.t_nounits, systems = [model, nn], name = :ude_sys)) + eqs, ModelingToolkit.t_nounits, systems = [model, nn], + name = :ude_sys, defaults = [nn.input.u => [0.0, 0.0]])) sys = structural_simplify(ude_sys)