add compile-time optimization based on parameter values

Author: baggepinnen

src/Multibody.jl

@@ -14,6 +14,7 @@ using StaticArrays
 export Rotational, Translational
 
 export render, render!
+export subs_constants
 
 """
 A helper function that calls `collect` on all parameters in a vector of parameters in a smart way
@@ -67,6 +68,64 @@ function get_systemtype(sys)
     eval(meta.type)
 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
+
+
 export multibody
 
 """

test/test_robot.jl

@@ -316,3 +316,36 @@ end
     @test maximum(abs, control_error) < 0.01
 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))
+
+    @test SciMLBase.successful_retcode(sol2)
+
+    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