diff --git a/urchin/urdf.py b/urchin/urdf.py index 80e8a05..70d0164 100644 --- a/urchin/urdf.py +++ b/urchin/urdf.py @@ -2388,37 +2388,23 @@ def get_child_pose(self, cfg=None): elif self.joint_type == "fixed": return self.origin elif self.joint_type in ["revolute", "continuous"]: - if cfg is None: - cfg = 0.0 - else: - cfg = float(cfg) + cfg = float(cfg) R = trimesh.transformations.rotation_matrix(cfg, self.axis) return self.origin.dot(R) elif self.joint_type == "prismatic": - if cfg is None: - cfg = 0.0 - else: - cfg = float(cfg) + cfg = float(cfg) translation = np.eye(4, dtype=np.float64) translation[:3, 3] = self.axis * cfg return self.origin.dot(translation) elif self.joint_type == "planar": - if cfg is None: - cfg = np.zeros(2, dtype=np.float64) - else: - cfg = np.asanyarray(cfg, dtype=np.float64) + cfg = np.asanyarray(cfg, dtype=np.float64) if cfg.shape != (2,): raise ValueError("(2,) float configuration required for planar joints") translation = np.eye(4, dtype=np.float64) translation[:3, 3] = self.origin[:3, :2].dot(cfg) return self.origin.dot(translation) elif self.joint_type == "floating": - if cfg is None: - cfg = np.zeros(6, dtype=np.float64) - else: - cfg = configure_origin(cfg) - if cfg is None: - raise ValueError("Invalid configuration for floating joint") + cfg = configure_origin(cfg) return self.origin.dot(cfg) else: raise ValueError("Invalid configuration") @@ -2437,8 +2423,10 @@ def get_child_poses(self, cfg, n_cfgs): - ``prismatic`` - a translation along the axis in meters. - ``revolute`` - a rotation about the axis in radians. - ``continuous`` - a rotation about the axis in radians. - - ``planar`` - Not implemented. - - ``floating`` - Not implemented. + - ``planar`` - the x and y translation values in the plane. + as an (n,2) matrix. + - ``floating`` - the xyz values followed by the rpy values, + as (n,6) or a (n,4,4) matrix. If ``cfg`` is ``None``, then this just returns the joint pose. @@ -2452,19 +2440,21 @@ def get_child_poses(self, cfg, n_cfgs): elif self.joint_type == "fixed": return np.tile(self.origin, (n_cfgs, 1, 1)) elif self.joint_type in ["revolute", "continuous"]: - if cfg is None: - cfg = np.zeros(n_cfgs) return np.matmul(self.origin, self._rotation_matrices(cfg, self.axis)) elif self.joint_type == "prismatic": - if cfg is None: - cfg = np.zeros(n_cfgs) translation = np.tile(np.eye(4), (n_cfgs, 1, 1)) translation[:, :3, 3] = self.axis * cfg[:, np.newaxis] return np.matmul(self.origin, translation) elif self.joint_type == "planar": - raise NotImplementedError() + cfg = np.asanyarray(cfg, dtype=np.float64) + if cfg.shape[-1] != 2: + raise ValueError("(...,2) float configuration required for planar joints") + translation = np.tile(np.eye(4), (n_cfgs, 1, 1)) + translation[:, :3, 3] = np.matmul(self.origin[np.newaxis, :3, :2], cfg[..., np.newaxis]).squeeze() + return np.matmul(self.origin, translation) elif self.joint_type == "floating": - raise NotImplementedError() + cfg = configure_origin(cfg) + return np.matmul(self.origin, cfg) else: raise ValueError("Invalid configuration") @@ -4029,12 +4019,27 @@ def _process_cfg(self, cfg): joint_cfg[joint] = cfg[joint] elif isinstance(cfg, (list, tuple, np.ndarray)): if len(cfg) != len(self.actuated_joints): - raise ValueError( - "Cfg must have same length as actuated joints " - "if specified as a numerical array" - ) - for joint, value in zip(self.actuated_joints, cfg): - joint_cfg[joint] = value + try: + cfg = np.asanyarray(cfg, dtype=np.float64) + except ValueError: + raise ValueError( + "Cfg must have same length as actuated joints or dof " + "if specified as a numerical array" + ) + if not isinstance(cfg, np.ndarray): + for joint, value in zip(self.actuated_joints, cfg): + joint_cfg[joint] = value + else: + sidx = 0 + for j in self.actuated_joints: + if j.joint_type == "planar": + eidx = sidx + 2 + elif j.joint_type == "floating": + eidx = sidx + 6 + else: + eidx = sidx + 1 + joint_cfg[j] = cfg[sidx:eidx].squeeze() + sidx = eidx else: raise TypeError("Invalid type for config") return joint_cfg @@ -4068,9 +4073,25 @@ def _process_cfgs(self, cfgs): elif cfgs[0] is None: pass else: - cfgs = np.asanyarray(cfgs, dtype=np.float64) - for i, j in enumerate(self.actuated_joints): - joint_cfg[j] = cfgs[:, i] + if isinstance(cfgs, (list, tuple)): + if len(cfgs) == len(self.actuated_joints): + for i, j in enumerate(self.actuated_joints): + joint_cfg[j].append(cfgs[i]) + else: + cfgs = np.asanyarray(cfgs, dtype=np.float64) + if isinstance(cfgs, np.ndarray): + sidx = 0 + for j in self.actuated_joints: + if j.joint_type == "planar": + eidx = sidx + 2 + elif j.joint_type == "floating": + eidx = sidx + 6 + else: + eidx = sidx + 1 + joint_cfg[j] = cfgs[:, sidx:eidx] + if joint_cfg[j].shape[-1] == 1: + joint_cfg[j] = joint_cfg[j].squeeze(-1) + sidx = eidx else: raise ValueError("Incorrectly formatted config array") diff --git a/urchin/utils.py b/urchin/utils.py index a822416..1d12e06 100644 --- a/urchin/utils.py +++ b/urchin/utils.py @@ -20,23 +20,30 @@ def rpy_to_matrix(coords): Parameters ---------- - coords : (3,) float + coords : (..., 3) float The roll-pitch-yaw coordinates in order (x-rot, y-rot, z-rot). Returns ------- - R : (3,3) float + R : (..., 3,3) float The corresponding homogenous 3x3 rotation matrix. """ coords = np.asanyarray(coords, dtype=np.float64) - c3, c2, c1 = np.cos(coords) - s3, s2, s1 = np.sin(coords) - - return np.array([ + coords_cos = np.cos(coords) + coords_sin = np.sin(coords) + c1 = coords_cos[..., 2] + c2 = coords_cos[..., 1] + c3 = coords_cos[..., 0] + s1 = coords_sin[..., 2] + s2 = coords_sin[..., 1] + s3 = coords_sin[..., 0] + + out = np.array([ [c1 * c2, (c1 * s2 * s3) - (c3 * s1), (s1 * s3) + (c1 * c3 * s2)], [c2 * s1, (c1 * c3) + (s1 * s2 * s3), (c3 * s1 * s2) - (c1 * s3)], [-s2, c2 * s3, c2 * c3] ], dtype=np.float64) + return np.moveaxis(out, (0, 1), (-2, -1)) def matrix_to_rpy(R, solution=1): @@ -113,17 +120,19 @@ def xyz_rpy_to_matrix(xyz_rpy): Parameters ---------- - xyz_rpy : (6,) float + xyz_rpy : (...,6) float The xyz_rpy vector. Returns ------- - matrix : (4,4) float + matrix : (...,4,4) float The homogenous transform matrix. """ - matrix = np.eye(4, dtype=np.float64) - matrix[:3,3] = xyz_rpy[:3] - matrix[:3,:3] = rpy_to_matrix(xyz_rpy[3:]) + xyz_rpy = np.asanyarray(xyz_rpy, dtype=np.float64) + matrix = np.zeros(xyz_rpy.shape[:-1] + (4,4), dtype=np.float64) + matrix[...,3,3] = 1.0 + matrix[...,:3,3] = xyz_rpy[...,:3] + matrix[...,:3,:3] = rpy_to_matrix(xyz_rpy[...,3:]) return matrix @@ -262,9 +271,9 @@ def configure_origin(value): value = np.eye(4, dtype=np.float64) elif isinstance(value, (list, tuple, np.ndarray)): value = np.asanyarray(value, dtype=np.float64) - if value.shape == (6,): + if value.shape[-1] == 6: value = xyz_rpy_to_matrix(value) - elif value.shape != (4,4): + elif value.shape[-2:] != (4,4): raise ValueError('Origin must be specified as a 4x4 ' 'homogenous transformation matrix') else: