Merge pull request #532 from JuliaRobotics/tk/clean-up-mass-matrix

Clean up `mass_matrix!` and `constraint_jacobian!`

src/RigidBodyDynamics.jl

@@ -51,6 +51,7 @@ export
     path,
     joints,
     tree_joints,
+    non_tree_joints,
     successor,
     predecessor,
     in_joints,

src/mechanism_algorithms.jl

@@ -246,32 +246,25 @@ The `out` argument must be an ``n_v \\times n_v`` lower triangular
 velocity vector ``v``.
 """
 function mass_matrix!(M::Symmetric, state::MechanismState)
-    @boundscheck size(M, 1) == num_velocities(state) || error("mass matrix has wrong size")
-    @boundscheck M.uplo == 'L' || error("expected a lower triangular symmetric matrix type as the mass matrix")
+    nv = num_velocities(state)
+    @boundscheck size(M, 1) == nv || throw(DimensionMismatch("mass matrix has wrong size"))
+    @boundscheck M.uplo == 'L' || throw(ArgumentError(("expected a lower triangular symmetric matrix")))
     update_motion_subspaces!(state)
     update_crb_inertias!(state)
-    fill!(M.data, 0)
     motion_subspaces = state.motion_subspaces.data
-    @inbounds for i in state.treejointids
-        bodyid = successorid(i, state)
-        Ici = crb_inertia(state, bodyid, false)
-        ancestor_joint_mask = values(state.ancestor_joint_masks[i]) # TODO
-        vrangei = velocity_range(state, i)
-        for coli in eachindex(vrangei)
-            vindexi = vrangei[coli]
-            Sicol = motion_subspaces[vindexi]
-            Ficol = Ici * Sicol
-            for j in Base.OneTo(i) # TODO: iterate directly over relevant joint ids
-                if ancestor_joint_mask[j]
-                    vrangej = velocity_range(state, j)
-                    for colj in eachindex(vrangej)
-                        vindexj = vrangej[colj]
-                        Sjcol = motion_subspaces[vindexj]
-                        # TODO: make nicer:
-                        @framecheck Ficol.frame Sjcol.frame
-                        M.data[vindexi, vindexj] = (transpose(angular(Ficol)) * angular(Sjcol) + transpose(linear(Ficol)) * linear(Sjcol))[1]
-                    end
-                end
+    @inbounds for i in Base.OneTo(nv)
+        jointi = velocity_index_to_joint_id(state, i)
+        bodyi = successorid(jointi, state)
+        Ici = crb_inertia(state, bodyi, false)
+        Si = motion_subspaces[i]
+        Fi = Ici * Si
+        for j in Base.OneTo(i)
+            jointj = velocity_index_to_joint_id(state, j)
+            M.data[i, j] = if supports(jointj, bodyi, state)
+                Sj = motion_subspaces[j]
+                (transpose(Fi) * Sj)[1]
+            else
+                zero(eltype(M))
             end
         end
     end
@@ -596,9 +589,7 @@ function constraint_jacobian!(jac::AbstractMatrix, rowranges, state::MechanismSt
                 for col in eachindex(vrange)
                     vindex = vrange[col]
                     Scol = state.motion_subspaces.data[vindex]
-                    # TODO: make nicer:
-                    @framecheck Tcol.frame Scol.frame
-                    jacelement = flipsign((transpose(angular(Tcol)) * angular(Scol) + transpose(linear(Tcol)) * linear(Scol))[1], sign)
+                    jacelement = flipsign((transpose(Tcol) * Scol)[1], sign)
                     jac[cindex, vindex] = jacelement
                 end
             end

src/mechanism_state.jl

@@ -47,7 +47,7 @@ struct MechanismState{X, M, C, JointCollection}
     qranges::JointDict{UnitRange{Int}}
     vranges::JointDict{UnitRange{Int}}
     constraintranges::IndexDict{JointID, UnitRange{JointID}, UnitRange{Int}}
-    ancestor_joint_masks::JointDict{JointDict{Bool}} # TODO: use a Matrix-backed type
+    support_set_masks::BodyDict{JointDict{Bool}} # TODO: use a Matrix-backed type
     constraint_jacobian_structure::JointDict{TreePath{RigidBody{M}, Joint{M}}} # TODO: use a Matrix-backed type
     q_index_to_joint_id::Vector{JointID}
     v_index_to_joint_id::Vector{JointID}
@@ -92,12 +92,10 @@ struct MechanismState{X, M, C, JointCollection}
         nontreejointids = lasttreejointid + 1 : lastjointid
         predecessor_and_successor_ids = JointDict{Pair{BodyID, BodyID}}(
             JointID(j) => (BodyID(predecessor(j, m)) => BodyID(successor(j, m))) for j in joints(m))
-        ancestor_joint_mask = function (joint)
-            JointDict{Bool}(
-                JointID(j) => j ∈ path(m, successor(joint, m), root_body(m)) for j in tree_joints(m)
-            )
+        support_set = function (b::RigidBody)
+            JointDict{Bool}(JointID(j) => (j ∈ tree_joints(m) ?  j ∈ path(m, b, root_body(m)) : false) for j in joints(m))
         end
-        ancestor_joint_masks = JointDict{JointDict{Bool}}(JointID(j) => ancestor_joint_mask(j) for j in tree_joints(m))
+        support_set_masks = BodyDict{JointDict{Bool}}(b => support_set(b) for b in bodies(m))
         constraint_jacobian_structure = JointDict{TreePath{RigidBody{M}, Joint{M}}}(
             JointID(j) => path(m, predecessor(j, m), successor(j, m)) for j in joints(m))
         qsegmented = SegmentedVector(q, tree_joints(m), num_positions)
@@ -164,7 +162,7 @@ struct MechanismState{X, M, C, JointCollection}
         new{X, M, C, JointCollection}(
             modcount(m), m, nonrootbodies, treejoints, nontreejoints,
             jointids, treejointids, nontreejointids,
-            predecessor_and_successor_ids, qranges, vranges, constraintranges, ancestor_joint_masks, constraint_jacobian_structure,
+            predecessor_and_successor_ids, qranges, vranges, constraintranges, support_set_masks, constraint_jacobian_structure,
             q_index_to_joint_id, v_index_to_joint_id,
             qsegmented, vsegmented, s,
             joint_transforms, joint_twists, joint_bias_accelerations, tree_joint_transforms, non_tree_joint_transforms,
@@ -551,41 +549,50 @@ $(SIGNATURES)
 
 Return the range of indices into the joint configuration vector ``q`` corresponding to joint `joint`.
 """
-Base.@propagate_inbounds configuration_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.qranges[joint]
+@propagate_inbounds configuration_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.qranges[joint]
 
 """
 $(SIGNATURES)
 
 Return the range of indices into the joint velocity vector ``v`` corresponding to joint `joint`.
 """
-Base.@propagate_inbounds velocity_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.vranges[joint]
+@propagate_inbounds velocity_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.vranges[joint]
 
 """
 $(SIGNATURES)
 
 Return the `JointID` of the joint associated with the given index into the configuration vector ``q``.
 """
-Base.@propagate_inbounds configuration_index_to_joint_id(state::MechanismState, qindex::Integer) = state.q_index_to_joint_id[qindex]
+@propagate_inbounds configuration_index_to_joint_id(state::MechanismState, qindex::Integer) = state.q_index_to_joint_id[qindex]
 
 """
 $(SIGNATURES)
 
 Return the `JointID` of the joint associated with the given index into the velocity vector ``v``.
 """
-Base.@propagate_inbounds velocity_index_to_joint_id(state::MechanismState, qindex::Integer) = state.v_index_to_joint_id[qindex]
+@propagate_inbounds velocity_index_to_joint_id(state::MechanismState, qindex::Integer) = state.v_index_to_joint_id[qindex]
 
 """
 $(SIGNATURES)
 
 Return the range of row indices into the constraint Jacobian corresponding to joint `joint`.
 """
-Base.@propagate_inbounds constraint_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.constraintranges[joint]
+@propagate_inbounds constraint_range(state::MechanismState, joint::Union{<:Joint, JointID}) = state.constraintranges[joint]
 
-## Accessor functions for cached variables
 
 """
 $(SIGNATURES)
 
+Return whether `joint` supports `body`, i.e., `joint` is a tree joint on the path between `body` and the root.
+"""
+@propagate_inbounds function supports(joint::Union{<:Joint, JointID}, body::Union{<:RigidBody, BodyID}, state::MechanismState)
+    state.support_set_masks[body][joint]
+end
+
+## Accessor functions for cached variables
+"""
+$(SIGNATURES)
+
 Return the joint transform for the given joint, i.e. the transform from
 `frame_after(joint)` to `frame_before(joint)`.
 """

src/spatial/motion_force_interaction.jl

@@ -308,6 +308,23 @@ for (MatrixType, VectorType) in (:WrenchMatrix => :(Union{Twist, SpatialAccelera
     end
 end
 
+for ForceSpaceMatrix in (:MomentumMatrix, :WrenchMatrix)
+    for (A, B) in ((ForceSpaceMatrix, :GeometricJacobian), (:GeometricJacobian, ForceSpaceMatrix))
+        @eval begin
+            function Base.:*(at::LinearAlgebra.Transpose{<:Any, <:$A}, b::$B)
+                a = parent(at)
+                @framecheck a.frame b.frame
+                transpose(angular(a)) * angular(b) + transpose(linear(a)) * linear(b)
+            end
+            function Base.:*(a::$A, bt::LinearAlgebra.Transpose{<:Any, <:$B})
+                b = parent(bt)
+                @framecheck a.frame b.frame
+                angular(a) * transpose(angular(b)) + linear(a) * transpose(linear(b))
+            end
+        end
+    end
+end
+
 """
 $(SIGNATURES)
 

test/test_mechanism_algorithms.jl

@@ -36,6 +36,24 @@ end
         show(devnull, x)
     end
 
+    @testset "supports" begin
+        Random.seed!(25)
+        mechanism = randmech()
+        mc_mechanism, = maximal_coordinates(mechanism)
+        for m in [mechanism, mc_mechanism]
+            state = MechanismState(m)
+            for body in bodies(m)
+                body_ancestors = RigidBodyDynamics.Graphs.ancestors(body, m.tree)
+                for joint in tree_joints(m)
+                    @test RigidBodyDynamics.supports(joint, body, state) == (successor(joint, m) ∈ body_ancestors)
+                end
+                for joint in non_tree_joints(m)
+                    @test !RigidBodyDynamics.supports(joint, body, state)
+                end
+            end
+        end
+    end
+
     @testset "basic stuff" begin
         Random.seed!(18)
         mechanism = randmech()