murrayrm · Dec 4, 2024 · Nov 30, 2024 · Nov 30, 2024 · Nov 30, 2024 · Dec 1, 2024
diff --git a/control/exception.py b/control/exception.py
@@ -52,6 +52,10 @@ class ControlArgument(TypeError):
     """Raised when arguments to a function are not correct"""
     pass
 
+class ControlIndexError(IndexError):
+    """Raised when arguments to an indexed object are not correct"""
+    pass
+
 class ControlMIMONotImplemented(NotImplementedError):
     """Function is not currently implemented for MIMO systems"""
     pass

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
 
@@ -20,7 +21,8 @@
 
 from . import config
 from .exception import pandas_check
-from .iosys import InputOutputSystem, _process_iosys_keywords, common_timebase
+from .iosys import InputOutputSystem, NamedSignal, _process_iosys_keywords, \
+    _process_subsys_index, common_timebase
 from .lti import LTI, _process_frequency_response
 
 __all__ = ['FrequencyResponseData', 'FRD', 'frd']
@@ -33,8 +35,8 @@ class FrequencyResponseData(LTI):
 
     The FrequencyResponseData (FRD) class is used to represent systems in
     frequency response data form.  It can be created manually using the
-    class constructor, using the :func:~~control.frd` factory function
-    (preferred), or via the :func:`~control.frequency_response` function.
+    class constructor, using the :func:`~control.frd` factory function or
+    via the :func:`~control.frequency_response` function.
 
     Parameters
     ----------
@@ -65,6 +67,28 @@ class constructor, using the :func:~~control.frd` factory function
         frequency point.
     dt : float, True, or None
         System timebase.
+    squeeze : bool
+        By default, if a system is single-input, single-output (SISO) then
+        the outputs (and inputs) are returned as a 1D array (indexed by
+        frequency) and if a system is multi-input or multi-output, then the
+        outputs are returned as a 2D array (indexed by output and
+        frequency) or a 3D array (indexed by output, trace, and frequency).
+        If ``squeeze=True``, access to the output response will remove
+        single-dimensional entries from the shape of the inputs and outputs
+        even if the system is not SISO. If ``squeeze=False``, the output is
+        returned as a 3D array (indexed by the output, input, and
+        frequency) even if the system is SISO. The default value can be set
+        using config.defaults['control.squeeze_frequency_response'].
+    ninputs, noutputs, nstates : int
+        Number of inputs, outputs, and states of the underlying system.
+    input_labels, output_labels : array of str
+        Names for the input and output variables.
+    sysname : str, optional
+        Name of the system.  For data generated using
+        :func:`~control.frequency_response`, stores the name of the system
+        that created the data.
+    title : str, optional
+        Set the title to use when plotting.
 
     See Also
     --------
@@ -89,6 +113,20 @@ class constructor, using the :func:~~control.frd` factory function
     the imaginary access).  See :meth:`~control.FrequencyResponseData.__call__`
     for a more detailed description.
 
+    A state space system is callable and returns the value of the transfer
+    function evaluated at a point in the complex plane.  See
+    :meth:`~control.StateSpace.__call__` for a more detailed description.
+
+    Subsystem response corresponding to selected input/output pairs can be
+    created by indexing the frequency response data object::
+
+        subsys = sys[output_spec, input_spec]
+
+    The input and output specifications can be single integers, lists of
+    integers, or slices.  In addition, the strings representing the names
+    of the signals can be used and will be replaced with the equivalent
+    signal offsets.
+
     """
     #
     # Class attributes
@@ -243,21 +281,72 @@ def __init__(self, *args, **kwargs):
 
     @property
     def magnitude(self):
-        return np.abs(self.fresp)
+        """Magnitude of the frequency response.
+
+        Magnitude of the frequency response, indexed by either the output
+        and frequency (if only a single input is given) or the output,
+        input, and frequency (for multi-input systems).  See
+        :attr:`FrequencyResponseData.squeeze` for a description of how this
+        can be modified using the `squeeze` keyword.
+
+        Input and output signal names can be used to index the data in
+        place of integer offsets.
+
+        :type: 1D, 2D, or 3D array
+
+        """
+        return NamedSignal(
+            np.abs(self.fresp), self.output_labels, self.input_labels)
 
     @property
     def phase(self):
-        return np.angle(self.fresp)
+        """Phase of the frequency response.
+
+        Phase of the frequency response in radians/sec, indexed by either
+        the output and frequency (if only a single input is given) or the
+        output, input, and frequency (for multi-input systems).  See
+        :attr:`FrequencyResponseData.squeeze` for a description of how this
+        can be modified using the `squeeze` keyword.
+
+        Input and output signal names can be used to index the data in
+        place of integer offsets.
+
+        :type: 1D, 2D, or 3D array
+
+        """
+        return NamedSignal(
+            np.angle(self.fresp), self.output_labels, self.input_labels)
 
     @property
     def frequency(self):
+        """Frequencies at which the response is evaluated.
+
+        :type: 1D array
+
+        """
         return self.omega
 
     @property
     def response(self):
-        return self.fresp
+        """Complex value of the frequency response.
+
+        Value of the frequency response as a complex number, indexed by
+        either the output and frequency (if only a single input is given)
+        or the output, input, and frequency (for multi-input systems).  See
+        :attr:`FrequencyResponseData.squeeze` for a description of how this
+        can be modified using the `squeeze` keyword.
+
+        Input and output signal names can be used to index the data in
+        place of integer offsets.
+
+        :type: 1D, 2D, or 3D array
+
+        """
+        return NamedSignal(
+            self.fresp, self.output_labels, self.input_labels)
 
     def __str__(self):
+
         """String representation of the transfer function."""
 
         mimo = self.ninputs > 1 or self.noutputs > 1
@@ -593,9 +682,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, level=1)  # ignore index checks
+        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/iosys.py b/control/iosys.py
@@ -13,9 +13,10 @@
 import numpy as np
 
 from . import config
+from .exception import ControlIndexError
 
-__all__ = ['InputOutputSystem', 'issiso', 'timebase', 'common_timebase',
-           'isdtime', 'isctime']
+__all__ = ['InputOutputSystem', 'NamedSignal', 'issiso', 'timebase',
+           'common_timebase', 'isdtime', 'isctime']
 
 # Define module default parameter values
 _iosys_defaults = {
@@ -33,6 +34,68 @@
 }
 
 
+# Named signal class
+class NamedSignal(np.ndarray):
+    def __new__(cls, input_array, signal_labels=None, trace_labels=None):
+        # See https://numpy.org/doc/stable/user/basics.subclassing.html
+        obj = np.asarray(input_array).view(cls)     # Cast to our class type
+        obj.signal_labels = signal_labels           # Save signal labels
+        obj.trace_labels = trace_labels             # Save trace labels
+        obj.data_shape = input_array.shape          # Save data shape
+        return obj                                  # Return new object
+
+    def __array_finalize__(self, obj):
+        # See https://numpy.org/doc/stable/user/basics.subclassing.html
+        if obj is None: return
-        if obj is None: return
+        if obj is None:
+            return
-        if obj is None: return
+        if obj is None:
+            return
+        self.signal_labels = getattr(obj, 'signal_labels', None)
+        self.trace_labels = getattr(obj, 'trace_labels', None)
+        self.data_shape = getattr(obj, 'data_shape', None)
+
+    def _parse_key(self, key, labels=None, level=0):
+        if labels is None:
+            labels = self.signal_labels
+        try:
+            if isinstance(key, str):
+                key = labels.index(item := key)
+                if level == 0 and len(self.data_shape) < 2:
+                    raise ControlIndexError
+            elif isinstance(key, list):
+                keylist = []
+                for item in key:        # use for loop to save item for error
+                    keylist.append(
+                        self._parse_key(item, labels=labels, level=level+1))
+                if level == 0 and key != keylist and len(self.data_shape) < 2:
+                    raise ControlIndexError
+                key = keylist
+            elif isinstance(key, tuple) and len(key) > 0:
+                keylist = []
+                keylist.append(
+                    self._parse_key(
+                        item := key[0], labels=self.signal_labels,
+                        level=level+1))
+                if len(key) > 1:
+                    keylist.append(
+                        self._parse_key(
+                            item := key[1], labels=self.trace_labels,
+                            level=level+1))
+                if level == 0 and key[:len(keylist)] != tuple(keylist) \
+                   and len(keylist) > len(self.data_shape) - 1:
+                    raise ControlIndexError
+                for i in range(2, len(key)):
+                    keylist.append(key[i])      # pass on remaining elements
+                key = tuple(keylist)
+        except ValueError:
+            raise ValueError(f"unknown signal name '{item}'")
+        except ControlIndexError:
+            raise ControlIndexError(
+                "signal name(s) not valid for squeezed data")
+
+        return key
+
+    def __getitem__(self, key):
+        return super().__getitem__(self._parse_key(key))
+
+
 class InputOutputSystem(object):
     """A class for representing input/output systems.
 
@@ -965,3 +1028,30 @@ def _parse_spec(syslist, spec, signame, dictname=None):
             ValueError(f"signal index '{index}' is out of range")
 
     return system_index, signal_indices, gain
+
+
+#
+# Utility function for processing subsystem indices
+#
+# This function processes an index specification (int, list, or slice) and
+# returns a index specification that can be used to create a subsystem
+#
+def _process_subsys_index(idx, sys_labels, slice_to_list=False):
+    if not isinstance(idx, (slice, list, int)):
+        raise TypeError(f"system indices must be integers, slices, or lists")
+
+    # Convert singleton lists to integers for proper slicing (below)
+    if isinstance(idx, (list, tuple)) and len(idx) == 1:
+        idx = idx[0]
+
+    # Convert int to slice so that numpy doesn't drop dimension
+    if isinstance(idx, int): idx = slice(idx, idx+1, 1)
-    if isinstance(idx, int): idx = slice(idx, idx+1, 1)
+    if isinstance(idx, int):
+        idx = slice(idx, idx+1, 1)
-    if isinstance(idx, int): idx = slice(idx, idx+1, 1)
+    if isinstance(idx, int):
+        idx = slice(idx, idx+1, 1)
+
+    # Get label names (taking care of possibility that we were passed a list)
+    labels = [sys_labels[i] for i in idx] if isinstance(idx, list) \
+        else sys_labels[idx]
+
+    if slice_to_list and isinstance(idx, slice):
+        idx = range(len(sys_labels))[idx]
+
+    return idx, labels
diff --git a/control/nlsys.py b/control/nlsys.py
@@ -193,8 +193,8 @@ def __call__(sys, u, params=None, squeeze=None):
 
         # Evaluate the function on the argument
         out = sys._out(0, np.array((0,)), np.asarray(u))
-        _, out = _process_time_response(
-            None, out, issiso=sys.issiso(), squeeze=squeeze)
+        out = _process_time_response(
+            out, issiso=sys.issiso(), squeeze=squeeze)
         return out
 
     def __mul__(self, other):