Add __repr__ methods for RobotTimestamp and BeaconObservation (#274)

Also make estimate a property to allow writing it in python
ewfuentes · Mar 6, 2024 · 9e10363 · 9e10363
1 parent 80ed93e
commit 9e10363
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 7 deletions.
diff --git a/common/liegroups/se2_python.cc b/common/liegroups/se2_python.cc
@@ -1,4 +1,6 @@
 
+#include <sstream>
+
 #include "common/liegroups/se2.hh"
 #include "pybind11/eigen.h"
 #include "pybind11/operators.h"
@@ -26,12 +28,28 @@ PYBIND11_MODULE(se2_python, m) {
         .def(py::self * Eigen::Vector2d())
         .def(
             "__mul__", [](const SE2 &a, const SE2 &b) -> SE2 { return a * b; }, py::is_operator())
+        .def(
+            "__matmul__",
+            [](const SE2 &a, const Eigen::Matrix2Xd &b) -> Eigen::Matrix2Xd {
+                Eigen::Matrix2Xd out(2, b.cols());
+                for (int i = 0; i < b.cols(); i++) {
+                    out.col(i) = a * b.col(i);
+                }
+                return out;
+            },
+            py::is_operator())
         .def(
             "__imul__",
             [](SE2 &a, const SE2 &b) -> SE2 & {
                 a *= b;
                 return a;
             },
-            py::is_operator());
+            py::is_operator())
+        .def("__repr__", [](const SE2 &a) {
+            const Eigen::IOFormat fmt(8);
+            std::ostringstream ss;
+            ss << a.matrix().format(fmt);
+            return ss.str();
+        });
 }
 }  // namespace robot::liegroups
diff --git a/common/time/robot_time_python.cc b/common/time/robot_time_python.cc
@@ -1,4 +1,7 @@
 
+#include <iomanip>
+#include <sstream>
+
 #include "common/time/robot_time.hh"
 #include "pybind11/chrono.h"
 #include "pybind11/operators.h"
@@ -19,14 +22,23 @@ PYBIND11_MODULE(robot_time_python, m) {
         .def_property_readonly_static("MIN", &RobotTimestamp::min)
         .def_property_readonly_static("MAX", &RobotTimestamp::max)
         .def(py::self - py::self)
+        .def(py::self - RobotTimestamp::duration())
         .def(py::self + RobotTimestamp::duration())
         .def(RobotTimestamp::duration() + py::self)
+        .def(RobotTimestamp::duration() - py::self)
+        .def(py::self -= RobotTimestamp::duration())
         .def(py::self += RobotTimestamp::duration())
         .def(py::self == py::self)
         .def(py::self != py::self)
         .def(py::self < py::self)
         .def(py::self > py::self)
         .def(py::self <= py::self)
-        .def(py::self >= py::self);
+        .def(py::self >= py::self)
+        .def("__repr__", [](const RobotTimestamp &time) {
+            std::ostringstream ss;
+            ss << std::fixed << std::setprecision(9)
+               << std::chrono::duration<double>(time.time_since_epoch()).count();
+            return ss.str();
+        });
 }
 }  // namespace robot::time
diff --git a/experimental/beacon_sim/ekf_slam_python.cc b/experimental/beacon_sim/ekf_slam_python.cc
@@ -60,8 +60,8 @@ PYBIND11_MODULE(ekf_slam_python, m) {
         .def("load_map", &EkfSlam::load_map)
         .def("predict", &EkfSlam::predict)
         .def("update", &EkfSlam::update)
-        .def("estimate", py::overload_cast<>(&EkfSlam::estimate, py::const_))
-        .def("estimate", py::overload_cast<>(&EkfSlam::estimate))
+        .def_property("estimate", py::overload_cast<>(&EkfSlam::estimate, py::const_),
+                      py::overload_cast<>(&EkfSlam::estimate))
         .def("config", &EkfSlam::config);
 }
 }  // namespace robot::experimental::beacon_sim
diff --git a/experimental/beacon_sim/ekf_slam_python_test.py b/experimental/beacon_sim/ekf_slam_python_test.py
@@ -41,7 +41,7 @@ def test_happy_case(self):
             ekf.predict(current_time, old_robot_from_new_robot)
 
             # Compute the observation
-            beacon_in_robot = ekf.estimate().local_from_robot().inverse() * BEACON_IN_LOCAL
+            beacon_in_robot = ekf.estimate.local_from_robot().inverse() * BEACON_IN_LOCAL
 
             range_m = np.linalg.norm(beacon_in_robot)
             bearing_rad = np.arctan2(beacon_in_robot[1], beacon_in_robot[0])
@@ -51,7 +51,7 @@ def test_happy_case(self):
             ekf.update([obs])
 
         # Verification
-        est_beacon_in_local = ekf.estimate().beacon_in_local(BEACON_ID)
+        est_beacon_in_local = ekf.estimate.beacon_in_local(BEACON_ID)
         self.assertAlmostEqual(est_beacon_in_local[0], BEACON_IN_LOCAL[0])
         self.assertAlmostEqual(est_beacon_in_local[1], BEACON_IN_LOCAL[1])
 

diff --git a/experimental/beacon_sim/generate_observations_python.cc b/experimental/beacon_sim/generate_observations_python.cc
@@ -1,4 +1,6 @@
 
+#include <sstream>
+
 #include "experimental/beacon_sim/generate_observations.hh"
 #include "pybind11/pybind11.h"
 #include "pybind11/stl.h"
@@ -11,6 +13,19 @@ PYBIND11_MODULE(generate_observations_python, m) {
         .def(py::init<std::optional<int>, std::optional<double>, std::optional<double>>())
         .def_readwrite("maybe_id", &BeaconObservation::maybe_id)
         .def_readwrite("maybe_range_m", &BeaconObservation::maybe_range_m)
-        .def_readwrite("maybe_bearing_rad", &BeaconObservation::maybe_bearing_rad);
+        .def_readwrite("maybe_bearing_rad", &BeaconObservation::maybe_bearing_rad)
+        .def("__repr__", [](const BeaconObservation &obs) {
+            std::ostringstream ss;
+            ss << "<BeaconObservation id:"
+               << (obs.maybe_id.has_value() ? std::to_string(obs.maybe_id.value()) : "Unknown")
+               << " range_m: "
+               << (obs.maybe_range_m.has_value() ? std::to_string(obs.maybe_range_m.value())
+                                                 : "Unknown")
+               << " bearing_rad: "
+               << (obs.maybe_bearing_rad.has_value() ? std::to_string(obs.maybe_bearing_rad.value())
+                                                     : "Unknown")
+               << ">";
+            return ss.str();
+        });
 }
 }  // namespace robot::experimental::beacon_sim