updated optimal constraint handling + unit tests

murrayrm · murrayrm · commit f807b338bc77 · 2022-03-12T10:28:44.000-08:00
diff --git a/control/optimal.py b/control/optimal.py
@@ -146,6 +146,11 @@ def __init__(
         else:
             self.trajectory_constraints = trajectory_constraints
 
+        # Make sure that we recognize all of the constraint types
+        for ctype, fun, lb, ub in self.trajectory_constraints:
+            if not ctype in [opt.LinearConstraint, opt.NonlinearConstraint]:
+                raise TypeError(f"unknown constraint type {ctype}")
+
         # Process terminal constraints
         if isinstance(terminal_constraints, tuple):
             self.terminal_constraints = [terminal_constraints]
@@ -154,6 +159,11 @@ def __init__(
         else:
             self.terminal_constraints = terminal_constraints
 
+        # Make sure that we recognize all of the constraint types
+        for ctype, fun, lb, ub in self.terminal_constraints:
+            if not ctype in [opt.LinearConstraint, opt.NonlinearConstraint]:
+                raise TypeError(f"unknown constraint type {ctype}")
+
         #
         # Compute and store constraints
         #
@@ -401,7 +411,8 @@ def _constraint_function(self, coeffs):
                     value.append(fun @ np.hstack([states[:, i], inputs[:, i]]))
                 elif ctype == opt.NonlinearConstraint:
                     value.append(fun(states[:, i], inputs[:, i]))
-                else:
+                else:      # pragma: no cover
+                    # Checked above => we should never get here
                     raise TypeError(f"unknown constraint type {ctype}")
 
         # Evaluate the terminal constraint functions
@@ -413,7 +424,8 @@ def _constraint_function(self, coeffs):
                 value.append(fun @ np.hstack([states[:, i], inputs[:, i]]))
             elif ctype == opt.NonlinearConstraint:
                 value.append(fun(states[:, i], inputs[:, i]))
-            else:
+            else:      # pragma: no cover
+                # Checked above => we should never get here
                 raise TypeError(f"unknown constraint type {ctype}")
 
         # Update statistics
@@ -485,7 +497,8 @@ def _eqconst_function(self, coeffs):
                     value.append(fun @ np.hstack([states[:, i], inputs[:, i]]))
                 elif ctype == opt.NonlinearConstraint:
                     value.append(fun(states[:, i], inputs[:, i]))
-                else:
+                else:      # pragma: no cover
+                    # Checked above => we should never get here
                     raise TypeError(f"unknown constraint type {ctype}")
 
         # Evaluate the terminal constraint functions
@@ -497,7 +510,8 @@ def _eqconst_function(self, coeffs):
                 value.append(fun @ np.hstack([states[:, i], inputs[:, i]]))
             elif ctype == opt.NonlinearConstraint:
                 value.append(fun(states[:, i], inputs[:, i]))
-            else:
+            else:      # pragma: no cover
+                # Checked above => we should never get here
                 raise TypeError("unknown constraint type {ctype}")
 
         # Update statistics
@@ -844,7 +858,7 @@ def __init__(
 
 # Compute the input for a nonlinear, (constrained) optimal control problem
 def solve_ocp(
-        sys, horizon, X0, cost, constraints=[], terminal_cost=None,
+        sys, horizon, X0, cost, trajectory_constraints=[], terminal_cost=None,
         terminal_constraints=[], initial_guess=None, basis=None, squeeze=None,
         transpose=None, return_states=False, log=False, **kwargs):
 
@@ -865,7 +879,7 @@ def solve_ocp(
         Function that returns the integral cost given the current state
         and input.  Called as `cost(x, u)`.
 
-    constraints : list of tuples, optional
+    trajectory_constraints : list of tuples, optional
         List of constraints that should hold at each point in the time vector.
         Each element of the list should consist of a tuple with first element
         given by :meth:`scipy.optimize.LinearConstraint` or
@@ -943,13 +957,18 @@ def solve_ocp(
     :func:`OptimalControlProblem` for more information.
 
     """
+    # Process keyword arguments
+    if trajectory_constraints is None:
+        # Backwards compatibility
+        trajectory_constraints = kwargs.pop('constraints', None)
+
     # Allow 'return_x` as a synonym for 'return_states'
     return_states = ct.config._get_param(
         'optimal', 'return_x', kwargs, return_states, pop=True)
 
     # Set up the optimal control problem
     ocp = OptimalControlProblem(
-        sys, horizon, cost, trajectory_constraints=constraints,
+        sys, horizon, cost, trajectory_constraints=trajectory_constraints,
         terminal_cost=terminal_cost, terminal_constraints=terminal_constraints,
         initial_guess=initial_guess, basis=basis, log=log, **kwargs)
 
diff --git a/control/tests/optimal_test.py b/control/tests/optimal_test.py
@@ -441,6 +441,17 @@ def test_ocp_argument_errors():
         res = opt.solve_ocp(
             sys, time, x0, cost, constraints, terminal_constraint=None)
 
+    # Unrecognized trajectory constraint type
+    constraints = [(None, np.eye(3), [0, 0, 0], [0, 0, 0])]
+    with pytest.raises(TypeError, match="unknown constraint type"):
+        res = opt.solve_ocp(
+            sys, time, x0, cost, trajectory_constraints=constraints)
+
+    # Unrecognized terminal constraint type
+    with pytest.raises(TypeError, match="unknown constraint type"):
+        res = opt.solve_ocp(
+            sys, time, x0, cost, terminal_constraints=constraints)
+
 
 def test_optimal_basis_simple():
     sys = ct.ss2io(ct.ss([[1, 1], [0, 1]], [[1], [0.5]], np.eye(2), 0, 1))
@@ -484,3 +495,57 @@ def test_optimal_basis_simple():
         basis=flat.BezierFamily(4, Tf), return_x=True, log=True)
     assert res3.success
     np.testing.assert_almost_equal(res3.inputs, res1.inputs, decimal=3)
+
+
+def test_equality_constraints():
+    """Test out the ability to handle equality constraints"""
+    # Create the system (double integrator, continuous time)
+    sys = ct.ss2io(ct.ss(np.zeros((2, 2)), np.eye(2), np.eye(2), 0))
+
+    # Shortest path to a point is a line
+    Q = np.zeros((2, 2))
+    R = np.eye(2)
+    cost = opt.quadratic_cost(sys, Q, R)
+
+    # Set up the terminal constraint to be the origin
+    final_point = [opt.state_range_constraint(sys, [0, 0], [0, 0])]
+
+    # Create the optimal control problem
+    time = np.arange(0, 3, 1)
+    optctrl = opt.OptimalControlProblem(
+        sys, time, cost, terminal_constraints=final_point)
+
+    # Find a path to the origin
+    x0 = np.array([4, 3])
+    res = optctrl.compute_trajectory(x0, squeeze=True, return_x=True)
+    t, u1, x1 = res.time, res.inputs, res.states
+
+    # Bug prior to SciPy 1.6 will result in incorrect results
+    if NumpyVersion(sp.__version__) < '1.6.0':
+        pytest.xfail("SciPy 1.6 or higher required")
+
+    np.testing.assert_almost_equal(x1[:,-1], 0, decimal=4)
+
+    # Set up terminal constraints as a nonlinear constraint
+    def final_point_eval(x, u):
+        return x
+    final_point = [
+        (sp.optimize.NonlinearConstraint, final_point_eval, [0, 0], [0, 0])]
+
+    optctrl = opt.OptimalControlProblem(
+        sys, time, cost, terminal_constraints=final_point)
+
+    # Find a path to the origin
+    x0 = np.array([4, 3])
+    res = optctrl.compute_trajectory(x0, squeeze=True, return_x=True)
+    t, u2, x2 = res.time, res.inputs, res.states
+    np.testing.assert_almost_equal(x2[:,-1], 0, decimal=4)
+    np.testing.assert_almost_equal(u1, u2)
+    np.testing.assert_almost_equal(x1, x2)
+
+    # Try passing and unknown constraint type
+    final_point = [(None, final_point_eval, [0, 0], [0, 0])]
+    with pytest.raises(TypeError, match="unknown constraint type"):
+        optctrl = opt.OptimalControlProblem(
+            sys, time, cost, terminal_constraints=final_point)
+        res = optctrl.compute_trajectory(x0, squeeze=True, return_x=True)
