Add slicing access for state-space models with tests

guptavaibhav0 · guptavaibhav0 · commit 6b61ed06a4ca · 2024-06-26T18:47:39.000+02:00
diff --git a/control/statesp.py b/control/statesp.py
@@ -50,6 +50,7 @@
 import math
 from copy import deepcopy
 from warnings import warn
+from collections.abc import Iterable
 
 import numpy as np
 import scipy as sp
@@ -1215,17 +1216,16 @@ def append(self, other):
 
     def __getitem__(self, indices):
         """Array style access"""
-        if len(indices) != 2:
+        if not isinstance(indices, Iterable) or len(indices) != 2:
             raise IOError('must provide indices of length 2 for state space')
-        outdx = indices[0] if isinstance(indices[0], list) else [indices[0]]
-        inpdx = indices[1] if isinstance(indices[1], list) else [indices[1]]
+        outdx, inpdx = indices
+        if not isinstance(outdx, (int, slice)) or not isinstance(inpdx, (int, slice)):
+            raise TypeError(f"system indices must be integers or slices")
         sysname = config.defaults['iosys.indexed_system_name_prefix'] + \
             self.name + config.defaults['iosys.indexed_system_name_suffix']
         return StateSpace(
             self.A, self.B[:, inpdx], self.C[outdx, :], self.D[outdx, inpdx],
-            self.dt, name=sysname,
-            inputs=[self.input_labels[i] for i in list(inpdx)],
-            outputs=[self.output_labels[i] for i in list(outdx)])
+            self.dt, name=sysname, inputs=self.input_labels[inpdx], outputs=self.output_labels[outdx])
 
     def sample(self, Ts, method='zoh', alpha=None, prewarp_frequency=None,
                name=None, copy_names=True, **kwargs):
diff --git a/control/tests/statesp_test.py b/control/tests/statesp_test.py
@@ -463,24 +463,38 @@ def test_append_tf(self):
         np.testing.assert_array_almost_equal(sys3c.A[:3, 3:], np.zeros((3, 2)))
         np.testing.assert_array_almost_equal(sys3c.A[3:, :3], np.zeros((2, 3)))
 
-    def test_array_access_ss(self):
-
+    def test_array_access_ss_failure(self):
+        sys1 = StateSpace(
+            [[1., 2.], [3., 4.]],
+            [[5., 6.], [6., 8.]],
+            [[9., 10.], [11., 12.]],
+            [[13., 14.], [15., 16.]], 1,
+            inputs=['u0', 'u1'], outputs=['y0', 'y1'])
+        with pytest.raises(IOError):
+            sys1[0]
+
+    @pytest.mark.parametrize("outdx, inpdx", 
+                             [(0, 1),
+                              (slice(0, 1, 1), 1),
+                              (0, slice(1, 2, 1)),
+                              (slice(0, 1, 1), slice(1, 2, 1))])
+    def test_array_access_ss(self, outdx, inpdx):
         sys1 = StateSpace(
             [[1., 2.], [3., 4.]],
             [[5., 6.], [6., 8.]],
             [[9., 10.], [11., 12.]],
             [[13., 14.], [15., 16.]], 1,
             inputs=['u0', 'u1'], outputs=['y0', 'y1'])
 
-        sys1_01 = sys1[0, 1]
+        sys1_01 = sys1[outdx, inpdx]
         np.testing.assert_array_almost_equal(sys1_01.A,
                                              sys1.A)
         np.testing.assert_array_almost_equal(sys1_01.B,
                                              sys1.B[:, 1:2])
         np.testing.assert_array_almost_equal(sys1_01.C,
                                              sys1.C[0:1, :])
         np.testing.assert_array_almost_equal(sys1_01.D,
-                                             sys1.D[0, 1])
+                                             sys1.D[0:1, 1:2])
 
         assert sys1.dt == sys1_01.dt
         assert sys1_01.input_labels == ['u1']
diff --git a/control/xferfcn.py b/control/xferfcn.py
@@ -47,6 +47,8 @@
 
 """
 
+from collections.abc import Iterable
+
 # External function declarations
 import numpy as np
 from numpy import angle, array, empty, finfo, ndarray, ones, \
@@ -758,7 +760,12 @@ def __pow__(self, other):
             return (TransferFunction([1], [1]) / self) * (self**(other + 1))
 
     def __getitem__(self, key):
+        if not isinstance(key, Iterable) or len(key) != 2:
+            raise IOError('must provide indices of length 2 for state space')
+        
         key1, key2 = key
+        if not isinstance(key1, (int, slice)) or not isinstance(key2, (int, slice)):
+            raise TypeError(f"system indices must be integers or slices")
 
         # pre-process
         if isinstance(key1, int):
