add signal name kwargs for create_mpc_iosysm + kwargs testing for ct.optimal

murrayrm · murrayrm · commit a1dfc5ee9290 · 2023-02-04T11:15:22.000-08:00
diff --git a/control/optimal.py b/control/optimal.py
@@ -18,7 +18,6 @@
 
 from . import config
 from .exception import ControlNotImplemented
-from .timeresp import TimeResponseData
 
 # Define module default parameter values
 _optimal_trajectory_methods = {'shooting', 'collocation'}
@@ -140,7 +139,8 @@ class OptimalControlProblem():
     def __init__(
             self, sys, timepts, integral_cost, trajectory_constraints=[],
             terminal_cost=None, terminal_constraints=[], initial_guess=None,
-            trajectory_method=None, basis=None, log=False, **kwargs):
+            trajectory_method=None, basis=None, log=False, kwargs_check=True,
+            **kwargs):
         """Set up an optimal control problem."""
         # Save the basic information for use later
         self.system = sys
@@ -183,7 +183,7 @@ def __init__(
                     " discrete time systems")
 
         # Make sure there were no extraneous keywords
-        if kwargs:
+        if kwargs_check and kwargs:
             raise TypeError("unrecognized keyword(s): ", str(kwargs))
 
         self.trajectory_constraints = _process_constraints(
@@ -829,7 +829,7 @@ def compute_mpc(self, x, squeeze=None):
         return res.inputs[:, 0]
 
     # Create an input/output system implementing an MPC controller
-    def create_mpc_iosystem(self):
+    def create_mpc_iosystem(self, **kwargs):
         """Create an I/O system implementing an MPC controller"""
         # Check to make sure we are in discrete time
         if self.system.dt == 0:
@@ -857,11 +857,17 @@ def _output(t, x, u, params={}):
             res = self.compute_trajectory(u, print_summary=False)
             return res.inputs[:, 0]
 
+        # Define signal names, if they are not already given
+        if not kwargs.get('inputs'):
+            kwargs['inputs'] = self.system.state_labels
+        if not kwargs.get('outputs'):
+            kwargs['outputs'] = self.system.input_labels
+        if not kwargs.get('states'):
+            kwargs['states'] = self.system.ninputs * \
+                (self.timepts.size if self.basis is None else self.basis.N)
+
         return ct.NonlinearIOSystem(
-            _update, _output, dt=self.system.dt,
-            inputs=self.system.nstates, outputs=self.system.ninputs,
-            states=self.system.ninputs * \
-                (self.timepts.size if self.basis is None else self.basis.N))
+            _update, _output, dt=self.system.dt, **kwargs)
 
 
 # Optimal control result
