solvers.py: Add vectorized version of EvalBinding (for python) (Robot…

…Locomotion#14064) * solvers.py: Add vectorized version of EvalBinding and CheckSatisfied
mcx · Sep 15, 2020 · 11d0d22 · 11d0d22
1 parent 4bb265a
commit 11d0d22
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 0 deletions.
diff --git a/bindings/pydrake/solvers/mathematicalprogram_py.cc b/bindings/pydrake/solvers/mathematicalprogram_py.cc
@@ -925,6 +925,22 @@ top-level documentation for :py:mod:`pydrake.math`.
           doc.MathematicalProgram.indeterminate.doc)
       .def("indeterminates_index", &MathematicalProgram::indeterminates_index,
           doc.MathematicalProgram.indeterminates_index.doc)
+      .def(
+          "EvalBindingVectorized",
+          [](const MathematicalProgram& prog,
+              const Binding<EvaluatorBase>& binding,
+              const MatrixX<double>& prog_var_vals) {
+            DRAKE_DEMAND(prog_var_vals.rows() == prog.num_vars());
+            MatrixX<double> Y(
+                binding.evaluator()->num_outputs(), prog_var_vals.cols());
+            for (int i = 0; i < prog_var_vals.cols(); ++i) {
+              Y.col(i) = prog.EvalBinding(binding, prog_var_vals.col(i));
+            }
+            return Y;
+          },
+          py::arg("binding"), py::arg("prog_var_vals"),
+          R"""(A "vectorized" version of EvalBinding.  It evaluates the binding 
+for every column of ``prog_var_vals``. )""")
       .def(
           "EvalBinding",
           [](const MathematicalProgram& prog,
@@ -1158,6 +1174,19 @@ top-level documentation for :py:mod:`pydrake.math`.
           doc.Constraint.lower_bound.doc)
       .def("upper_bound", &Constraint::upper_bound,
           doc.Constraint.upper_bound.doc)
+      .def(
+          "CheckSatisfiedVectorized",
+          [](Constraint& self, const Eigen::Ref<const Eigen::MatrixXd>& x,
+              double tol) {
+            DRAKE_DEMAND(x.rows() == self.num_vars());
+            std::vector<bool> satisfied(x.cols());
+            for (int i = 0; i < x.cols(); ++i) {
+              satisfied[i] = self.CheckSatisfied(x.col(i), tol);
+            }
+            return satisfied;
+          },
+          py::arg("x"), py::arg("tol"),
+          R"""(A "vectorized" version of CheckSatisfied.  It evaluates the constraint for every column of ``x``. )""")
       .def(
           "CheckSatisfied",
           [](Constraint& self, const Eigen::Ref<const Eigen::VectorXd>& x,

diff --git a/bindings/pydrake/solvers/test/mathematicalprogram_test.py b/bindings/pydrake/solvers/test/mathematicalprogram_test.py
@@ -284,6 +284,9 @@ def test_constraint_api(self):
         ce = prog.AddLinearEqualityConstraint(2*x0, 1).evaluator()
 
         self.assertTrue(c.CheckSatisfied([2.], tol=1e-3))
+        satisfied = c.CheckSatisfiedVectorized(
+            np.array([1., 2., 3.]).reshape((1, 3)), tol=1e-3)
+        self.assertEqual(satisfied, [False, True, True])
         self.assertFalse(c.CheckSatisfied([AutoDiffXd(1.)]))
         self.assertIsInstance(c.CheckSatisfied([x0]), sym.Formula)
 
@@ -351,11 +354,26 @@ def test_eval_binding(self):
         for (constraint, value_expected) in enum:
             value = result.EvalBinding(constraint)
             self.assertTrue(np.allclose(value, value_expected))
+            value = prog.EvalBinding(constraint, x_expected)
+            self.assertTrue(np.allclose(value, value_expected))
+            value = prog.EvalBindingVectorized(
+                constraint,
+                np.vstack((x_expected, x_expected)).T)
+            a = np.vstack((value_expected, value_expected)).T
+            self.assertTrue(np.allclose(
+                value, np.vstack((value_expected, value_expected)).T))
 
         enum = zip(costs, cost_values_expected)
         for (cost, value_expected) in enum:
             value = result.EvalBinding(cost)
             self.assertTrue(np.allclose(value, value_expected))
+            value = prog.EvalBinding(cost, x_expected)
+            self.assertTrue(np.allclose(value, value_expected))
+            value = prog.EvalBindingVectorized(
+                cost,
+                np.vstack((x_expected, x_expected)).T)
+            self.assertTrue(np.allclose(
+                value, np.vstack((value_expected, value_expected)).T))
 
         self.assertIsInstance(
             result.EvalBinding(costs[0]), np.ndarray)