From 658bc99ba5a1861859d4a562491d037b0dd38694 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Thu, 15 Aug 2024 10:07:44 +0200
Subject: [PATCH] add compile-time optimization based on parameter values

---
 src/Multibody.jl   | 58 ++++++++++++++++++++++++++++++++++++++++++++++
 test/test_robot.jl | 31 +++++++++++++++++++++++++
 2 files changed, 89 insertions(+)

diff --git a/src/Multibody.jl b/src/Multibody.jl
index 72145521..0c83ca9a 100644
--- a/src/Multibody.jl
+++ b/src/Multibody.jl
@@ -11,6 +11,7 @@ using StaticArrays
 export Rotational, Translational
 
 export render, render!
+export subs_constants
 
 """
 Find parameters that occur both scalarized and not scalarized
@@ -43,6 +44,63 @@ function benchmark_f(prob)
     @eval Main @btime $(prob.f)($dx, $x, $p, 0.0)
 end
 
+"""
+    subs_constants(model, c=[0, 1]; ssys = structural_simplify(IRSystem(model)), defs = defaults(model))
+
+A value-dependent compile-time optimization. Replace parameters in the model that have a default value contained in `c` with their value.
+
+Many parameters in multibody models have a sparse structure, e.g., rotation axes like `[1, 0, 0]`, or diagonal mass matrices etc. Specializing the generated code to this structure may increase performance, sometimes up to 2x. This function replaces all parameters that have a value contained in `c`, which defaults to `[0, 1]`, with their value.
+
+This performance optimization is primarily beneficial when the runtime of the simulation exceeds the compilation time of the model. For non-multibody models, sparse parameter structure is less common and this optimization is this not likely to be beneficial.
+
+# Drawbacks
+There are two main drawbacks to performing this optimization
+- Parameters that have been replaced **cannot be changed** after the optimization has been performed, without recompiling the model using `structural_simplify`.
+- The value of the repalced parameters are no longer accessible in the solution object. This typically means that **3D animations cannot be rendered** for models with replaced parameters.
+
+# Example usage
+```
+@named robot = Robot6DOF()
+robot = complete(robot)
+ssys = structural_simplify(IRSystem(robot))
+ssys = Multibody.subs_constants(model, [0, 1]; ssys)
+```
+
+If this optimization is to be performed repeatedly for several simulations of the same model, the indices of the substituted parameters can be stored and reused, call the lower-level function `Multibody.find_defaults_with_val` with the same signature as this function to obtain these indices, and then call `JuliaSimCompiler.freeze_parameters(ssys, inds)` with the indices to freeze the parameters.
+"""
+function subs_constants(model, c=[0, 1]; ssys = structural_simplify(IRSystem(model)), kwargs...)
+    inds = find_defaults_with_val(model, c; ssys, kwargs...)
+    ssys = JuliaSimCompiler.freeze_parameters(ssys, inds)
+end
+
+function find_defaults_with_val(model, c=[0, 1]; defs = defaults(model), ssys = structural_simplify(IRSystem(model)))
+    kvpairs = map(collect(pairs(defs))) do (key, val)
+        if val isa AbstractArray
+            string.(collect(key)) .=> collect(val)
+        else
+            string(key) => val
+        end
+    end
+    kvpairs = reduce(vcat, kvpairs)
+
+    p = parameters(ssys)
+    stringp = string.(p)
+
+    inds = map(c) do ci
+        sdefs = map(kvpairs) do (key, val)
+            if (val isa Bool) && !(ci isa Bool)
+                # We do not want to treat bools as an integer to prevent subsituting false when c contains 0
+                return key => false
+            end
+            key => isequal(val, ci)
+        end |> Dict
+        map(eachindex(p)) do i
+            get(sdefs, stringp[i], false)
+        end |> findall
+    end
+    sort(reduce(vcat, inds))
+end
+
 """
     scene, time = render(model, sol, t::Real; framerate = 30, traces = [])
     path        = render(model, sol, timevec = range(sol.t[1], sol.t[end], step = 1 / framerate); framerate = 30, timescale=1, display=false, loop=1)
diff --git a/test/test_robot.jl b/test/test_robot.jl
index f1c1ff08..6444887b 100644
--- a/test/test_robot.jl
+++ b/test/test_robot.jl
@@ -314,3 +314,34 @@ end
     @test maximum(abs, control_error) < 0.002
 end
 
+
+
+@testset "subs constants" begin
+    @info "Testing subs constants"
+    @named robot = Robot6DOF()
+    robot = complete(robot)
+    ssys = structural_simplify(IRSystem(robot))
+    ssys = Multibody.subs_constants(robot; ssys)
+    prob = ODEProblem(ssys, [
+        robot.mechanics.r1.phi => deg2rad(-60)
+        robot.mechanics.r2.phi => deg2rad(20)
+        robot.mechanics.r3.phi => deg2rad(90)
+        robot.mechanics.r4.phi => deg2rad(0)
+        robot.mechanics.r5.phi => deg2rad(-110)
+        robot.mechanics.r6.phi => deg2rad(0)
+
+        robot.axis1.motor.Jmotor.phi => deg2rad(-60) * (-105) # Multiply by gear ratio
+        robot.axis2.motor.Jmotor.phi => deg2rad(20) * (210)
+        robot.axis3.motor.Jmotor.phi => deg2rad(90) * (60)
+    ], (0.0, 2.0))
+    sol2 = solve(prob, Rodas5P(autodiff=false))
+
+    tv = 0:0.1:2
+    control_error = sol2(tv, idxs=robot.pathPlanning.controlBus.axisControlBus1.angle_ref-robot.pathPlanning.controlBus.axisControlBus1.angle)
+    @test maximum(abs, control_error) < 0.002
+
+    # using BenchmarkTools
+    # @btime solve(prob, Rodas5P(autodiff=false));
+    # 152.225 ms (2272926 allocations: 40.08 MiB)
+    # 114.113 ms (1833534 allocations: 33.38 MiB) # sub 0, 1
+end