@@ -923,7 +929,7 @@ def __init__(
             print("* Final cost:", self.cost)
 
         # Process data as a time response (with "outputs" = inputs)
-        response = TimeResponseData(
+        response = ct.TimeResponseData(
             ocp.timepts, inputs, states, issiso=ocp.system.issiso(),
             transpose=transpose, return_x=return_states, squeeze=squeeze)
 
@@ -1129,10 +1135,10 @@ def create_mpc_iosystem(
     ocp = OptimalControlProblem(
         sys, horizon, cost, trajectory_constraints=constraints,
         terminal_cost=terminal_cost, terminal_constraints=terminal_constraints,
-        log=log, **kwargs)
+        log=log, kwargs_check=False, **kwargs)
 
     # Return an I/O system implementing the model predictive controller
-    return ocp.create_mpc_iosystem()
+    return ocp.create_mpc_iosystem(**kwargs)
 
 
 #
diff --git a/control/tests/kwargs_test.py b/control/tests/kwargs_test.py
@@ -22,12 +22,13 @@
 import control.tests.flatsys_test as flatsys_test
 import control.tests.frd_test as frd_test
 import control.tests.interconnect_test as interconnect_test
+import control.tests.optimal_test as optimal_test
 import control.tests.statefbk_test as statefbk_test
 import control.tests.trdata_test as trdata_test
 
 
 @pytest.mark.parametrize("module, prefix", [
-    (control, ""), (control.flatsys, "flatsys.")
+    (control, ""), (control.flatsys, "flatsys."), (control.optimal, "optimal.")
 ])
 def test_kwarg_search(module, prefix):
     # Look through every object in the package
@@ -158,7 +159,7 @@ def test_matplotlib_kwargs(function, nsysargs, moreargs, kwargs, mplcleanup):
 kwarg_unittest = {
     'bode': test_matplotlib_kwargs,
     'bode_plot': test_matplotlib_kwargs,
-    'create_statefbk_iosystem': statefbk_test.TestStatefbk.test_statefbk_iosys,
+    'create_statefbk_iosystem': statefbk_test.TestStatefbk.test_statefbk_errors,
     'describing_function_plot': test_matplotlib_kwargs,
     'dlqe': test_unrecognized_kwargs,
     'dlqr': test_unrecognized_kwargs,
@@ -191,6 +192,8 @@ def test_matplotlib_kwargs(function, nsysargs, moreargs, kwargs, mplcleanup):
         flatsys_test.TestFlatSys.test_point_to_point_errors,
     'flatsys.solve_flat_ocp':
         flatsys_test.TestFlatSys.test_solve_flat_ocp_errors,
+    'optimal.create_mpc_iosystem': optimal_test.test_mpc_iosystem_rename,
+    'optimal.solve_ocp': optimal_test.test_ocp_argument_errors,
     'FrequencyResponseData.__init__':
         frd_test.TestFRD.test_unrecognized_keyword,
     'InputOutputSystem.__init__': test_unrecognized_kwargs,
@@ -205,7 +208,13 @@ def test_matplotlib_kwargs(function, nsysargs, moreargs, kwargs, mplcleanup):
     'StateSpace.sample': test_unrecognized_kwargs, 
     'TimeResponseData.__call__': trdata_test.test_response_copy,
     'TransferFunction.__init__': test_unrecognized_kwargs,
-    'TransferFunction.sample': test_unrecognized_kwargs, 
+    'TransferFunction.sample': test_unrecognized_kwargs,
+    'optimal.OptimalControlProblem.__init__':
+        optimal_test.test_ocp_argument_errors,
+    'optimal.OptimalControlProblem.compute_trajectory':
+        optimal_test.test_ocp_argument_errors,
+    'optimal.OptimalControlProblem.create_mpc_iosystem':
+        optimal_test.test_mpc_iosystem_rename,
 }
 
 #
diff --git a/control/tests/optimal_test.py b/control/tests/optimal_test.py
@@ -214,6 +214,35 @@ def test_mpc_iosystem_aircraft():
         xout[0:sys.nstates, -1], xd, atol=0.1, rtol=0.01)
 
 
+def test_mpc_iosystem_rename():
+    # Create a discrete time system (double integrator) + cost function
+    sys = ct.ss([[1, 1], [0, 1]], [[0], [1]], np.eye(2), 0, dt=True)
+    cost = opt.quadratic_cost(sys, np.eye(2), np.eye(1))
+    timepts = np.arange(0, 5)
+
+    # Create the default optimal control problem and check labels
+    mpc = opt.create_mpc_iosystem(sys, timepts, cost)
+    assert mpc.input_labels == sys.state_labels
+    assert mpc.output_labels == sys.input_labels
+
+    # Change the signal names
+    input_relabels = ['x1', 'x2']
+    output_relabels = ['u']
+    state_relabels = [f'x_[{i}]' for i in timepts]
+    mpc_relabeled = opt.create_mpc_iosystem(
+        sys, timepts, cost, inputs=input_relabels, outputs=output_relabels,
+        states=state_relabels, name='mpc_relabeled')
+    assert mpc_relabeled.input_labels == input_relabels
+    assert mpc_relabeled.output_labels == output_relabels
+    assert mpc_relabeled.state_labels == state_relabels
+    assert mpc_relabeled.name == 'mpc_relabeled'
+
+    # Make sure that unknown keywords are caught
+    # Unrecognized arguments
+    with pytest.raises(TypeError, match="unrecognized keyword"):
+        mpc = opt.create_mpc_iosystem(sys, timepts, cost, unknown=None)
+
+
 def test_mpc_iosystem_continuous():
     # Create a random state space system
     sys = ct.rss(2, 1, 1)
@@ -492,6 +521,14 @@ def test_ocp_argument_errors():
         res = opt.solve_ocp(
             sys, time, x0, cost, constraints, terminal_constraint=None)
 
+    with pytest.raises(TypeError, match="unrecognized keyword"):
+        ocp = opt.OptimalControlProblem(
+            sys, time, x0, cost, constraints, terminal_constraint=None)
+
+    with pytest.raises(TypeError, match="unrecognized keyword"):
+        ocp = opt.OptimalControlProblem(sys, time, cost, constraints)
+        ocp.compute_trajectory(x0, unknown=None)
+            
     # Unrecognized trajectory constraint type
     constraints = [(None, np.eye(3), [0, 0, 0], [0, 0, 0])]
     with pytest.raises(TypeError, match="unknown trajectory constraint type"):
diff --git a/control/tests/statefbk_test.py b/control/tests/statefbk_test.py
@@ -778,7 +778,7 @@ def test_statefbk_errors(self):
                 ctrl, clsys = ct.create_statefbk_iosystem(
                     sys, K, type='nonlinear', controller_type='nonlinear')
 
-        with pytest.raises(TypeError, match="unrecognized keywords"):
+        with pytest.raises(TypeError, match="unrecognized keyword"):
             ctrl, clsys = ct.create_statefbk_iosystem(sys, K, typo='nonlinear')
 
         with pytest.raises(ControlArgument, match="unknown controller_type"):
