def test_joints_api(self):
        # Verify construction from both Isometry3d and 4x4 arrays,
        # and sanity-check that the accessors function.
        name = "z"
        prismatic_joint_np = PrismaticJoint(name, np.eye(4),
                                            np.array([0., 0., 1.]))
        prismatic_joint_isom = PrismaticJoint(name, Isometry3.Identity(),
                                              np.array([0., 0., 1.]))
        self.assertEqual(prismatic_joint_isom.get_num_positions(), 1)
        self.assertEqual(prismatic_joint_isom.get_name(), name)

        name = "theta"
        revolute_joint_np = RevoluteJoint(name, np.eye(4),
                                          np.array([0., 0., 1.]))
        revolute_joint_isom = RevoluteJoint(name, Isometry3.Identity(),
                                            np.array([0., 0., 1.]))
        self.assertEqual(revolute_joint_isom.get_num_positions(), 1)
        self.assertEqual(revolute_joint_isom.get_name(), name)

        name = "fixed"
        fixed_joint_np = FixedJoint(name, np.eye(4))
        fixed_joint_isom = FixedJoint(name, Isometry3.Identity())
        self.assertEqual(fixed_joint_isom.get_num_positions(), 0)
        self.assertEqual(fixed_joint_isom.get_name(), name)

        name = "rpy"
        rpy_floating_joint_np = RollPitchYawFloatingJoint(
            name, np.eye(4))
        rpy_floating_joint_isom = RollPitchYawFloatingJoint(
            name, Isometry3.Identity())
        self.assertEqual(rpy_floating_joint_isom.get_num_positions(), 6)
        self.assertEqual(rpy_floating_joint_isom.get_name(), name)
    def test_rigid_body_api(self):
        # Tests as much of the RigidBody API as is possible in isolation.
        # Adding collision geometry is *not* tested here, as it needs to
        # be done in the context of the RigidBodyTree.
        body = RigidBody()
        name = "body"
        body.set_name(name)
        self.assertEqual(body.get_name(), name)
        inertia = np.eye(6)
        body.set_spatial_inertia(inertia)
        self.assertTrue(np.allclose(inertia, body.get_spatial_inertia()))

        # Try adding a joint to a dummy body.
        body_joint = PrismaticJoint("z", np.eye(4),
                                    np.array([0., 0., 1.]))
        self.assertFalse(body.has_joint())
        dummy_body = RigidBody()
        body.add_joint(dummy_body, body_joint)
        self.assertEqual(body.getJoint(), body_joint)
        self.assertTrue(body.has_joint())

        # Try adding visual geometry.
        box_element = shapes.Box([1.0, 1.0, 1.0])
        box_visual_element = shapes.VisualElement(
            box_element, np.eye(4), [1., 0., 0., 1.])
        body.AddVisualElement(box_visual_element)
        body_visual_elements = body.get_visual_elements()
        self.assertEqual(len(body_visual_elements), 1)
        self.assertEqual(body_visual_elements[0].getGeometry().getShape(),
                         box_visual_element.getGeometry().getShape())

        # Test collision-related methods.
        self.assertEqual(body.get_num_collision_elements(), 0)
        self.assertEqual(len(body.get_collision_element_ids()), 0)
    def test_rigid_body_tree_collision_api(self):
        # This test creates a simple RBT with two spheres in simple
        # collision and verifies collision-detection APIs.
        rbt = RigidBodyTree()
        world_body = rbt.world()

        # Both bodies are unit-diameter spheres
        # spaced 0.5m from each other along the x axis.
        for k in range(2):
            body = RigidBody()
            body.set_name("body_{}".format(k))
            joint = PrismaticJoint("x", np.eye(4), np.array([1.0, 0., 0.]))
            body.add_joint(world_body, joint)
            body.set_spatial_inertia(np.eye(6))
            rbt.add_rigid_body(body)

            sphere_element = shapes.Sphere(0.5)
            sphere_collision_element = CollisionElement(
                sphere_element, np.eye(4))
            sphere_collision_element.set_body(body)
            rbt.addCollisionElement(sphere_collision_element, body, "default")

        rbt.compile()
        self.assertTrue(rbt.initialized())
        self.assertEqual(rbt.get_num_positions(), 2)

        q0 = np.array([-0.25, 0.25])
        kinsol = rbt.doKinematics(q0)
        # One point in each region of overlap:
        # Sphere one is (), sphere two is []
        # p0  (  p2  [  p3  ) p4 ] p5
        points = np.array([[-1., -0.6, 0., 0.6, 1.], [0., 0., 0., 0., 0],
                           [0., 0., 0., 0., 0]])
        n_pts = points.shape[1]

        phi, normal, x, body_x, body_idx = rbt.collisionDetectFromPoints(
            cache=kinsol, points=points, use_margins=False)
        # Check phi, but leave the other fields as API checks.
        self.assertTrue(
            np.allclose(phi, np.array([0.25, -0.15, -0.25, -0.15, 0.25])))
        self.assertEqual(phi.shape[0], n_pts)
        self.assertEqual(normal.shape, (3, n_pts))
        self.assertEqual(x.shape, (3, n_pts))
        self.assertEqual(normal.shape, (3, n_pts))

        point_pairs = rbt.ComputeMaximumDepthCollisionPoints(
            cache=kinsol, use_margins=False, throw_if_missing_gradient=True)
        self.assertEqual(len(point_pairs), 1)
        pp = point_pairs[0]
        self.assertEqual(rbt.FindCollisionElement(pp.idA), pp.elementA)
        self.assertEqual(rbt.FindCollisionElement(pp.idB), pp.elementB)
        possible_points = [np.array([0.5, 0., 0.]), np.array([-0.5, 0.0, 0.0])]
        for pt in [pp.ptA, pp.ptB]:
            self.assertTrue(
                np.allclose(pt, np.array([0.5, 0., 0.]))
                or np.allclose(pt, np.array([-0.5, 0., 0.])))
        self.assertTrue(np.allclose(np.abs(pp.normal), np.array([1., 0., 0.])))
        self.assertEqual(pp.distance, -0.5)
    def test_rigid_body_tree_programmatic_construction(self):
        # Tests RBT programmatic construction methods by assembling
        # a simple RBT with a prismatic and revolute joint, with
        # both visual and collision geometry on the last joint.
        rbt = RigidBodyTree()
        world_body = rbt.world()

        # body_1 is connected to the world via a prismatic joint along
        # the +z axis.
        body_1 = RigidBody()
        body_1.set_name("body_1")
        body_1_joint = PrismaticJoint("z", np.eye(4),
                                      np.array([0., 0., 1.]))
        body_1.add_joint(world_body, body_1_joint)
        rbt.add_rigid_body(body_1)

        # body_2 is connected to body_1 via a revolute joint around the z-axis.
        body_2 = RigidBody()
        body_2.set_name("body_2")
        body_2_joint = RevoluteJoint("theta", np.eye(4),
                                     np.array([0., 0., 1.]))
        body_2.add_joint(body_1, body_2_joint)
        box_element = shapes.Box([1.0, 1.0, 1.0])
        box_visual_element = shapes.VisualElement(
            box_element, np.eye(4), [1., 0., 0., 1.])
        body_2.AddVisualElement(box_visual_element)
        body_2_visual_elements = body_2.get_visual_elements()
        rbt.add_rigid_body(body_2)

        box_collision_element = CollisionElement(box_element, np.eye(4))
        box_collision_element.set_body(body_2)
        rbt.addCollisionElement(box_collision_element, body_2, "default")

        # Define a collision filter group containing bodies 1 and 2 and make
        # that group ignore itself.
        rbt.DefineCollisionFilterGroup(name="test_group")
        rbt.AddCollisionFilterGroupMember(
            group_name="test_group", body_name="body_1", model_id=0)
        rbt.AddCollisionFilterGroupMember(
            group_name="test_group", body_name="body_2", model_id=0)
        rbt.AddCollisionFilterIgnoreTarget("test_group", "test_group")

        self.assertFalse(rbt.initialized())
        rbt.compile()
        self.assertTrue(rbt.initialized())

        # The RBT's position vector should now be [z, theta].
        self.assertEqual(body_1.get_position_start_index(), 0)
        self.assertEqual(body_2.get_position_start_index(), 1)

        self.assertIsNotNone(
            rbt.FindCollisionElement(
                body_2.get_collision_element_ids()[0]))