add string-based indexing of transfer functions

murrayrm · murrayrm · commit e52fca6104e0 · 2024-11-30T21:52:56.000-08:00
diff --git a/control/tests/xferfcn_test.py b/control/tests/xferfcn_test.py
@@ -390,19 +390,20 @@ def test_pow(self):
         with pytest.raises(ValueError):
             TransferFunction.__pow__(sys1, 0.5)
 
-    def test_slice(self):
+    @pytest.mark.parametrize("named", [False, True])
+    def test_slice(self, named):
         sys = TransferFunction(
             [ [   [1],    [2],    [3]], [   [3],    [4],    [5]] ],
             [ [[1, 2], [1, 3], [1, 4]], [[1, 4], [1, 5], [1, 6]] ],
             inputs=['u0', 'u1', 'u2'], outputs=['y0', 'y1'], name='sys')
 
-        sys1 = sys[1:, 1:]
+        sys1 = sys[1:, 1:] if not named else sys['y1', ['u1', 'u2']]
         assert (sys1.ninputs, sys1.noutputs) == (2, 1)
         assert sys1.input_labels == ['u1', 'u2']
         assert sys1.output_labels == ['y1']
         assert sys1.name == 'sys$indexed'
 
-        sys2 = sys[:2, :2]
+        sys2 = sys[:2, :2] if not named else sys[['y0', 'y1'], ['u0', 'u1']]
         assert (sys2.ninputs, sys2.noutputs) == (2, 2)
         assert sys2.input_labels == ['u0', 'u1']
         assert sys2.output_labels == ['y0', 'y1']
@@ -411,7 +412,7 @@ def test_slice(self):
         sys = TransferFunction(
             [ [   [1],    [2],    [3]], [   [3],    [4],    [5]] ],
             [ [[1, 2], [1, 3], [1, 4]], [[1, 4], [1, 5], [1, 6]] ], 0.5)
-        sys1 = sys[1:, 1:]
+        sys1 = sys[1:, 1:] if not named else sys[['y[1]'], ['u[1]', 'u[2]']]
         assert (sys1.ninputs, sys1.noutputs) == (2, 1)
         assert sys1.dt == 0.5
         assert sys1.input_labels == ['u[1]', 'u[2]']
diff --git a/control/xferfcn.py b/control/xferfcn.py
@@ -48,25 +48,26 @@
 """
 
 from collections.abc import Iterable
+from copy import deepcopy
+from itertools import chain
+from re import sub
+from warnings import warn
 
 # External function declarations
 import numpy as np
-from numpy import angle, array, empty, finfo, ndarray, ones, \
-    polyadd, polymul, polyval, roots, sqrt, zeros, squeeze, exp, pi, \
-    where, delete, real, poly, nonzero
 import scipy as sp
-from scipy.signal import tf2zpk, zpk2tf, cont2discrete
+from numpy import angle, array, delete, empty, exp, finfo, ndarray, nonzero, \
+    ones, pi, poly, polyadd, polymul, polyval, real, roots, sqrt, squeeze, \
+    where, zeros
 from scipy.signal import TransferFunction as signalTransferFunction
-from copy import deepcopy
-from warnings import warn
-from itertools import chain
-from re import sub
-from .lti import LTI, _process_frequency_response
-from .iosys import InputOutputSystem, common_timebase, isdtime, \
-    _process_iosys_keywords
+from scipy.signal import cont2discrete, tf2zpk, zpk2tf
+
+from . import config
 from .exception import ControlMIMONotImplemented
 from .frdata import FrequencyResponseData
-from . import config
+from .iosys import InputOutputSystem, NamedSignal, _process_iosys_keywords, \
+    _process_subsys_index, common_timebase, isdtime
+from .lti import LTI, _process_frequency_response
 
 __all__ = ['TransferFunction', 'tf', 'zpk', 'ss2tf', 'tfdata']
 
@@ -761,48 +762,30 @@ def __pow__(self, other):
 
     def __getitem__(self, key):
         if not isinstance(key, Iterable) or len(key) != 2:
-            raise IOError('must provide indices of length 2 for transfer functions')
+            raise IOError(
+                "must provide indices of length 2 for transfer functions")
+
+        # Convert signal names to integer offsets (via NamedSignal object)
+        iomap = NamedSignal(
+            np.empty((self.noutputs, self.ninputs)),
+            self.output_labels, self.input_labels)
+        indices = iomap._parse_key(key)
+        outdx, outputs = _process_subsys_index(
+            indices[0], self.output_labels, slice_to_list=True)
+        inpdx, inputs = _process_subsys_index(
+            indices[1], self.input_labels, slice_to_list=True)
         
-        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):
-            key1 = slice(key1, key1 + 1, 1)
-        if isinstance(key2, int):
-            key2 = slice(key2, key2 + 1, 1)
-        # dim1
-        start1, stop1, step1 = key1.start, key1.stop, key1.step
-        if step1 is None:
-            step1 = 1
-        if start1 is None:
-            start1 = 0
-        if stop1 is None:
-            stop1 = len(self.num)
-        # dim1
-        start2, stop2, step2 = key2.start, key2.stop, key2.step
-        if step2 is None:
-            step2 = 1
-        if start2 is None:
-            start2 = 0
-        if stop2 is None:
-            stop2 = len(self.num[0])
-
+        # Construct the transfer function for the subsyste
         num, den = [], []
-        for i in range(start1, stop1, step1):
+        for i in outdx:
             num_i = []
             den_i = []
-            for j in range(start2, stop2, step2):
+            for j in inpdx:
                 num_i.append(self.num[i][j])
                 den_i.append(self.den[i][j])
             num.append(num_i)
             den.append(den_i)
 
-        # Save the label names
-        outputs = [self.output_labels[i] for i in range(start1, stop1, step1)]
-        inputs = [self.input_labels[j] for j in range(start2, stop2, step2)]
-
         # Create the system name
         sysname = config.defaults['iosys.indexed_system_name_prefix'] + \
             self.name + config.defaults['iosys.indexed_system_name_suffix']
