vectorize jacobian calculation

src/adam/core/rbd_algorithms.py

@@ -167,16 +167,22 @@ def joints_jacobian(
         chain = self.model.get_joints_chain(self.root_link, frame)
         eye = self.math.factory.eye(4)
         B_H_j = eye
-        J = self.math.factory.zeros(6, self.NDoF)
         B_H_L = self.forward_kinematics(frame, eye, joint_positions)
         L_H_B = self.math.homogeneous_inverse(B_H_L)
-        for joint in chain:
-            q = joint_positions[joint.idx] if joint.idx is not None else 0.0
-            H_j = joint.homogeneous(q=q)
-            B_H_j = B_H_j @ H_j
-            L_H_j = L_H_B @ B_H_j
-            if joint.idx is not None:
-                J[:, joint.idx] = self.math.adjoint(L_H_j) @ joint.motion_subspace()
+        J_list = []
+        for i in range(self.NDoF):
+            joint = next((joint for joint in chain if joint.idx == i), None)
+            #Check if the joint with i as idx is in the chain
+            if (joint is not None) and (joint.idx is not None):
+                q = joint_positions[joint.idx] if joint.idx is not None else 0.0
+                H_j = joint.homogeneous(q=q)
+                B_H_j = B_H_j @ H_j
+                L_H_j = L_H_B @ B_H_j
+                J_list.append(self.math.adjoint(L_H_j) @ joint.motion_subspace())
+            else:
+                J_list.append(self.math.factory.zeros(6))
+
+        J = self.math.stack_list(J_list)
         return J
 
     def jacobian(
@@ -186,9 +192,7 @@ def jacobian(
         J = self.joints_jacobian(frame, joint_positions)
         B_H_L = self.forward_kinematics(frame, eye, joint_positions)
         L_X_B = self.math.adjoint_inverse(B_H_L)
-        J_tot = self.math.factory.zeros(6, self.NDoF + 6)
-        J_tot[:6, :6] = L_X_B
-        J_tot[:, 6:] = J
+        J_tot = self.math.vcat_inputs(L_X_B, J)
 
         if (
             self.frame_velocity_representation
@@ -202,7 +206,8 @@ def jacobian(
             w_H_L = base_transform @ B_H_L
             LI_X_L = self.math.adjoint_mixed(w_H_L)
             X = self.math.factory.eye(6 + self.NDoF)
-            X[:6, :6] = self.math.adjoint_mixed_inverse(base_transform)
+            X = self.math.compose_blocks(self.math.adjoint_mixed_inverse(base_transform), self.math.factory.zeros(6, self.NDoF),
+                                         self.math.factory.zeros(self.NDoF, 6), self.math.factory.eye(self.NDoF))
             J_tot = LI_X_L @ J_tot @ X
             return J_tot
         else:

src/adam/core/spatial_math.py

@@ -239,12 +239,9 @@ def H_revolute_joint(
         Returns:
             npt.ArrayLike: Homogeneous transform
         """
-        T = self.factory.eye(4)
         R = self.R_from_RPY(rpy) @ self.R_from_axis_angle(axis, q)
-        T[:3, :3] = R
-        T[0, 3] = xyz[0]
-        T[1, 3] = xyz[1]
-        T[2, 3] = xyz[2]
+        T = self.transform(xyz, R)
+
         return T
 
     def H_prismatic_joint(
@@ -460,10 +457,8 @@ def adjoint(self, H: npt.ArrayLike) -> npt.ArrayLike:
         """
         R = H[:3, :3]
         p = H[:3, 3]
-        X = self.factory.eye(6)
-        X[:3, :3] = R
-        X[3:6, 3:6] = R
-        X[:3, 3:6] = self.skew(p) @ R
+        X = self.compose_blocks(R, self.skew(p) @ R,
+                                self.factory.zeros(3, 3), R)
         return X
 
     def adjoint_derivative(self, H: npt.ArrayLike, v: npt.ArrayLike) -> npt.ArrayLike:
@@ -494,10 +489,9 @@ def adjoint_inverse(self, H: npt.ArrayLike) -> npt.ArrayLike:
         """
         R = H[:3, :3]
         p = H[:3, 3]
-        X = self.factory.eye(6)
-        X[:3, :3] = R.T
-        X[3:6, 3:6] = R.T
-        X[:3, 3:6] = -R.T @ self.skew(p)
+        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
 
     def adjoint_inverse_derivative(
@@ -528,9 +522,9 @@ def adjoint_mixed(self, H: npt.ArrayLike) -> npt.ArrayLike:
             npt.ArrayLike: adjoint matrix
         """
         R = H[:3, :3]
-        X = self.factory.eye(6)
-        X[:3, :3] = R
-        X[3:6, 3:6] = R
+        R = H[:3, :3]
+        X = self.compose_blocks(R, self.factory.zeros(3,3),
+                                self.factory.zeros(3,3), R)
         return X
 
     def adjoint_mixed_inverse(self, H: npt.ArrayLike) -> npt.ArrayLike:
@@ -541,9 +535,8 @@ def adjoint_mixed_inverse(self, H: npt.ArrayLike) -> npt.ArrayLike:
             npt.ArrayLike: adjoint matrix
         """
         R = H[:3, :3]
-        X = self.factory.eye(6)
-        X[:3, :3] = R.T
-        X[3:6, 3:6] = R.T
+        X = self.compose_blocks(R.T, self.factory.zeros(3,3),
+                                self.factory.zeros(3,3), R.T)
         return X
 
     def adjoint_mixed_derivative(
@@ -589,7 +582,5 @@ def homogeneous_inverse(self, H: npt.ArrayLike) -> npt.ArrayLike:
         """
         R = H[:3, :3]
         p = H[:3, 3]
-        T = self.factory.eye(4)
-        T[:3, :3] = R.T
-        T[:3, 3] = -R.T @ p
+        T = self.transform(-R.T @ p, R.T)
         return T

src/adam/pytorch/torch_like.py

@@ -23,6 +23,7 @@ def __post_init__(self):
         if self.array.dtype != torch.float64:
             self.array = self.array.double()
 
+    #TODO retry the index_put_ override, the error with no len() is due to it already being a tensor usually
     def __setitem__(self, idx, value: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
         """Overrides set item operator"""
         if type(self) is type(value):
@@ -149,6 +150,9 @@ def array(x: ntp.ArrayLike) -> "TorchLike":
         Returns:
             TorchLike: vector wrapping x
         """
+
+        if isinstance(x, torch.Tensor):
+            return TorchLike(x)
         return TorchLike(torch.tensor(x))
 
 
@@ -205,13 +209,17 @@ def skew(x: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
             TorchLike: skew matrix from x
         """
         if not isinstance(x, TorchLike):
+            x0, x1, x2 = x
             return TorchLike(
-                torch.tensor([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
+                torch.tensor([[0, -x2, x1], [x2, 0, -x0], [-x1, x0, 0]])
             )
+
         x = x.array
-        return TorchLike(
-            torch.tensor([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
-        )
+        x0, x1, x2 = x.unbind(dim=-1)
+        skew_x = torch.stack([torch.tensor(0), -x2, x1,
+                              x2, torch.tensor(0), -x0,
+                              -x1, x0, torch.tensor(0)], dim=-1).reshape(x.shape[:-1] + (3, 3))
+        return TorchLike(skew_x)
 
     @staticmethod
     def vertcat(*x: ntp.ArrayLike) -> "TorchLike":
@@ -236,3 +244,70 @@ def horzcat(*x: ntp.ArrayLike) -> "TorchLike":
         else:
             v = torch.tensor(x)
         return TorchLike(v)
+
+    @staticmethod
+    def transform(p: Union["TorchLike", ntp.ArrayLike], R: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """
+        Returns:
+            TorchLike: composition of 4x4 transformation matrix from translation vector and rotation matrix
+        """
+        R = R.array
+
+        if not isinstance(p, TorchLike):
+            T = torch.cat([
+                torch.cat([R, torch.tensor([p]).T], dim=-1),
+                torch.tensor([[0, 0, 0, 1]])
+                ], dim=-2)
+            return TorchLike(T)
+
+        p = p.array.unsqueeze(0)
+
+        T = torch.cat([
+                torch.cat([R, p.T], dim=-1),
+                torch.tensor([[0, 0, 0, 1]])
+                ], dim=-2)
+        return TorchLike(T)
+
+    @staticmethod
+    def compose_blocks(B00: Union["TorchLike", ntp.ArrayLike],
+                       B01: Union["TorchLike", ntp.ArrayLike],
+                       B10: Union["TorchLike", ntp.ArrayLike],
+                       B11: Union["TorchLike", ntp.ArrayLike],) -> "TorchLike":
+        """
+        Returns:
+            TorchLike: put together 4 matrix blocks in a single matrix
+        """
+        first_row = torch.cat([B00.array, B01.array], dim=-1)
+        second_row = torch.cat([B10.array, B11.array], dim=-1)
+
+        X = torch.cat([
+                first_row,
+                second_row
+                ], dim=-2)
+
+        return TorchLike(X)
+
+    @staticmethod
+    def stack_list(tensor_list: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """
+        Returns:
+            TorchLike: take a list of tensors and stack them together
+        """
+        if isinstance(tensor_list[0], TorchLike):
+            tensor_list = [t.array for t in tensor_list]
+
+        stacked_list = torch.stack(tensor_list, dim=1)
+
+        return TorchLike(stacked_list)
+
+    @staticmethod
+    def vcat_inputs(a: Union["TorchLike", ntp.ArrayLike], b: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
+        """
+        Returns:
+            TorchLike: take the base and joint components of the Jacobian to build the full Jacobian matrix
+        """
+        a = a.array
+        b = b.array
+        c = torch.cat([a, b], dim=-1)
+
+        return TorchLike(c)