[pydrake] Improve BsplineTrajectory usability (RobotLocomotion#18182)

mqnc · Oct 26, 2022 · aa2d1b8 · aa2d1b8
1 parent 8fde778
commit aa2d1b8
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 10 deletions.
diff --git a/bindings/pydrake/test/trajectories_test.py b/bindings/pydrake/test/trajectories_test.py
@@ -75,9 +75,16 @@ def test_bspline_trajectory(self, T):
         BsplineBasis = BsplineBasis_[T]
         BsplineTrajectory = BsplineTrajectory_[T]
 
+        # Call the default constructor.
         bspline = BsplineTrajectory()
         self.assertIsInstance(bspline, BsplineTrajectory)
         self.assertEqual(BsplineBasis().num_basis_functions(), 0)
+        # Call the vector<vector<T>> constructor.
+        bspline = BsplineTrajectory(basis=BsplineBasis(2, [0, 1, 2, 3]),
+                                    control_points=np.zeros((4, 2)))
+        self.assertEqual(bspline.rows(), 4)
+        self.assertEqual(bspline.cols(), 1)
+        # Call the vector<MatrixX<T>> constructor.
         bspline = BsplineTrajectory(
             basis=BsplineBasis(2, [0, 1, 2, 3]),
             control_points=[np.zeros((3, 4)), np.ones((3, 4))])
@@ -101,9 +108,9 @@ def test_bspline_trajectory(self, T):
             bspline.CopyBlock(start_row=1, start_col=2,
                               block_rows=2, block_cols=1),
             BsplineTrajectory)
-        bspline = BsplineTrajectory(
-            basis=BsplineBasis(2, [0, 1, 2, 3]),
-            control_points=[np.zeros(3), np.ones(3)])
+        bspline = BsplineTrajectory(basis=BsplineBasis(2, [0, 1, 2, 3]),
+                                    control_points=np.array([[0, 1], [0, 1],
+                                                             [0, 1]]))
         self.assertIsInstance(bspline.CopyHead(n=2), BsplineTrajectory)
         # Ensure we can copy.
         self.assertEqual(copy.copy(bspline).rows(), 3)

diff --git a/bindings/pydrake/trajectories_py.cc b/bindings/pydrake/trajectories_py.cc
@@ -153,6 +153,15 @@ struct Impl {
           m, "BsplineTrajectory", param, cls_doc.doc);
       cls  // BR
           .def(py::init<>())
+          // This overload will match 2d numpy arrays before
+          // std::vector<MatrixX<T>>. We want each column of the numpy array as
+          // a MatrixX of control points, but the std::vectors here are
+          // associated with the rows in numpy.
+          .def(py::init([](math::BsplineBasis<T> basis,
+                            std::vector<std::vector<T>> control_points) {
+            return Class(basis, MakeEigenFromRowMajorVectors(control_points));
+          }),
+              py::arg("basis"), py::arg("control_points"), cls_doc.ctor.doc)
           .def(py::init<math::BsplineBasis<T>, std::vector<MatrixX<T>>>(),
               py::arg("basis"), py::arg("control_points"), cls_doc.ctor.doc)
           .def("Clone", &Class::Clone, cls_doc.Clone.doc)

diff --git a/common/trajectories/bspline_trajectory.cc b/common/trajectories/bspline_trajectory.cc
@@ -23,7 +23,7 @@ template <typename T>
 BsplineTrajectory<T>::BsplineTrajectory(BsplineBasis<T> basis,
                                         std::vector<MatrixX<T>> control_points)
     : basis_(std::move(basis)), control_points_(std::move(control_points)) {
-  DRAKE_DEMAND(CheckInvariants());
+  CheckInvariants();
 }
 
 template <typename T>
@@ -234,9 +234,9 @@ boolean<T> BsplineTrajectory<T>::operator==(
 }
 
 template <typename T>
-bool BsplineTrajectory<T>::CheckInvariants() const {
-  return static_cast<int>(control_points_.size()) ==
-      basis_.num_basis_functions();
+void BsplineTrajectory<T>::CheckInvariants() const {
+  DRAKE_THROW_UNLESS(static_cast<int>(control_points_.size()) ==
+                     basis_.num_basis_functions());
 }
 
 DRAKE_DEFINE_CLASS_TEMPLATE_INSTANTIATIONS_ON_DEFAULT_SCALARS(

diff --git a/common/trajectories/bspline_trajectory.h b/common/trajectories/bspline_trajectory.h
@@ -129,7 +129,7 @@ class BsplineTrajectory final : public trajectories::Trajectory<T> {
   Serialize(Archive* a) {
     a->Visit(MakeNameValue("basis", &basis_));
     a->Visit(MakeNameValue("control_points", &control_points_));
-    DRAKE_THROW_UNLESS(CheckInvariants());
+    CheckInvariants();
   }
 
  private:
@@ -140,7 +140,7 @@ class BsplineTrajectory final : public trajectories::Trajectory<T> {
   std::unique_ptr<trajectories::Trajectory<T>> DoMakeDerivative(
       int derivative_order) const override;
 
-  bool CheckInvariants() const;
+  void CheckInvariants() const;
 
   math::BsplineBasis<T> basis_;
   std::vector<MatrixX<T>> control_points_;

diff --git a/common/trajectories/test/bspline_trajectory_test.cc b/common/trajectories/test/bspline_trajectory_test.cc
@@ -372,7 +372,7 @@ const char* const not_enough_control_points = R"""(
 GTEST_TEST(BsplineTrajectorySerializeTests, NotEnoughControlPointsTest) {
     DRAKE_EXPECT_THROWS_MESSAGE(
       LoadYamlString<BsplineTrajectory<double>>(not_enough_control_points),
-      ".*CheckInvariants.*");
+      ".*num_basis_functions.*");
 }
 
 }  // namespace trajectories