allow vectorization of jacobian dot function

ami-iit · Apr 26, 2024 · c5338b7 · c5338b7
1 parent b793e2c
commit c5338b7
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 19 deletions.
diff --git a/src/adam/core/rbd_algorithms.py b/src/adam/core/rbd_algorithms.py
@@ -288,8 +288,10 @@ def jacobian_dot(
 
         v = self.math.adjoint(L_H_B) @ B_v_IB
         a = self.math.adjoint_derivative(L_H_B, v) @ B_v_IB
-        J[:, :6] = self.math.adjoint(L_H_B)
-        J_dot[:, :6] = self.math.adjoint_derivative(L_H_B, v)
+        joint_J_list = []
+        joint_J_dict = {}
+        joint_J_dot_list = []
+        joint_J_dot_dict = {}
         for joint in chain:
             q = joint_positions[joint.idx] if joint.idx is not None else 0.0
             q_dot = joint_velocities[joint.idx] if joint.idx is not None else 0.0
@@ -301,8 +303,21 @@ def jacobian_dot(
             J_dot_j = self.math.adjoint_derivative(L_H_j, v) @ joint.motion_subspace()
             a += J_dot_j * q_dot
             if joint.idx is not None:
-                J[:, joint.idx + 6] = J_j
-                J_dot[:, joint.idx + 6] = J_dot_j
+                joint_J_dict[joint.idx] = J_j
+                joint_J_dot_dict[joint.idx] = J_dot_j
+
+        for i in range(self.NDoF):
+            if i in joint_J_dict:
+                joint_J_list.append(joint_J_dict[i])
+                joint_J_dot_list.append(joint_J_dot_dict[i])
+            else:
+                joint_J_list.append(self.math.factory.zeros(6))
+                joint_J_dot_list.append(self.math.factory.zeros(6))
+        joint_J = self.math.stack_list(joint_J_list)
+        joint_J_dot = self.math.stack_list(joint_J_dot_list)
+
+        J = self.math.vcat_inputs(self.math.adjoint(L_H_B), joint_J)
+        J_dot = self.math.vcat_inputs(self.math.adjoint_derivative(L_H_B, v), joint_J_dot)
 
         if (
             self.frame_velocity_representation

diff --git a/src/adam/core/spatial_math.py b/src/adam/core/spatial_math.py
@@ -474,10 +474,8 @@ def adjoint_derivative(self, H: npt.ArrayLike, v: npt.ArrayLike) -> npt.ArrayLik
         p = H[:3, 3]
         R_dot = self.skew(v[3:]) @ R
         p_dot = v[:3] - self.skew(p) @ v[3:]
-        X = self.factory.zeros(6, 6)
-        X[:3, :3] = R_dot
-        X[3:6, 3:6] = R_dot
-        X[:3, 3:6] = self.skew(p_dot) @ R + self.skew(p) @ R_dot
+        X = self.compose_blocks(R_dot, self.skew(p_dot) @ R + self.skew(p) @ R_dot,
+                                self.factory.zeros(3, 3), R_dot)
         return X
 
     def adjoint_inverse(self, H: npt.ArrayLike) -> npt.ArrayLike:
@@ -489,7 +487,6 @@ def adjoint_inverse(self, H: npt.ArrayLike) -> npt.ArrayLike:
         """
         R = H[:3, :3]
         p = H[:3, 3]
-        skew_block = -R.T @ self.skew(p)
         X = self.compose_blocks(R.T, -R.T @ self.skew(p),
                                 self.factory.zeros(3, 3), R.T)
         return X
@@ -508,10 +505,8 @@ def adjoint_inverse_derivative(
         p = H[:3, 3]
         R_dot = self.skew(v[3:]) @ R
         p_dot = v[:3] - self.skew(p) @ v[3:]
-        X = self.factory.zeros(6, 6)
-        X[:3, :3] = R_dot.T
-        X[3:6, 3:6] = R_dot.T
-        X[:3, 3:6] = -R_dot.T @ self.skew(p) - R.T @ self.skew(p_dot)
+        X = self.compose_blocks(R_dot.T, -R_dot.T @ self.skew(p) - R.T @ self.skew(p_dot),
+                                self.factory.zeros(3, 3), R_dot.T)
         return X
 
     def adjoint_mixed(self, H: npt.ArrayLike) -> npt.ArrayLike:
@@ -551,9 +546,8 @@ def adjoint_mixed_derivative(
         """
         R = H[:3, :3]
         R_dot = self.skew(v[3:]) @ R
-        X = self.factory.zeros(6, 6)
-        X[:3, :3] = R_dot
-        X[3:6, 3:6] = R_dot
+        X = self.compose_blocks(R_dot, self.factory.zeros(3,3),
+                                self.factory.zeros(3,3), R_dot)
         return X
 
     def adjoint_mixed_inverse_derivative(
@@ -568,9 +562,12 @@ def adjoint_mixed_inverse_derivative(
         """
         R = H[:3, :3]
         R_dot = self.skew(v[3:]) @ R
-        X = self.factory.zeros(6, 6)
-        X[:3, :3] = R_dot.T
-        X[3:6, 3:6] = R_dot.T
+        # X = self.factory.zeros(6, 6)
+        # X[:3, :3] = R_dot.T
+        # X[3:6, 3:6] = R_dot.T
+
+        X = self.compose_blocks(R_dot.T, self.factory.zeros(3,3),
+                                self.factory.zeros(3,3), R_dot.T)
         return X
 
     def homogeneous_inverse(self, H: npt.ArrayLike) -> npt.ArrayLike: