Merge pull request #144 from JuliaComputing/optimize_rendering

Optimize rendering

Author: baggepinnen

docs/src/examples/free_motion.md

@@ -68,6 +68,7 @@ eqs = [connect(bar2.frame_a, world.frame_b)
                              spring1,
                              spring2,
                          ])
+model = complete(model)
 ssys = structural_simplify(IRSystem(model))
 prob = ODEProblem(ssys, [
     collect(body.body.v_0 .=> 0);

docs/src/examples/pendulum.md

@@ -42,6 +42,7 @@ nothing # hide
 With all components and connections defined, we can create an `ODESystem` like so:
 ```@example pendulum
 @named model = ODESystem(connections, t, systems=[world, joint, body])
+model = complete(model)
 nothing # hide
 ```
 The `ODESystem` is the fundamental model type in ModelingToolkit used for multibody-type models.
@@ -88,6 +89,7 @@ connections = [connect(world.frame_b, joint.frame_a)
                connect(body.frame_a, joint.frame_b)]
 
 @named model = ODESystem(connections, t, systems = [world, joint, body, damper])
+model = complete(model)
 ssys = structural_simplify(IRSystem(model))
 
 prob = ODEProblem(ssys, [damper.phi_rel => 1], (0, 10))
@@ -120,6 +122,7 @@ connections = [connect(world.frame_b, joint.frame_a)
                connect(body_0.frame_a, joint.frame_b)]
 
 @named model = ODESystem(connections, t, systems = [world, joint, body_0, damper, spring])
+model = complete(model)
 ssys = structural_simplify(IRSystem(model))
 
 prob = ODEProblem(ssys, [], (0, 10))
@@ -146,6 +149,7 @@ connections = [connect(world.frame_b, multibody_spring.frame_a)
                 connect(root_body.frame_a, multibody_spring.frame_b)]
 
 @named model = ODESystem(connections, t, systems = [world, multibody_spring, root_body])
+model = complete(model)
 ssys = structural_simplify(IRSystem(model))
 
 defs = Dict(collect(root_body.r_0) .=> [0, 1e-3, 0]) # The spring has a singularity at zero length, so we start some distance away
@@ -163,6 +167,7 @@ push!(connections, connect(multibody_spring.spring2d.flange_a, damper.flange_a))
 push!(connections, connect(multibody_spring.spring2d.flange_b, damper.flange_b))
 
 @named model = ODESystem(connections, t, systems = [world, multibody_spring, root_body, damper])
+model = complete(model)
 ssys = structural_simplify(IRSystem(model))
 prob = ODEProblem(ssys, defs, (0, 10))
 

docs/src/examples/quad.md

@@ -82,19 +82,16 @@ function RotorCraft(; closed_loop = true, addload=true)
     end
 
     thrusters = [Thruster(name = Symbol("thruster$i")) for i = 1:num_arms]
-    @named body = Body(m = body_mass, state_priority = 0, I_11=0.01, I_22=0.01, I_33=0.01, air_resistance=1)
-    @named freemotion = FreeMotion(state=true, isroot=true, quat=false) # We use Euler angles to describe the orientation of the rotorcraft.
+    @named body = Body(m = body_mass, state_priority = 0, I_11=0.01, I_22=0.01, I_33=0.01, air_resistance=1, isroot=true)
 
     connections = [
-        connect(world.frame_b, freemotion.frame_a)
-        connect(freemotion.frame_b, body.frame_a)
         y_alt ~ body.r_0[2]
-        y_roll ~ freemotion.phi[3]
-        y_pitch ~ freemotion.phi[1]
+        y_roll ~ body.phi[3]
+        y_pitch ~ body.phi[1]
         [connect(body.frame_a, arms[i].frame_a) for i = 1:num_arms]
         [connect(arms[i].frame_b, thrusters[i].frame_b) for i = 1:num_arms]
     ]
-    systems = [world; arms; body; thrusters; freemotion]
+    systems = [world; arms; body; thrusters]
     if addload
         @named load = Body(m = load_mass, air_resistance=0.1)
         @named cable = Rope(
@@ -142,15 +139,16 @@ function RotorCraft(; closed_loop = true, addload=true)
         # append!(connections, [feedback_gain.input.u[i] ~ arms[i].frame_b.r_0[2] for i = 1:num_arms ]) # Connect positions to controller
         # append!(connections, [feedback_gain.input.u[i+num_arms] ~ D(arms[i].frame_b.r_0[2]) for i = 1:num_arms]) # Connect velocities to controller
         # append!(connections, [feedback_gain.input.u[i+2num_arms] ~ Ie[i] for i = 1:num_arms]) #
-        # append!(connections, [feedback_gain.input.u[i] ~ freemotion.phi[[1,3][i]] for i = 1:2 ]) # Connect positions to controller
-        # append!(connections, [feedback_gain.input.u[i+2] ~ freemotion.phid[[1,3][i]] for i = 1:2]) # Connect velocities to controller
+        # append!(connections, [feedback_gain.input.u[i] ~ body.phi[[1,3][i]] for i = 1:2 ]) # Connect positions to controller
+        # append!(connections, [feedback_gain.input.u[i+2] ~ body.phid[[1,3][i]] for i = 1:2]) # Connect velocities to controller
         # push!(systems, feedback_gain)
         =#
     end
     @named model = ODESystem(connections, t; systems)
     complete(model)
 end
 model = RotorCraft(closed_loop=true, addload=true)
+model = complete(model)
 ssys = structural_simplify(IRSystem(model))
 
 op = [

docs/src/examples/ropes_and_cables.md

@@ -30,6 +30,7 @@ connections = [connect(world.frame_b, rope.frame_a)
 
 @named stiff_rope = ODESystem(connections, t, systems = [world, body, rope])
 
+stiff_rope = complete(stiff_rope)
 ssys = structural_simplify(IRSystem(stiff_rope))
 prob = ODEProblem(ssys, [], (0, 5))
 sol = solve(prob, Rodas4(autodiff=false))
@@ -53,6 +54,7 @@ connections = [connect(world.frame_b, rope.frame_a)
 
 @named flexible_rope = ODESystem(connections, t, systems = [world, body, rope])
 
+flexible_rope = complete(flexible_rope)
 ssys = structural_simplify(IRSystem(flexible_rope))
 prob = ODEProblem(ssys, [], (0, 8))
 sol = solve(prob, Rodas4(autodiff=false));
@@ -85,7 +87,7 @@ connections = [connect(world.frame_b, fixed.frame_a, chain.frame_a)
                connect(spring.frame_b, fixed.frame_b)]
 
 @named mounted_chain = ODESystem(connections, t, systems = [systems; world])
-
+mounted_chain = complete(mounted_chain)
 ssys = structural_simplify(IRSystem(mounted_chain))
 prob = ODEProblem(ssys, [
     collect(chain.link_8.body.w_a) .=> [0,0,0]; 

docs/src/examples/spherical_pendulum.md

@@ -27,6 +27,7 @@ connections = [connect(world.frame_b, joint.frame_a)
             connect(bar.frame_b, body.frame_a)]
 
 @named model = ODESystem(connections, t, systems = [world; systems])
+model = complete(model)
 ssys = structural_simplify(IRSystem(model))
 
 prob = ODEProblem(ssys, [], (0, 5))

docs/src/examples/wheel.md

@@ -332,7 +332,7 @@ prob = ODEProblem(ssys, [
     D(model.prismatic.r0[2]) => 0,
 ], (0.0, 15.0))
 sol = solve(prob, Rodas5Pr())
-render(model, sol, show_axis=false, x=0, y=0, z=4, traces=[model.slipBasedWheelJoint.frame_a], filename="slipwheel.gif")
+render(model, sol, show_axis=false, x=0, y=0, z=4, traces=[model.slipBasedWheelJoint.frame_a], filename="slipwheel.gif", cache=fale)
 nothing # hide
 ```
 
@@ -452,7 +452,7 @@ defs = merge(
 prob = ODEProblem(ssys, defs, (0.0, 5.0))
 sol = solve(prob, Rodas5P(autodiff=false))
 @test SciMLBase.successful_retcode(sol)
-Multibody.render(model, sol, show_axis=false, x=0, y=0, z=5, filename="twotrack.gif")
+Multibody.render(model, sol, show_axis=false, x=0, y=0, z=5, filename="twotrack.gif", cache=false)
 nothing # hide
 ```
 

ext/Render.jl

@@ -40,7 +40,6 @@ function recursive_extract(sol, s, depth=0)
     end
 end
 
-
 function getu(sol::FakeSol, syms)
     t->[recursive_extract(sol, s) for s in syms]
 end
@@ -75,6 +74,82 @@ function Base.getproperty(sol::FakeSol, s::Symbol)
     end
 end
 
+mutable struct CacheSol
+    model
+    sol
+    cache::Dict{Any, Any}
+    vars
+    last_t::Float64
+    function CacheSol(model, sol)
+        vars = get_all_vars(model) |> unique
+        Main.vars = vars
+        values = sol(0.0, idxs=vars)
+        new(model, sol, Dict(vars .=> values), vars, 0)
+    end
+end
+
+function get_all_vars(model, vars = Multibody.collect_all(unknowns(model)))
+    for sys in model.systems
+        if ModelingToolkit.isframe(sys)
+            newvars = Multibody.ModelingToolkit.renamespace.(model.name, Multibody.Symbolics.unwrap.(vec(ori(sys).R)))
+            append!(vars, newvars)
+        else
+            subsys_ns = getproperty(model, sys.name)
+            get_all_vars(subsys_ns, vars)
+        end
+    end
+    vars
+end
+
+
+get_cached(cs::CacheSol, t::AbstractArray, idxs) = cs.sol(t; idxs)
+function get_cached(cs::CacheSol, t::Real, idxs)
+    if ModelingToolkit.isparameter(idxs[1])
+        return cs.prob.ps[idxs]
+    end
+    if idxs isa AbstractArray{Num}
+        idxs = Multibody.Symbolics.unwrap.(idxs)
+    end
+    if !haskey(cs.cache, idxs[1])
+        # Fallback for things not in cache
+        return cs.sol(t; idxs)
+    end
+    if t != cs.last_t     
+        values = cs.sol(t, idxs=cs.vars)
+        cs.cache = Dict(cs.vars .=> values)
+        cs.last_t = t
+    end
+    if idxs isa Real
+        return cs.cache[idxs]
+    else
+        return [cs.cache[i] for i in idxs]
+    end
+end
+
+
+function getu(cs::CacheSol, syms)
+    t->get_cached(cs::CacheSol, t.t, syms)
+end
+
+function ModelingToolkit.parameter_values(cs::CacheSol)
+    ModelingToolkit.parameter_values(cs.sol)
+    # pars = Multibody.collect_all(parameters(cs.model))
+    # cs.prob.ps[pars]
+end
+
+function (cs::CacheSol)(t; idxs=nothing)
+    if idxs === nothing
+        cs.sol(t)
+    else
+        get_cached(cs, t, idxs)
+    end
+end
+
+function Base.getproperty(cs::CacheSol, s::Symbol)
+    s ∈ fieldnames(typeof(cs)) && return getfield(cs, s)
+    return getproperty(getfield(cs, :sol), s)
+end
+
 
 
 """
@@ -199,11 +274,15 @@ function render(model, sol,
     traces = nothing,
     display = false,
     loop = 1,
+    cache = true,
     kwargs...
     )
+    ModelingToolkit.iscomplete(model) || (model = complete(model))
     if sol isa ODEProblem
         sol = FakeSol(model, sol)
         return render(model, sol, 0; x, y, z, lookat, up, show_axis, kwargs...)[1]
+    elseif cache
+        sol = CacheSol(model, sol)
     end
     scene, fig = default_scene(x,y,z; lookat,up,show_axis)
     if timevec === nothing
@@ -258,14 +337,19 @@ function render(model, sol, time::Real;
     x = 2,
     y = 0.5,
     z = 2,
+    cache = true,
     kwargs...,
     )
 
+    ModelingToolkit.iscomplete(model) || (model = complete(model))
     slider = !(sol isa Union{ODEProblem, FakeSol})
 
     if sol isa ODEProblem
         sol = FakeSol(model, sol)
     end
+    if cache
+        sol = CacheSol(model, sol)
+    end
 
     # fig = Figure()
     # scene = LScene(fig[1, 1]).scene
@@ -332,13 +416,11 @@ end
 render!(scene, ::Any, args...) = false # Fallback for systems that have no rendering
 
 function render!(scene, ::typeof(Body), sys, sol, t)
-    sol(sol.t[1], idxs=sys.render)==true || return true # yes, == true
+    render, radius, length_fraction, cylinder_radius = sol(sol.t[1], idxs=[sys.render, sys.radius, sys.length_fraction, sys.cylinder_radius]) .|> Float32
+    render==true || return true # yes, == true
     color = get_color(sys, sol, :purple)
     r_cm = get_fun(sol, collect(sys.r_cm))
     framefun = get_frame_fun(sol, sys.frame_a)
-    radius = sol(sol.t[1], idxs=sys.radius) |> Float32
-    length_fraction = sol(sol.t[1], idxs=sys.length_fraction) |> Float32
-    cylinder_radius = sol(sol.t[1], idxs=sys.cylinder_radius) |> Float32
     thing = @lift begin # Sphere
         Ta = framefun($t)
         coords = (Ta*[r_cm($t); 1])[1:3] # TODO: make use of a proper transformation library instead of rolling own?
@@ -504,8 +586,6 @@ end
 function render!(scene, ::typeof(BodyShape), sys, sol, t)
     color = get_color(sys, sol, :purple)
     shapepath = get_shape(sys, sol)
-    Tshape = reshape(sol(sol.t[1], idxs=sys.shape_transform), 4, 4)
-    scale = Vec3f(Float32(sol(sol.t[1], idxs=sys.shape_scale))*ones(Float32, 3))
     if isempty(shapepath)
         radius = Float32(sol(sol.t[1], idxs=sys.radius))
         r_0a = get_fun(sol, collect(sys.frame_a.r_0))
@@ -520,7 +600,9 @@ function render!(scene, ::typeof(BodyShape), sys, sol, t)
         mesh!(scene, thing; color, specular = Vec3f(1.5), shininess=20f0, diffuse=Vec3f(1), transparency=true)
     else
         T = get_frame_fun(sol, sys.frame_a)
-
+        scale = Vec3f(Float32(sol(sol.t[1], idxs=sys.shape_scale))*ones(Float32, 3))
+        Tshape = reshape(sol(sol.t[1], idxs=sys.shape_transform), 4, 4)
+        
         @info "Loading shape mesh $shapepath"
         shapemesh = FileIO.load(shapepath)
         m = mesh!(scene, shapemesh; color, specular = Vec3f(1.5))
@@ -570,9 +652,7 @@ function render!(scene, ::typeof(BodyBox), sys, sol, t)
 
     # NOTE: This draws a solid box without the hole in the middle. Cannot figure out how to render a hollow box
     color = get_color(sys, sol, [1, 0.2, 1, 0.9])
-    width = Float32(sol(sol.t[1], idxs=sys.width))
-    height = Float32(sol(sol.t[1], idxs=sys.height))
-    length = Float32(sol(sol.t[1], idxs=sys.render_length))
+    width, height, length = Float32.(sol(sol.t[1], idxs=[sys.width, sys.height, sys.render_length]))
 
     length_dir = sol(sol.t[1], idxs=collect(sys.render_length_dir))
     width_dir = sol(sol.t[1], idxs=collect(sys.render_width_dir))
@@ -782,6 +862,7 @@ end
 
 Multibody.render!(scene, ::typeof(Multibody.URDFRevolute), sys, sol, t) = false
 Multibody.render!(scene, ::typeof(Multibody.URDFPrismatic), sys, sol, t) = false
+Multibody.render!(scene, ::typeof(Multibody.NullJoint), sys, sol, t) = false
 
 # ==============================================================================
 ## PlanarMechanics