require StateSpace for lqr/lqe first arg + process uniformly

murrayrm · murrayrm · commit fdd399f6cacf · 2021-12-26T18:22:55.000-08:00
diff --git a/control/statefbk.py b/control/statefbk.py
@@ -44,7 +44,7 @@
 
 from . import statesp
 from .mateqn import care, dare, _check_shape
-from .statesp import _ssmatrix, _convert_to_statespace
+from .statesp import StateSpace, _ssmatrix, _convert_to_statespace
 from .lti import LTI, isdtime, isctime
 from .exception import ControlSlycot, ControlArgument, ControlDimension, \
     ControlNotImplemented
@@ -366,19 +366,18 @@ def lqe(*args, **keywords):
         # Call dlqe
         return dlqe(*args, **keywords)
 
-    try:
-        sys = args[0]           # Treat the first argument as a system
-        if isinstance(sys, LTI):
-            # Convert LTI system to state space
-            sys = _convert_to_statespace(sys)
-
-        # Extract A, G (assume disturbances come through input), and C
-        A = np.array(sys.A, ndmin=2, dtype=float)
-        G = np.array(sys.B, ndmin=2, dtype=float)
-        C = np.array(sys.C, ndmin=2, dtype=float)
+    # If we were passed a state space  system, use that to get system matrices
+    if isinstance(args[0], StateSpace):
+        A = np.array(args[0].A, ndmin=2, dtype=float)
+        G = np.array(args[0].B, ndmin=2, dtype=float)
+        C = np.array(args[0].C, ndmin=2, dtype=float)
         index = 1
 
-    except AttributeError:
+    elif isinstance(args[0], LTI):
+        # Don't allow other types of LTI systems
+        raise ControlArgument("LTI system must be in state space form")
+
+    else:
         # Arguments should be A and B matrices
         A = np.array(args[0], ndmin=2, dtype=float)
         G = np.array(args[1], ndmin=2, dtype=float)
@@ -490,14 +489,18 @@ def dlqe(*args, **keywords):
     if isinstance(args[0], LTI) and isctime(args[0], strict=True):
         raise ControlArgument("dlqr() called with a continuous time system")
 
-    try:
-        # If this works, we were (probably) passed a system as the
-        # first argument; extract A and B
+    # If we were passed a state space  system, use that to get system matrices
+    if isinstance(args[0], StateSpace):
         A = np.array(args[0].A, ndmin=2, dtype=float)
         G = np.array(args[0].B, ndmin=2, dtype=float)
         C = np.array(args[0].C, ndmin=2, dtype=float)
         index = 1
-    except AttributeError:
+
+    elif isinstance(args[0], LTI):
+        # Don't allow other types of LTI systems
+        raise ControlArgument("LTI system must be in state space form")
+
+    else:
         # Arguments should be A and B matrices
         A = np.array(args[0], ndmin=2, dtype=float)
         G = np.array(args[1], ndmin=2, dtype=float)
@@ -517,6 +520,9 @@ def dlqe(*args, **keywords):
         NN = np.array(args[index+2], ndmin=2, dtype=float)
         raise ControlNotImplemented("cross-covariance not yet implememented")
 
+    # Check dimensions of G (needed before calling care())
+    _check_shape("QN", QN, G.shape[1], G.shape[1])
+
     # Compute the result (dimension and symmetry checking done in dare())
     P, E, LT = dare(A.T, C.T, G @ QN @ G.T, RN, method=method,
                     B_s="C", Q_s="QN", R_s="RN", S_s="NN")
@@ -650,13 +656,17 @@ def lqr(*args, **keywords):
         # Call dlqr
         return dlqr(*args, **keywords)
 
-    try:
-        # If this works, we were (probably) passed a system as the
-        # first argument; extract A and B
+    # If we were passed a state space  system, use that to get system matrices
+    if isinstance(args[0], StateSpace):
         A = np.array(args[0].A, ndmin=2, dtype=float)
         B = np.array(args[0].B, ndmin=2, dtype=float)
         index = 1
-    except AttributeError:
+
+    elif isinstance(args[0], LTI):
+        # Don't allow other types of LTI systems
+        raise ControlArgument("LTI system must be in state space form")
+
+    else:
         # Arguments should be A and B matrices
         A = np.array(args[0], ndmin=2, dtype=float)
         B = np.array(args[1], ndmin=2, dtype=float)
@@ -746,13 +756,17 @@ def dlqr(*args, **keywords):
     if isinstance(args[0], LTI) and isctime(args[0], strict=True):
         raise ControlArgument("dsys must be discrete time (dt != 0)")
 
-    try:
-        # If this works, we were (probably) passed a system as the
-        # first argument; extract A and B
+    # If we were passed a state space  system, use that to get system matrices
+    if isinstance(args[0], StateSpace):
         A = np.array(args[0].A, ndmin=2, dtype=float)
         B = np.array(args[0].B, ndmin=2, dtype=float)
         index = 1
-    except AttributeError:
+
+    elif isinstance(args[0], LTI):
+        # Don't allow other types of LTI systems
+        raise ControlArgument("LTI system must be in state space form")
+
+    else:
         # Arguments should be A and B matrices
         A = np.array(args[0], ndmin=2, dtype=float)
         B = np.array(args[1], ndmin=2, dtype=float)
diff --git a/control/tests/statefbk_test.py b/control/tests/statefbk_test.py
@@ -345,16 +345,18 @@ def test_DLQR_4args(self, matarrayin, matarrayout):
         K, S, poles = dlqr(A, B, Q, R)
         self.check_DLQR(K, S, poles, Q, R)
 
-    def test_lqr_badmethod(self):
+    @pytest.mark.parametrize("cdlqr", [lqr, dlqr])
+    def test_lqr_badmethod(self, cdlqr):
         A, B, Q, R = 0, 1, 10, 2
         with pytest.raises(ControlArgument, match="Unknown method"):
-            K, S, poles = lqr(A, B, Q, R, method='nosuchmethod')
+            K, S, poles = cdlqr(A, B, Q, R, method='nosuchmethod')
 
-    def test_lqr_slycot_not_installed(self):
+    @pytest.mark.parametrize("cdlqr", [lqr, dlqr])
+    def test_lqr_slycot_not_installed(self, cdlqr):
         A, B, Q, R = 0, 1, 10, 2
         if not slycot_check():
             with pytest.raises(ControlSlycot, match="Can't find slycot"):
-                K, S, poles = lqr(A, B, Q, R, method='slycot')
+                K, S, poles = cdlqr(A, B, Q, R, method='slycot')
 
     @pytest.mark.xfail(reason="warning not implemented")
     def testLQR_warning(self):
@@ -375,37 +377,49 @@ def testLQR_warning(self):
         with pytest.warns(UserWarning):
             (K, S, E) = lqr(A, B, Q, R, N)
 
-    def test_lqr_call_format(self):
+    @pytest.mark.parametrize("cdlqr", [lqr, dlqr])
+    def test_lqr_call_format(self, cdlqr):
         # Create a random state space system for testing
         sys = rss(2, 3, 2)
+        sys.dt = None           # treat as either continuous or discrete time
 
         # Weighting matrices
         Q = np.eye(sys.nstates)
         R = np.eye(sys.ninputs)
         N = np.zeros((sys.nstates, sys.ninputs))
 
         # Standard calling format
-        Kref, Sref, Eref = lqr(sys.A, sys.B, Q, R)
+        Kref, Sref, Eref = cdlqr(sys.A, sys.B, Q, R)
 
         # Call with system instead of matricees
-        K, S, E = lqr(sys, Q, R)
+        K, S, E = cdlqr(sys, Q, R)
         np.testing.assert_array_almost_equal(Kref, K)
         np.testing.assert_array_almost_equal(Sref, S)
         np.testing.assert_array_almost_equal(Eref, E)
 
         # Pass a cross-weighting matrix
-        K, S, E = lqr(sys, Q, R, N)
+        K, S, E = cdlqr(sys, Q, R, N)
         np.testing.assert_array_almost_equal(Kref, K)
         np.testing.assert_array_almost_equal(Sref, S)
         np.testing.assert_array_almost_equal(Eref, E)
 
         # Inconsistent system dimensions
         with pytest.raises(ct.ControlDimension, match="Incompatible dimen"):
-            K, S, E = lqr(sys.A, sys.C, Q, R)
+            K, S, E = cdlqr(sys.A, sys.C, Q, R)
 
