From c5338b753e0d7a3fcb65fbc2f58752d40f20d3a0 Mon Sep 17 00:00:00 2001 From: Evelyn D'Elia <17877131+evelyd@users.noreply.github.com> Date: Fri, 26 Apr 2024 17:16:33 +0200 Subject: [PATCH] allow vectorization of jacobian dot function --- src/adam/core/rbd_algorithms.py | 23 +++++++++++++++++++---- src/adam/core/spatial_math.py | 27 ++++++++++++--------------- 2 files changed, 31 insertions(+), 19 deletions(-) diff --git a/src/adam/core/rbd_algorithms.py b/src/adam/core/rbd_algorithms.py index 9300ba55..feb9e4ff 100644 --- 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 index ba1efc86..f0acedb7 100644 --- 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: