diff --git a/tests/unit/test_joint.py b/tests/unit/test_joint.py new file mode 100644 index 0000000..9e2062d --- /dev/null +++ b/tests/unit/test_joint.py @@ -0,0 +1,38 @@ +import numpy as np +import pytest + +from urchin import Joint, xyz_rpy_to_matrix + + +@pytest.mark.parametrize("joint_type", ["continuous", "planar", "floating"]) +def test_none_configuration_uses_joint_origin(joint_type): + origin = xyz_rpy_to_matrix([1.0, 2.0, 3.0, 0.1, 0.2, 0.3]) + joint = Joint("joint", joint_type, "parent", "child", origin=origin) + + assert np.allclose(joint.get_child_pose(None), origin) + assert np.allclose(joint.get_child_poses(None, 2), np.tile(origin, (2, 1, 1))) + + +def test_planar_get_child_poses_matches_scalar_poses(): + origin = xyz_rpy_to_matrix([1.0, 2.0, 3.0, 0.1, 0.2, 0.3]) + joint = Joint("joint", "planar", "parent", "child", origin=origin) + cfgs = np.array([[0.5, 1.5], [-1.0, 2.0]]) + + expected = np.stack([joint.get_child_pose(cfg) for cfg in cfgs]) + + assert np.allclose(joint.get_child_poses(cfgs, len(cfgs)), expected) + + +def test_floating_get_child_poses_matches_scalar_poses(): + origin = xyz_rpy_to_matrix([1.0, 2.0, 3.0, 0.1, 0.2, 0.3]) + joint = Joint("joint", "floating", "parent", "child", origin=origin) + cfgs = np.array( + [ + [0.5, 1.5, 2.5, 0.1, 0.2, 0.3], + [-1.0, 2.0, 0.0, -0.3, 0.0, 0.2], + ] + ) + expected = np.stack([joint.get_child_pose(cfg) for cfg in cfgs]) + + assert np.allclose(joint.get_child_poses(cfgs, len(cfgs)), expected) + assert np.allclose(joint.get_child_poses(xyz_rpy_to_matrix(cfgs), len(cfgs)), expected) diff --git a/tests/unit/test_urdf.py b/tests/unit/test_urdf.py index 114769c..90be5db 100644 --- a/tests/unit/test_urdf.py +++ b/tests/unit/test_urdf.py @@ -2,7 +2,22 @@ import pytest import trimesh -from urchin import URDF, Joint, Link, Material, Transmission +from urchin import URDF, Joint, Link, Material, Transmission, xyz_rpy_to_matrix + + +def _make_multi_dof_urdf(): + links = [ + Link("base", None, None, None), + Link("floating_link", None, None, None), + Link("planar_link", None, None, None), + Link("tip", None, None, None), + ] + joints = [ + Joint("floating_joint", "floating", "base", "floating_link"), + Joint("planar_joint", "planar", "floating_link", "planar_link"), + Joint("continuous_joint", "continuous", "planar_link", "tip"), + ] + return URDF("multi_dof", links, joints) def test_urchin(tmpdir): @@ -109,3 +124,47 @@ def test_urchin(tmpdir): assert isinstance(x, URDF) x = x.copy(scale=[1, 1, 3]) assert isinstance(x, URDF) + + +def test_multi_dof_flat_configurations(): + robot = _make_multi_dof_urdf() + cfgs = np.array( + [ + [0.5, 1.5, 2.5, 0.1, 0.2, 0.3, 3.5, 4.5, 0.4], + [-1.0, 2.0, 0.0, -0.3, 0.0, 0.2, 1.0, -2.0, -0.5], + ] + ) + + expected = [] + for cfg in cfgs: + expected.append( + robot.link_fk( + { + "floating_joint": cfg[:6], + "planar_joint": cfg[6:8], + "continuous_joint": cfg[8], + }, + link="tip", + ) + ) + assert np.allclose(robot.link_fk(cfg.tolist(), link="tip"), expected[-1]) + assert np.allclose(robot.link_fk([cfg[:6], cfg[6:8], cfg[8]], link="tip"), expected[-1]) + + expected_batch = np.stack(expected) + assert np.allclose(robot.link_fk_batch(cfgs, link="tip"), expected_batch) + assert np.allclose( + robot.link_fk_batch( + { + "floating_joint": xyz_rpy_to_matrix(cfgs[:, :6]), + "planar_joint": cfgs[:, 6:8], + "continuous_joint": cfgs[:, 8], + }, + link="tip", + ), + expected_batch, + ) + + with pytest.raises(ValueError, match="degree of freedom"): + robot.link_fk(np.zeros(3)) + with pytest.raises(ValueError, match="degree of freedom"): + robot.link_fk_batch(np.zeros((2, 3))) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 888cdf5..a9a1b81 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -14,6 +14,16 @@ def test_rpy_to_matrix(): assert np.allclose(rpy_to_matrix([0, np.pi / 2, 0]), yr45) assert np.allclose(rpy_to_matrix([0, 0, np.pi / 2]), zr45) assert np.allclose(rpy_to_matrix(np.pi / 2 * np.ones(3)), c) + assert np.allclose( + rpy_to_matrix( + [ + [np.pi / 2, 0, 0], + [0, np.pi / 2, 0], + [0, 0, np.pi / 2], + ] + ), + np.stack([xr45, yr45, zr45]), + ) def test_matrix_to_rpy(): @@ -51,3 +61,13 @@ def test_xyz_rpy_to_matrix(): assert np.allclose(xyz_rpy_to_matrix([1, 2, 3, np.pi / 2, 0, 0]), xr45) assert np.allclose(xyz_rpy_to_matrix([2, 3, 1, 0, np.pi / 2, 0]), yr45) assert np.allclose(xyz_rpy_to_matrix([3, 1, 2, 0, 0, np.pi / 2]), zr45) + assert np.allclose( + xyz_rpy_to_matrix( + [ + [1, 2, 3, np.pi / 2, 0, 0], + [2, 3, 1, 0, np.pi / 2, 0], + [3, 1, 2, 0, 0, np.pi / 2], + ] + ), + np.stack([xr45, yr45, zr45]), + ) diff --git a/urchin/joint.py b/urchin/joint.py index af4efeb..7df5951 100644 --- a/urchin/joint.py +++ b/urchin/joint.py @@ -748,7 +748,7 @@ def get_child_poses(self, cfg: Optional[npt.ArrayLike], n_cfgs: int) -> np.ndarr Parameters ---------- - cfg : (n,) float or None + cfg : (n,) float, (n,2) float, (n,6) float, (n,4,4) float, or None The configuration values for this joint. They are interpreted based on the joint type as follows: @@ -756,8 +756,10 @@ def get_child_poses(self, cfg: Optional[npt.ArrayLike], n_cfgs: int) -> np.ndarr - ``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 an (n,6) matrix, or an (n,4,4) matrix. If ``cfg`` is ``None``, then this just returns the joint pose. @@ -783,9 +785,19 @@ def get_child_poses(self, cfg: Optional[npt.ArrayLike], n_cfgs: int) -> np.ndarr translation[:, :3, 3] = self.axis * cfg_arr[:, np.newaxis] return np.matmul(self.origin, translation) elif self.joint_type == "planar": - raise NotImplementedError() + cfg_arr = np.asanyarray(cfg, dtype=np.float64) + if cfg_arr.shape != (n_cfgs, 2): + raise ValueError("(n,2) float configuration required for planar joints") + translation = np.tile(np.eye(4, dtype=np.float64), (n_cfgs, 1, 1)) + translation[:, :3, 3] = cfg_arr.dot(self.origin[:3, :2].T) + return np.matmul(self.origin, translation) elif self.joint_type == "floating": - raise NotImplementedError() + cfg_arr = configure_origin(cfg) + if cfg_arr.shape != (n_cfgs, 4, 4): + raise ValueError( + "(n,6) or (n,4,4) float configuration required for floating joints" + ) + return np.matmul(self.origin, cfg_arr) else: raise ValueError("Invalid configuration") diff --git a/urchin/urdf.py b/urchin/urdf.py index a7625c7..bd5ea60 100644 --- a/urchin/urdf.py +++ b/urchin/urdf.py @@ -364,8 +364,9 @@ def link_fk( ---------- cfg : dict or (n), float A map from joints or joint names to configuration values for - each joint, or a list containing a value for each actuated joint - in sorted order from the base link. + each joint, a list containing a value for each actuated joint, + or a flat list containing a value for each degree of freedom in + sorted order from the base link. If not specified, all joints are assumed to be in their default configurations. link : str or :class:`.Link` @@ -452,7 +453,7 @@ def link_fk_batch(self, cfgs=None, link=None, links=None, use_names=False): One of the following: (A) a map from joints or joint names to vectors of joint configuration values, (B) a list of maps from joints or joint names to single configuration values, or (C) a list of ``n`` configuration vectors, - each of which has a vector with an entry for each actuated joint. + each of which has an entry for each degree of freedom. link : str or :class:`.Link` A single link or link name to return a pose for. links : list of str or list of :class:`.Link` @@ -1350,6 +1351,14 @@ def _validate_graph(self) -> tuple[Link, list[Link]]: raise ValueError("URDF has no base link") return base_link, end_links + @classmethod + def _joint_dof(cls, joint: Joint) -> int: + if joint.joint_type == "planar": + return 2 + elif joint.joint_type == "floating": + return 6 + return 1 + def _process_cfg( self, cfg: Union[ @@ -1358,11 +1367,11 @@ def _process_cfg( npt.ArrayLike, None, ], - ) -> dict[Joint, float]: + ) -> dict[Joint, Union[float, npt.NDArray[np.float64]]]: """Process a joint configuration spec into a dictionary mapping joints to configuration values. """ - joint_cfg: dict[Joint, float] = {} + joint_cfg: dict[Joint, Union[float, npt.NDArray[np.float64]]] = {} if cfg is None: return joint_cfg if isinstance(cfg, dict): @@ -1372,12 +1381,36 @@ def _process_cfg( elif isinstance(joint, Joint): joint_cfg[joint] = cfg[joint] elif isinstance(cfg, (list, tuple, np.ndarray)): - if len(cfg) != len(self.actuated_joints): + # Preserve the existing per-joint nested form before treating scalar + # sequences as flattened values for each degree of freedom. + if len(cfg) == len(self.actuated_joints) and any(np.ndim(value) > 0 for value in cfg): + for joint, value in zip(self.actuated_joints, cfg): + joint_cfg[joint] = value + return joint_cfg + + total_dofs = sum(self._joint_dof(joint) for joint in self.actuated_joints) + try: + cfg_arr = np.asanyarray(cfg, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise ValueError( + "Cfg must contain one value per actuated joint or one value " + "per degree of freedom" + ) from exc + if cfg_arr.shape != (total_dofs,): raise ValueError( - "Cfg must have same length as actuated joints if specified as a numerical array" + "Cfg must contain one value per actuated joint or one value " + "per degree of freedom" ) - for joint, value in zip(self.actuated_joints, cfg): - joint_cfg[joint] = value + + start = 0 + for joint in self.actuated_joints: + dof = self._joint_dof(joint) + end = start + dof + if dof == 1: + joint_cfg[joint] = float(cfg_arr[start]) + else: + joint_cfg[joint] = cfg_arr[start:end] + start = end else: raise TypeError("Invalid type for config") return joint_cfg @@ -1415,7 +1448,9 @@ def _process_cfgs( n_cfgs = len(cfgs[joint]) elif isinstance(cfgs, (list, tuple, np.ndarray)): n_cfgs = len(cfgs) - if isinstance(cfgs[0], dict): + if n_cfgs == 0: + pass + elif isinstance(cfgs[0], dict): for cfg in cfgs: for joint in cfg: if isinstance(joint, str): @@ -1431,20 +1466,39 @@ def _process_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] = cast(npt.NDArray[np.float64], cfgs[:, i]) + cfgs_arr = np.asanyarray(cfgs, dtype=np.float64) + total_dofs = sum(self._joint_dof(joint) for joint in self.actuated_joints) + if cfgs_arr.ndim != 2 or cfgs_arr.shape[1] != total_dofs: + raise ValueError( + "Cfgs must contain one column per degree of freedom " + f"(expected {total_dofs})" + ) + + start = 0 + for joint in self.actuated_joints: + dof = self._joint_dof(joint) + end = start + dof + batch_joint_values = cfgs_arr[:, start:end] + joint_cfg[joint] = batch_joint_values[:, 0] if dof == 1 else batch_joint_values + start = end else: raise ValueError("Incorrectly formatted config array") for j in joint_cfg: - if isinstance(joint_cfg[j], list): + configured_values = joint_cfg[j] + if isinstance(configured_values, list): from typing import cast as _cast - if len(_cast(list[float], joint_cfg[j])) == 0: + if len(_cast(list[float], configured_values)) == 0: joint_cfg[j] = None - elif n_cfgs is not None and len(_cast(list[float], joint_cfg[j])) != n_cfgs: + elif n_cfgs is not None and len(_cast(list[float], configured_values)) != n_cfgs: raise ValueError("Inconsistent number of configurations for joints") + elif ( + configured_values is not None + and n_cfgs is not None + and len(configured_values) != n_cfgs + ): + raise ValueError("Inconsistent number of configurations for joints") from typing import cast as _cast diff --git a/urchin/utils.py b/urchin/utils.py index 36be7c1..92ba271 100644 --- a/urchin/utils.py +++ b/urchin/utils.py @@ -24,25 +24,40 @@ def rpy_to_matrix(coords: npt.ArrayLike) -> npt.NDArray[np.float64]: 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( + if coords.shape[-1:] != (3,): + raise ValueError("Roll-pitch-yaw coordinates must have shape (...,3)") + + 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] + + return np.stack( [ - [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], + np.stack( + [c1 * c2, (c1 * s2 * s3) - (c3 * s1), (s1 * s3) + (c1 * c3 * s2)], + axis=-1, + ), + np.stack( + [c2 * s1, (c1 * c3) + (s1 * s2 * s3), (c3 * s1 * s2) - (c1 * s3)], + axis=-1, + ), + np.stack([-s2, c2 * s3, c2 * c3], axis=-1), ], - dtype=np.float64, + axis=-2, ) @@ -121,18 +136,22 @@ def xyz_rpy_to_matrix(xyz_rpy: npt.ArrayLike) -> npt.NDArray[np.float64]: 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) arr = np.asanyarray(xyz_rpy, dtype=np.float64) - matrix[:3, 3] = arr[:3] - matrix[:3, :3] = rpy_to_matrix(arr[3:]) + if arr.shape[-1:] != (6,): + raise ValueError("XYZ-RPY coordinates must have shape (...,6)") + + matrix = np.zeros(arr.shape[:-1] + (4, 4), dtype=np.float64) + matrix[..., 3, 3] = 1.0 + matrix[..., :3, 3] = arr[..., :3] + matrix[..., :3, :3] = rpy_to_matrix(arr[..., 3:]) return matrix @@ -260,23 +279,26 @@ def configure_origin( Parameters ---------- - value : None, (6,) float, or (4,4) float + value : None, (...,6) float, or (...,4,4) float The value to turn into the matrix. - If (6,), interpreted as xyzrpy coordinates. + If (...,6), interpreted as xyzrpy coordinates. Returns ------- - matrix : (4,4) float + matrix : (...,4,4) float The created matrix. If ``value`` is ``None``, returns the identity. """ if value is None: 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): - raise ValueError("Origin must be specified as a 4x4 homogenous transformation matrix") + elif value.shape[-2:] != (4, 4): + raise ValueError( + "Origin must be specified as an (...,6) vector or (...,4,4) " + "homogenous transformation matrix" + ) else: raise TypeError("Invalid type for origin, expect 4x4 matrix") return value