-        # incorrect covariance matrix dimensions
+        # Incorrect covariance matrix dimensions
         with pytest.raises(ct.ControlDimension, match="Q must be a square"):
-            K, S, E = lqr(sys.A, sys.B, sys.C, R, Q)
+            K, S, E = cdlqr(sys.A, sys.B, sys.C, R, Q)
+
+        # Too few input arguments
+        with pytest.raises(ct.ControlArgument, match="not enough input"):
+            K, S, E = cdlqr(sys.A, sys.B)
+
+        # First argument is the wrong type (use SISO for non-slycot tests)
+        sys_tf = tf(rss(3, 1, 1))
+        sys_tf.dt = None        # treat as either continuous or discrete time
+        with pytest.raises(ct.ControlArgument, match="LTI system must be"):
+            K, S, E = cdlqr(sys_tf, Q, R)
 
     @pytest.mark.xfail(reason="warning not implemented")
     def testDLQR_warning(self):
@@ -443,41 +457,47 @@ def test_LQE(self, matarrayin, method):
         L, P, poles = lqe(A, G, C, QN, RN, method=method)
         self.check_LQE(L, P, poles, G, QN, RN)
 
-    def test_lqe_call_format(self):
+    @pytest.mark.parametrize("cdlqe", [lqe, dlqe])
+    def test_lqe_call_format(self, cdlqe):
         # Create a random state space system for testing
         sys = rss(4, 3, 2)
+        sys.dt = None           # treat as either continuous or discrete time
 
         # Covariance matrices
         Q = np.eye(sys.ninputs)
         R = np.eye(sys.noutputs)
         N = np.zeros((sys.ninputs, sys.noutputs))
 
         # Standard calling format
-        Lref, Pref, Eref = lqe(sys.A, sys.B, sys.C, Q, R)
+        Lref, Pref, Eref = cdlqe(sys.A, sys.B, sys.C, Q, R)
 
         # Call with system instead of matricees
-        L, P, E = lqe(sys, Q, R)
+        L, P, E = cdlqe(sys, Q, R)
         np.testing.assert_array_almost_equal(Lref, L)
         np.testing.assert_array_almost_equal(Pref, P)
         np.testing.assert_array_almost_equal(Eref, E)
 
-        # Compare state space and transfer function (SISO only)
-        sys_siso = rss(4, 1, 1)
-        L_ss, P_ss, E_ss = lqe(sys_siso, np.eye(1), np.eye(1))
-        L_tf, P_tf, E_tf = lqe(tf(sys_siso), np.eye(1), np.eye(1))
-        np.testing.assert_array_almost_equal(np.sort(E_ss), np.sort(E_tf))
-
         # Make sure we get an error if we specify N
         with pytest.raises(ct.ControlNotImplemented):
-            L, P, E = lqe(sys, Q, R, N)
+            L, P, E = cdlqe(sys, Q, R, N)
 
         # Inconsistent system dimensions
         with pytest.raises(ct.ControlDimension, match="Incompatible"):
-            L, P, E = lqe(sys.A, sys.C, sys.B, Q, R)
+            L, P, E = cdlqe(sys.A, sys.C, sys.B, Q, R)
 
-        # incorrect covariance matrix dimensions
+        # Incorrect covariance matrix dimensions
         with pytest.raises(ct.ControlDimension, match="Incompatible"):
-            L, P, E = lqe(sys.A, sys.B, sys.C, R, Q)
+            L, P, E = cdlqe(sys.A, sys.B, sys.C, R, Q)
+
+        # Too few input arguments
+        with pytest.raises(ct.ControlArgument, match="not enough input"):
+            L, P, E = cdlqe(sys.A, sys.C)
+
+        # First argument is the wrong type (use SISO for non-slycot tests)
+        sys_tf = tf(rss(3, 1, 1))
+        sys_tf.dt = None        # treat as either continuous or discrete time
+        with pytest.raises(ct.ControlArgument, match="LTI system must be"):
+            L, P, E = cdlqe(sys_tf, Q, R)
 
     def check_DLQE(self, L, P, poles, G, QN, RN):
         P_expected = asmatarrayout(G.dot(QN).dot(G))
