Add wheel constraints, wheel set and some examples (#119)

* move wheels to own file, add vertical weheel constraint

* add test for wheel constraint

* add wheel examples

* tidy up

* rm float param

docs/make.jl

@@ -38,6 +38,7 @@ makedocs(;
                  "Swing" => "examples/swing.md",
                  "Bodies in space" => "examples/space.md",
                  "Gyroscopic effects" => "examples/gyroscopic_effects.md",
+                 "Wheels" => "examples/wheel.md",
                  "Quadrotor with cable-suspended load" => "examples/quad.md",
              ],
              "Rotations and orientation" => "rotations.md",

docs/src/examples/wheel.md

@@ -0,0 +1,99 @@
+# Wheels
+
+## Rolling wheel
+```@example WHEEL
+using Multibody
+using ModelingToolkit
+import ModelingToolkitStandardLibrary.Mechanical.Rotational
+import ModelingToolkitStandardLibrary.Blocks
+using Plots
+using OrdinaryDiffEq
+using LinearAlgebra
+using JuliaSimCompiler
+using Test
+
+t = Multibody.t
+D = Differential(t)
+W(args...; kwargs...) = Multibody.world
+
+@mtkmodel WheelInWorld begin
+    @components begin
+        world = W()
+        wheel = RollingWheel(
+            radius = 0.3,
+            m = 2,
+            I_axis = 0.06,
+            I_long = 0.12,
+            x0 = 0.2,
+            z0 = 0.2,
+            der_angles = [0, 5, 1],
+        )
+    end
+end
+
+@named worldwheel = WheelInWorld()
+worldwheel = complete(worldwheel)
+
+defs = Dict([
+    worldwheel.wheel.body.r_0[1] => 0.2;
+    worldwheel.wheel.body.r_0[2] => 0.3;
+    worldwheel.wheel.body.r_0[3] => 0.2;
+])
+
+ssys = structural_simplify(IRSystem(worldwheel))
+prob = ODEProblem(ssys, defs, (0, 4))
+sol = solve(prob, Tsit5())
+@test SciMLBase.successful_retcode(sol)
+```
+
+```@example WHEEL
+import GLMakie
+Multibody.render(worldwheel, sol; filename = "worldwheel.gif")
+nothing # hide
+```
+
+![wheel animation](worldwheel.gif)
+
+## Wheel set
+```@example WHEEL
+@mtkmodel DrivingWheelSet begin
+    @components begin
+        sine1 = Blocks.Sine(frequency=1, amplitude=2)
+        sine2 = Blocks.Sine(frequency=1, amplitude=2, phase=pi/2)
+        torque1 = Rotational.Torque()
+        torque2 = Rotational.Torque()
+        wheels = RollingWheelSet(radius=0.1, m_wheel=0.5, I_axis=0.01, I_long=0.02, track=0.5, state_priority=100)
+        bar = FixedTranslation(r = [0.2, 0, 0])
+        body = Body(m=0.01, state_priority=1)
+        world = W()
+    end
+    @equations begin
+        connect(sine1.output, torque1.tau)
+        connect(sine2.output, torque2.tau)
+        connect(torque1.flange, wheels.axis1)
+        connect(torque2.flange, wheels.axis2)
+        connect(wheels.frame_middle, bar.frame_a)
+        connect(bar.frame_b, body.frame_a)
+    end
+end
+
+@named model = DrivingWheelSet()
+model = complete(model)
+ssys = structural_simplify(IRSystem(model))
+# display(unknowns(ssys))
+prob = ODEProblem(ssys, [
+    model.wheels.wheelSetJoint.prismatic1.s => 0.1
+    model.wheels.wheelSetJoint.prismatic2.s => 0.1
+], (0, 3))
+sol = solve(prob, Tsit5())
+@test SciMLBase.successful_retcode(sol)
+```
+
+```@example WHEEL
+import GLMakie
+Multibody.render(model, sol; filename = "wheelset.gif")
+nothing # hide
+```
+
+![wheelset animation](wheelset.gif)
+

ext/Render.jl

@@ -239,6 +239,7 @@ 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
     color = get_color(sys, sol, :purple)
     r_cm = get_fun(sol, collect(sys.r_cm))
     framefun = get_frame_fun(sol, sys.frame_a)
@@ -616,7 +617,12 @@ function render!(scene, ::typeof(Multibody.WorldForce), sys, sol, t)
 end
 
 function render!(scene, ::Function, sys, sol, t, args...) # Fallback for systems that have at least two frames
-    count(ModelingToolkit.isframe, sys.systems) == 2 || return false
+    frameinds = findall(ModelingToolkit.isframe, collect(sys.systems))
+    length(frameinds) == 2 || return false
+
+    nameof(sys.systems[frameinds[1]]) ∈ (:frame_a, :frame_b) || return false
+    nameof(sys.systems[frameinds[2]]) ∈ (:frame_a, :frame_b) || return false
+
     r_0a = get_fun(sol, collect(sys.frame_a.r_0))
     r_0b = get_fun(sol, collect(sys.frame_b.r_0))
     color = get_color(sys, sol, :green)

src/Multibody.jl

@@ -1,3 +1,5 @@
+# Find variables that are both array form and scalarized / collected
+# foreach(println, sort(unknowns(IRSystem(model)), by=string))
 module Multibody
 
 using LinearAlgebra
@@ -10,6 +12,22 @@ export Rotational, Translational
 
 export render, render!
 
+"""
+Find parameters that occur both scalarized and not scalarized
+"""
+function find_arry_problems(model)
+    # foreach(println, sort(unknowns(IRSystem(model)), by=string))
+    xs = string.(unknowns(IRSystem(model)))
+    for x in xs
+        endswith(x, ']') && continue # Only look at non-array vars
+        l = ncodeunits(x)
+        inds = findall(y->startswith(y, x*"["), xs)
+        isempty(inds) && continue
+        println(x)
+        println(xs[inds])
+    end
+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)
@@ -155,12 +173,15 @@ export World, world, Mounting1D, Fixed, FixedTranslation, FixedRotation, Body, B
 include("components.jl")
 
 export Revolute, Prismatic, Planar, Spherical, Universal,
-GearConstraint, RollingWheelJoint, RollingWheel, FreeMotion, RevolutePlanarLoopConstraint, Cylindrical
+GearConstraint, FreeMotion, RevolutePlanarLoopConstraint, Cylindrical
 include("joints.jl")
 
 export SphericalSpherical, UniversalSpherical, JointUSR, JointRRR
 include("fancy_joints.jl")
 
+export RollingWheelJoint, RollingWheel, RollingWheelSet, RollingConstraintVerticalWheel
+include("wheels.jl")
+
 export Spring, Damper, SpringDamperParallel, Torque, Force, WorldForce, WorldTorque
 include("forces.jl")
 

src/components.jl

@@ -79,15 +79,16 @@ function inner_gravity(point_gravity, mu, g, n, r)
 end
 
 
-@component function Fixed(; name, r = [0, 0, 0])
+@component function Fixed(; name, r = [0, 0, 0], render = true)
     systems = @named begin frame_b = Frame() end
-    @parameters begin r[1:3] = r,
-                               [
-                                   description = "Position vector from world frame to frame_b, resolved in world frame",
-                               ] end
+    @parameters begin
+        r[1:3] = r, [description = "Position vector from world frame to frame_b, resolved in world frame"]
+        render = render, [description = "Render the component in animations"]
+    end
     eqs = [collect(frame_b.r_0 .~ r)
            ori(frame_b) ~ nullrotation()]
-    compose(ODESystem(eqs, t; name), systems...)
+    sys = compose(ODESystem(eqs, t; name=:nothing), systems...)
+    add_params(sys, [render]; name)
 end
 
 @component function Mounting1D(; name, n = [1, 0, 0], phi0 = 0)
@@ -120,7 +121,7 @@ Fixed translation of `frame_b` with respect to `frame_a` with position vector `r
 
 Can be thought of as a massless rod. For a massive rod, see [`BodyShape`](@ref) or [`BodyCylinder`](@ref).
 """
-@component function FixedTranslation(; name, r, radius=0.02f0, color = purple)
+@component function FixedTranslation(; name, r, radius=0.02f0, color = purple, render = true)
     @named frame_a = Frame()
     @named frame_b = Frame()
     @parameters r[1:3]=r [
@@ -130,6 +131,7 @@ Can be thought of as a massless rod. For a massive rod, see [`BodyShape`](@ref)
     @parameters begin
         radius = radius, [description = "Radius of the body in animations"]
         color[1:4] = color, [description = "Color of the body in animations (RGBA)"]
+        render = render, [description = "Render the component in animations"]
     end
     fa = frame_a.f |> collect
     fb = frame_b.f |> collect
@@ -139,7 +141,7 @@ Can be thought of as a massless rod. For a massive rod, see [`BodyShape`](@ref)
                    (ori(frame_b) ~ ori(frame_a))
                    collect(0 .~ fa + fb)
                    (0 .~ taua + taub + cross(r, fb))]
-    pars = [r; radius; color]
+    pars = [r; radius; color; render]
     vars = []
     compose(ODESystem(eqs, t, vars, pars; name), frame_a, frame_b)
 end
@@ -252,6 +254,7 @@ This component has a single frame, `frame_a`. To represent bodies with more than
               air_resistance = 0.0,
               color = [1,0,0,1],
               state_priority = 2,
+              render = true,
               quat=false,)
     if state
         # @warn "Make the body have state variables by using isroot=true rather than state=true"
@@ -289,6 +292,7 @@ This component has a single frame, `frame_a`. To represent bodies with more than
     ]
     @parameters color[1:4] = color [description = "Color of the body in animations (RGBA)"]
     @parameters length_fraction=length_fraction, [description = "Fraction of the length of the body that is the cylinder from frame to COM in animations"]
+    @parameters render = render [description = "Render the component in animations"]
     # @parameters I[1:3, 1:3]=I [description="inertia tensor"]
 
     @parameters I_11=I_11 [description = "Element (1,1) of inertia tensor"]
@@ -361,7 +365,7 @@ This component has a single frame, `frame_a`. To represent bodies with more than
     # pars = [m;r_cm;radius;I_11;I_22;I_33;I_21;I_31;I_32;color]
 
     sys = ODESystem(eqs, t; name=:nothing, metadata = Dict(:isroot => isroot), systems = [frame_a])
-    add_params(sys, [radius; cylinder_radius; color; length_fraction]; name)
+    add_params(sys, [radius; cylinder_radius; color; length_fraction; render]; name)
 end
 
 

src/fancy_joints.jl

@@ -278,7 +278,7 @@ This joint aggregation can be used in cases where in reality a rod with spherica
         eRod_a(t)[1:3], [description="Unit vector in direction of rRod_a, resolved in frame_a (needed for analytic loop handling)"]
         rRod_0(t)[1:3] = rRod_ia, [description="Position vector from frame_a to frame_b resolved in world frame"]
         rRod_a(t)[1:3] = rRod_ia, [description="Position vector from frame_a to frame_b resolved in frame_a"]
-        (constraintResidue(t) = 0), [description="Constraint equation of joint in residue form: Either length constraint (= default) or equation to compute rod force (for analytic solution of loops in combination with Internal.RevoluteWithLengthConstraint/PrismaticWithLengthConstraint)"]
+        (constraintResidue(t) = 0), [description="Constraint equation of joint in residue form: Either length constraint (= default) or equation to compute rod force (for analytic solution of loops in combination with RevoluteWithLengthConstraint)"]
         f_b_a(t)[1:3], [description="frame_b.f resolved in frame_a"]
         f_ia_a(t)[1:3], [description="frame_ia.f resolved in frame_a"]
         t_ia_a(t)[1:3], [description="frame_ia.t resolved in frame_a"]