implement NamedSignal's for frequency responses

murrayrm · murrayrm · commit b382abf66403 · 2024-11-30T10:29:15.000-08:00
diff --git a/control/frdata.py b/control/frdata.py
@@ -20,7 +20,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, \
+    common_timebase
 from .lti import LTI, _process_frequency_response
 
 __all__ = ['FrequencyResponseData', 'FRD', 'frd']
@@ -243,19 +244,22 @@ def __init__(self, *args, **kwargs):
 
     @property
     def magnitude(self):
-        return np.abs(self.fresp)
+        return NamedSignal(
+            np.abs(self.fresp), self.output_labels, self.input_labels)
 
     @property
     def phase(self):
-        return np.angle(self.fresp)
+        return NamedSignal(
+            np.angle(self.fresp), self.output_labels, self.input_labels)
 
     @property
     def frequency(self):
         return self.omega
 
     @property
     def response(self):
-        return self.fresp
+        return NamedSignal(
+            self.fresp, self.output_labels, self.input_labels)
 
     def __str__(self):
         """String representation of the transfer function."""
diff --git a/control/tests/frd_test.py b/control/tests/frd_test.py
@@ -609,3 +609,49 @@ def test_frequency_response():
     assert mag_nosq_sq.shape == mag_default.shape
     assert phase_nosq_sq.shape == phase_default.shape
     assert omega_nosq_sq.shape == omega_default.shape
+
+
+def test_signal_labels():
+    # Create a system response for a SISO system
+    sys = ct.rss(4, 1, 1)
+    fresp = ct.frequency_response(sys)
+
+    # Make sure access via strings works
+    np.testing.assert_equal(
+        fresp.magnitude['y[0]'], fresp.magnitude[0])
+    np.testing.assert_equal(
+        fresp.phase['y[0]'], fresp.phase[0])
+
+    # Make sure errors are generated if key is unknown
+    with pytest.raises(ValueError, match="unknown signal name 'bad'"):
+        fresp.magnitude['bad']
+
+    # Create a system response for a MIMO system
+    sys = ct.rss(4, 2, 2)
+    fresp = ct.frequency_response(sys)
+
+    # Make sure access via strings works
+    np.testing.assert_equal(
+        fresp.magnitude['y[0]', 'u[1]'],
+        fresp.magnitude[0, 1])
+    np.testing.assert_equal(
+        fresp.phase['y[0]', 'u[1]'],
+        fresp.phase[0, 1])
+    np.testing.assert_equal(
+        fresp.response['y[0]', 'u[1]'],
+        fresp.response[0, 1])
+
+    # Make sure access via lists of strings works
+    np.testing.assert_equal(
+        fresp.response[['y[1]', 'y[0]'], 'u[0]'],
+        fresp.response[[1, 0], 0])
+
+    # Make sure errors are generated if key is unknown
+    with pytest.raises(ValueError, match="unknown signal name 'bad'"):
+        fresp.magnitude['bad']
+
+    with pytest.raises(ValueError, match="unknown signal name 'bad'"):
+        fresp.response[['y[1]', 'bad']]
+
+    with pytest.raises(ValueError, match=r"unknown signal name 'y\[0\]'"):
+        fresp.response['y[1]', 'y[0]']         # second index = input name
