Merge pull request #19 from SciML/sb/docs

Add Friction tutorial

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"

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"

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 <[email protected]> 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"
     ]
 )
 

docs/src/api.md

@@ -0,0 +1,5 @@
+# API
+
+```@docs
+NeuralNetworkBlock
+```

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!

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
 ```

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]
 

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)