add gainsched label processing + unit tests for errors

murrayrm · murrayrm · commit 6e56b733830c · 2022-12-28T11:33:44.000-08:00
diff --git a/control/statefbk.py b/control/statefbk.py
@@ -788,6 +788,9 @@ def create_statefbk_iosystem(
         # Generate the list of labels using the argument as a format string
         ud_labels = [ud_labels.format(i=i) for i in range(sys.ninputs)]
 
+    # Create the string of labels for the control system
+    input_labels = xd_labels + ud_labels + estimator.output_labels
+
     # Process gainscheduling variables, if present
     if gainsched:
         # Create a copy of the scheduling variable indices (default = empty)
@@ -797,10 +800,13 @@ def create_statefbk_iosystem(
         # Make sure the scheduling variable indices are the right length
         if len(gainsched_indices) != points.shape[1]:
             raise ControlArgument(
-                "Length of gainsched_indices must match dimension of"
+                "length of gainsched_indices must match dimension of"
                 " scheduling variables")
 
-        # TODO: Process scheduling variables
+        # Process scheduling variables
+        for i, idx in enumerate(gainsched_indices):
+            if isinstance(idx, str):
+                gainsched_indices[i] = input_labels.index(gainsched_indices[i])
 
         # Create interpolating function
         if method == 'nearest':
@@ -855,9 +861,8 @@ def _control_output(t, states, inputs, params):
         params = {} if gainsched else {'K': K}
         ctrl = NonlinearIOSystem(
             _control_update, _control_output, name='control',
-            inputs=xd_labels + ud_labels + estimator.output_labels,
-            outputs=list(sys.input_index.keys()), params=params,
-            states=nintegrators)
+            inputs=input_labels, outputs=list(sys.input_index.keys()),
+            params=params, states=nintegrators)
 
     elif type == 'linear' or type is None:
         # Create the matrices implementing the controller
diff --git a/control/tests/statefbk_test.py b/control/tests/statefbk_test.py
@@ -849,7 +849,8 @@ def test_gainsched_unicycle(unicycle, method):
     # Make sure that gains are different from 'nearest'
     if method is not None and method != 'nearest':
         ctrl_nearest, clsys_nearest = ct.create_statefbk_iosystem(
-            unicycle, (gains, points, 'nearest'), gainsched_indices=[3, 2])
+            unicycle, (gains, points, 'nearest'),
+            gainsched_indices=['ud[0]', 2])
         nearest_lin = clsys_nearest.linearize(xe, [xd, ud])
         assert not np.allclose(
             np.sort(clsys_lin.poles()), np.sort(nearest_lin.poles()), rtol=1e-2)
@@ -880,7 +881,8 @@ def test_gainsched_unicycle(unicycle, method):
 
     # Create gain scheduled controller
     ctrl, clsys = ct.create_statefbk_iosystem(
-        unicycle, (gains, points), gainsched_indices=[3, 7])
+        unicycle, (gains, points),
+        ud_labels=['vd', 'phid'], gainsched_indices=['vd', 'theta'])
 
     # Check the gain at the selected points
     for speed, angle in points:
@@ -903,3 +905,43 @@ def test_gainsched_unicycle(unicycle, method):
     resp = ct.input_output_response(clsys, timepts, [Xd, Ud], X0)
     np.testing.assert_allclose(
         resp.states[:, -1], Xd[:, -1], atol=1e-2, rtol=1e-2)
+
+def test_gainsched_errors(unicycle):
+    # Set up gain schedule (same as previous test)
+    speeds = [1, 5, 10]
+    angles = np.linspace(0, pi/2, 4)
+    points = list(itertools.product(speeds, angles))
+
+    Q = np.identity(unicycle.nstates)
+    R = np.identity(unicycle.ninputs)
+    gains = [np.array(ct.lqr(unicycle.linearize(
+        [0, 0, angle], [speed, 0]), Q, R)[0]) for speed, angle in points]
+
+    # Make sure the generic case works OK
+    ctrl, clsys = ct.create_statefbk_iosystem(
+        unicycle, (gains, points), gainsched_indices=[3, 2])
+    xd, ud = np.array([0, 0, angles[0]]), np.array([speeds[0], 0])
+    ctrl_lin = ctrl.linearize([], [xd, ud, xd*0])
+    K, S, E = ct.lqr(unicycle.linearize(xd, ud), Q, R)
+    np.testing.assert_allclose(
+        ctrl_lin.D[-xd.size:, -xd.size:], -K, rtol=1e-2)
+
+    # Wrong type of gain schedule argument
+    with pytest.raises(ControlArgument, match="gain must be an array"):
+        ctrl, clsys = ct.create_statefbk_iosystem(
+            unicycle, [gains, points], gainsched_indices=[3, 2])
+
+    # Mismatched dimensions for gains and points
+    with pytest.raises(ControlArgument, match="length of gainsched_indices"):
+        ctrl, clsys = ct.create_statefbk_iosystem(
+            unicycle, (gains, [speeds]), gainsched_indices=[3, 2])
+
+    # Unknown gain scheduling variable label
+    with pytest.raises(ValueError, match=".* not in list"):
+        ctrl, clsys = ct.create_statefbk_iosystem(
+            unicycle, (gains, points), gainsched_indices=['stuff', 2])
+
+    # Unknown gain scheduling method
+    with pytest.raises(ControlArgument, match="unknown gain scheduling method"):
+        ctrl, clsys = ct.create_statefbk_iosystem(
+            unicycle, (gains, points, 'stuff'), gainsched_indices=[3, 2])
