Add UpdateBound for PYFunctionConstraint (RobotLocomotion#14410)

wangdongnwpu · Dec 8, 2020 · 5e98e70 · 5e98e70
1 parent 098e830
commit 5e98e70
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
@@ -207,6 +207,10 @@ class PyFunctionConstraint : public Constraint {
         double_func_(Wrap<double, DoubleFunc>(func)),
         autodiff_func_(Wrap<AutoDiffXd, AutoDiffFunc>(func)) {}
 
+  using Constraint::set_bounds;
+  using Constraint::UpdateLowerBound;
+  using Constraint::UpdateUpperBound;
+
  protected:
   void DoEval(const Eigen::Ref<const Eigen::VectorXd>& x,
       Eigen::VectorXd* y) const override {
@@ -1263,6 +1267,17 @@ for every column of ``prog_var_vals``. )""")
           },
           py::arg("x"), doc.Constraint.CheckSatisfied.doc);
 
+  py::class_<PyFunctionConstraint, Constraint,
+      std::shared_ptr<PyFunctionConstraint>>(m, "PyFunctionConstraint",
+      "Constraint with its evaluator as a Python function")
+      .def("UpdateLowerBound", &PyFunctionConstraint::UpdateLowerBound,
+          py::arg("new_lb"), "Update the lower bound of the constraint.")
+      .def("UpdateUpperBound", &PyFunctionConstraint::UpdateUpperBound,
+          py::arg("new_ub"), "Update the upper bound of the constraint.")
+      .def("set_bounds", &PyFunctionConstraint::set_bounds,
+          py::arg("lower_bound"), py::arg("upper_bound"),
+          "Set both the lower and upper bounds of the constraint.");
+
   py::class_<LinearConstraint, Constraint, std::shared_ptr<LinearConstraint>>(
       m, "LinearConstraint", doc.LinearConstraint.doc)
       .def(py::init([](const Eigen::MatrixXd& A, const Eigen::VectorXd& lb,

diff --git a/bindings/pydrake/solvers/test/mathematicalprogram_test.py b/bindings/pydrake/solvers/test/mathematicalprogram_test.py
@@ -4,6 +4,7 @@
 from pydrake.solvers.mathematicalprogram import (
     LinearConstraint,
     MathematicalProgramResult,
+    PyFunctionConstraint,
     SolverOptions,
     SolverType,
     SolverId,
@@ -577,6 +578,37 @@ def test_linear_constraints(self):
         prog.AddLinearEqualityConstraint(
             2 * x[:2] + np.array([0, 1]), np.array([3, 2]))
 
+    def test_constraint_set_bounds(self):
+        prog = mp.MathematicalProgram()
+        x = prog.NewContinuousVariables(2, "x")
+
+        def constraint(x):
+            return x[1] ** 2
+        binding = prog.AddConstraint(constraint, [0], [1], vars=x)
+        self.assertIsInstance(binding.evaluator(), PyFunctionConstraint)
+        np.testing.assert_array_equal(
+            binding.evaluator().lower_bound(), np.array([0.]))
+        np.testing.assert_array_equal(
+            binding.evaluator().upper_bound(), np.array([1.]))
+        # Test UpdateLowerBound()
+        binding.evaluator().UpdateLowerBound(new_lb=[-1.])
+        np.testing.assert_array_equal(
+            binding.evaluator().lower_bound(), np.array([-1.]))
+        np.testing.assert_array_equal(
+            binding.evaluator().upper_bound(), np.array([1.]))
+        # Test UpdateLowerBound()
+        binding.evaluator().UpdateUpperBound(new_ub=[2.])
+        np.testing.assert_array_equal(
+            binding.evaluator().lower_bound(), np.array([-1.]))
+        np.testing.assert_array_equal(
+            binding.evaluator().upper_bound(), np.array([2.]))
+        # Test set_bounds()
+        binding.evaluator().set_bounds(lower_bound=[-3.], upper_bound=[4.])
+        np.testing.assert_array_equal(
+            binding.evaluator().lower_bound(), np.array([-3.]))
+        np.testing.assert_array_equal(
+            binding.evaluator().upper_bound(), np.array([4.]))
+
     def test_constraint_gradient_sparsity(self):
         prog = mp.MathematicalProgram()
         x = prog.NewContinuousVariables(2, "x")