add compile-time optimization based on parameter values

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
@@ -60,6 +61,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 = render(model, prob)
  scene, time = render(model, sol, t::Real; framerate = 30, traces = [])

test/test_robot.jl

@@ -316,3 +316,34 @@ 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))
+
+ 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