add string-based indexing of FRD systems

murrayrm · murrayrm · commit cca4d13ef7ed · 2024-11-30T21:54:54.000-08:00
diff --git a/control/frdata.py b/control/frdata.py
@@ -10,6 +10,7 @@
 FRD data.
 """
 
+from collections.abc import Iterable
 from copy import copy
 from warnings import warn
 
@@ -21,7 +22,7 @@
 from . import config
 from .exception import pandas_check
 from .iosys import InputOutputSystem, NamedSignal, _process_iosys_keywords, \
-    common_timebase
+    _process_subsys_index, common_timebase
 from .lti import LTI, _process_frequency_response
 
 __all__ = ['FrequencyResponseData', 'FRD', 'frd']
@@ -597,9 +598,25 @@ def __iter__(self):
             return iter((self.omega, fresp))
         return iter((np.abs(fresp), np.angle(fresp), self.omega))
 
-    # Implement (thin) getitem to allow access via legacy indexing
-    def __getitem__(self, index):
-        return list(self.__iter__())[index]
+    def __getitem__(self, key):
+        if not isinstance(key, Iterable) or len(key) != 2:
+            # Implement (thin) getitem to allow access via legacy indexing
+            return list(self.__iter__())[key]
+
+        # Convert signal names to integer offsets (via NamedSignal object)
+        iomap = NamedSignal(
+            self.fresp[:, :, 0], self.output_labels, self.input_labels)
+        indices = iomap._parse_key(key)
+        outdx, outputs = _process_subsys_index(indices[0], self.output_labels)
+        inpdx, inputs = _process_subsys_index(indices[1], self.input_labels)
+
+        # Create the system name
+        sysname = config.defaults['iosys.indexed_system_name_prefix'] + \
+            self.name + config.defaults['iosys.indexed_system_name_suffix']
+
+        return FrequencyResponseData(
+            self.fresp[outdx, :][:, inpdx], self.omega, self.dt,
+            inputs=inputs, outputs=outputs, name=sysname)
 
     # Implement (thin) len to emulate legacy testing interface
     def __len__(self):
diff --git a/control/tests/lti_test.py b/control/tests/lti_test.py
@@ -303,3 +303,49 @@ def test_squeeze_exceptions(self, fcn):
             sys([[0.1j, 1j], [1j, 10j]])
         with pytest.raises(ValueError, match="must be 1D"):
             evalfr(sys, [[0.1j, 1j], [1j, 10j]])
+
+
+@pytest.mark.parametrize(
+    "outdx, inpdx, key",
+    [('y[0]', 'u[1]', (0, 1)),
+     (['y[0]'], ['u[1]'], (0, 1)),
+     (slice(0, 1, 1), slice(1, 2, 1), (0, 1)),
+     (['y[0]', 'y[1]'], ['u[1]', 'u[2]'], ([0, 1], [1, 2])),
+     ([0, 'y[1]'], ['u[1]', 2], ([0, 1], [1, 2])),
+     (slice(0, 2, 1), slice(1, 3, 1), ([0, 1], [1, 2])),
+     (['y[2]', 'y[1]'], ['u[2]', 'u[0]'], ([2, 1], [2, 0])),
+    ])
+@pytest.mark.parametrize("fcn", [ct.ss, ct.tf, ct.frd])
+def test_subsys_indexing(fcn, outdx, inpdx, key):
+    # Construct the base system and subsystem
+    sys = ct.rss(4, 3, 3)
+    subsys = sys[key]
+
+    # Construct the system to be test
+    match fcn:
+        case ct.frd:
+            omega = np.logspace(-1, 1)
+            sys = fcn(sys, omega)
+            subsys_chk = fcn(subsys, omega)
+        case _:
+            sys = fcn(sys)
+            subsys_chk = fcn(subsys)
+
+    # Construct the subsystem
+    subsys_fcn = sys[outdx, inpdx]
+
+    # Check to make sure everythng matches up
+    match fcn:
+        case ct.frd:
+            np.testing.assert_almost_equal(
+                subsys_fcn.response, subsys_chk.response)
+        case ct.ss:
+            np.testing.assert_almost_equal(subsys_fcn.A, subsys_chk.A)
+            np.testing.assert_almost_equal(subsys_fcn.B, subsys_chk.B)
+            np.testing.assert_almost_equal(subsys_fcn.C, subsys_chk.C)
+            np.testing.assert_almost_equal(subsys_fcn.D, subsys_chk.D)
+        case ct.tf:
+            omega = np.logspace(-1, 1)
+            np.testing.assert_almost_equal(
+                subsys_fcn.frequency_response(omega).response,
+                subsys_chk.frequency_response(omega).response)
