docs: add friction tutorial

Author: Sathvik Bhagavan

docs/make.jl

@@ -10,13 +10,17 @@ 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/assets/favicon.ico

@@ -0,0 +1,178 @@
+# 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 UDEComponents
+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
+```
+
+## Problem Setup
+
+Lets 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
+```
+
+Lets 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, lets simulate it from 0 to 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)
+```
+
+Lets plot it.
+
+```@example friction
+scatter(sol_ref, label = "velocity")
+```
+
+That was the velocity. Lets 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, lets 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 `RealInput` and `RealOutput` blocks.
+
+```@example friction
+function friction_ude(Fu)
+    @variables y(t) = 0.0
+    @constants Fu = Fu
+    @named nn_in = UDEComponents.RealInput2(nin = 1)
+    @named nn_out = UDEComponents.RealOutput2(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! Lets 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, lets check the simulation before the training as well to get an idea of the starting point of the network.
+
+```@example friction
+initial_sol = solve(prob, Rodas4(), saveat = sol_ref.t)
+```
+
+Lets 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! Lets also check 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!