Python bindings for ContactResults (RobotLocomotion#8736)

YinjieLin · May 2, 2018 · 7c3d615 · 7c3d615
1 parent ae44866
commit 7c3d615
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 0 deletions.
diff --git a/bindings/pydrake/multibody/BUILD.bazel b/bindings/pydrake/multibody/BUILD.bazel
@@ -58,6 +58,7 @@ drake_pybind_library(
 drake_pybind_library(
     name = "rigid_body_plant_py",
     cc_deps = [
+        "//bindings/pydrake/systems:systems_pybind",
         "//lcmtypes:viewer",
     ],
     cc_so_name = "rigid_body_plant",

diff --git a/bindings/pydrake/multibody/rigid_body_plant_py.cc b/bindings/pydrake/multibody/rigid_body_plant_py.cc
@@ -2,6 +2,7 @@
 #include "pybind11/pybind11.h"
 
 #include "drake/bindings/pydrake/pydrake_pybind.h"
+#include "drake/bindings/pydrake/systems/systems_pybind.h"
 #include "drake/multibody/rigid_body_plant/drake_visualizer.h"
 #include "drake/multibody/rigid_body_plant/rigid_body_plant.h"
 
@@ -46,6 +47,59 @@ PYBIND11_MODULE(rigid_body_plant, m) {
         Class::kDefaultCharacteristicRadius;
   }
 
+  {
+    using Class = ContactInfo<T>;
+    py::class_<Class> cls(m, "ContactInfo");
+    cls
+        .def("get_element_id_1", &Class::get_element_id_1)
+        .def("get_element_id_2", &Class::get_element_id_2)
+        .def("get_resultant_force", &Class::get_resultant_force,
+             py_reference_internal);
+  }
+
+  {
+    using Class = ContactForce<T>;
+    py::class_<Class> cls(m, "ContactForce");
+    cls
+        .def(
+            py::init<
+                const Vector3<T>&,
+                const Vector3<T>&,
+                const Vector3<T>&,
+                const Vector3<T>&>(),
+            py::arg("application_point"),
+            py::arg("normal"),
+            py::arg("force"),
+            py::arg("torque") = Vector3<T>::Zero())
+        .def("get_reaction_force", &Class::get_reaction_force)
+        .def("get_application_point", &Class::get_application_point)
+        .def("get_force", &Class::get_force)
+        .def("get_normal_force", &Class::get_normal_force)
+        .def("get_tangent_force", &Class::get_tangent_force)
+        .def("get_torque", &Class::get_torque)
+        .def("get_normal", &Class::get_normal);
+  }
+
+  {
+    using Class = ContactResults<T>;
+    py::class_<Class> cls(m, "ContactResults");
+    cls
+        .def(py::init<>())
+        .def("get_num_contacts", &Class::get_num_contacts)
+        .def("get_contact_info",
+             &Class::get_contact_info, py_reference_internal)
+        .def("set_generalized_contact_force",
+            [](Class* self, const Eigen::VectorXd& f) {
+                self->set_generalized_contact_force(f);
+            })
+        .def("get_generalized_contact_force",
+             &Class::get_generalized_contact_force, py_reference_internal)
+        .def("AddContact", &Class::AddContact, py_reference_internal,
+            py::arg("element_a"), py::arg("element_b"))
+        .def("Clear", &Class::Clear);
+    pysystems::AddValueInstantiation<Class>(m);
+  }
+
   {
     using Class = CompliantMaterial;
     py::class_<Class> cls(m, "CompliantMaterial");

diff --git a/bindings/pydrake/multibody/test/rigid_body_plant_test.py b/bindings/pydrake/multibody/test/rigid_body_plant_test.py
@@ -113,6 +113,46 @@ def test_contact_parameters_api(self):
         param = cls(0.1, 0.2)
         param = cls(v_stiction_tolerance=0.1, characteristic_radius=0.2)
 
+    def test_contact_result_api(self):
+        # Test contact result bindings in isolation
+        # by constructing dummy contact results
+        results = mut.ContactResults()
+        self.assertEquals(results.get_num_contacts(), 0)
+        f = np.array([0., 1., 2.])
+        results.set_generalized_contact_force(f)
+        self.assertTrue(np.allclose(
+            results.get_generalized_contact_force(), f))
+
+        id_1 = 42
+        id_2 = 43
+        info = results.AddContact(element_a=id_1, element_b=id_2)
+        self.assertEquals(results.get_num_contacts(), 1)
+        info_dup = results.get_contact_info(0)
+        self.assertIsInstance(info_dup, mut.ContactInfo)
+        self.assertIs(info, info_dup)
+        self.assertEquals(info.get_element_id_1(), id_1)
+        self.assertEquals(info.get_element_id_2(), id_2)
+
+        pt = np.array([0.1, 0.2, 0.3])
+        normal = np.array([0.0, 1.0, 0.0])
+        force = np.array([0.9, 0.1, 0.9])
+        torque = np.array([0.6, 0.6, 0.6])
+        cf = mut.ContactForce(application_point=pt,
+                              normal=normal,
+                              force=force,
+                              torque=torque)
+        self.assertTrue(np.allclose(cf.get_application_point(), pt))
+        self.assertTrue(np.allclose(cf.get_force(), force))
+        self.assertTrue(np.allclose(cf.get_normal_force(),
+                                    normal*force.dot(normal)))
+        self.assertTrue(np.allclose(cf.get_tangent_force(),
+                                    force - normal*force.dot(normal)))
+        self.assertTrue(np.allclose(cf.get_torque(), torque))
+        self.assertTrue(np.allclose(cf.get_normal(), normal))
+
+        results.Clear()
+        self.assertEquals(results.get_num_contacts(), 0)
+
     def test_compliant_material_api(self):
 
         def flat_values(material):