implement FrequencyResponseList class with plot() method

murrayrm · murrayrm · commit 2d6a779247d5 · 2023-09-16T09:49:26.000-07:00
diff --git a/control/frdata.py b/control/frdata.py
@@ -211,8 +211,9 @@ def __init__(self, *args, **kwargs):
         self.sysname = kwargs.pop('sysname', None)
 
         # Keep track of default properties for plotting
-        self.plot_phase=kwargs.pop('plot_phase', None)
-        self.title=kwargs.pop('title', None)
+        self.plot_phase = kwargs.pop('plot_phase', None)
+        self.title = kwargs.pop('title', None)
+        self.plot_type = kwargs.pop('plot_type', 'bode')
 
         # Keep track of return type
         self.return_magphase=kwargs.pop('return_magphase', False)
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -111,9 +111,41 @@
     'freqplot.share_frequency': 'col',
 }
 
+#
+# Frequency response data list class
+#
+# This class is a subclass of list that adds a plot() method, enabling
+# direct plotting from routines returning a list of FrequencyResponseData
+# objects.
+#
+
+class FrequencyResponseList(list):
+    def plot(self, *args, plot_type=None, **kwargs):
+        if plot_type == None:
+            for response in self:
+                if plot_type is not None and response.plot_type != plot_type:
+                    raise TypeError(
+                        "inconsistent plot_types in data; set plot_type "
+                        "to 'bode', 'svplot', or 'nyquist'")
+                plot_type = response.plot_type
+
+        if plot_type == 'bode':
+            bode_plot(self, *args, **kwargs)
+        elif plot_type == 'svplot':
+            singular_values_plot(self, *args, **kwargs)
+        elif plot_type == 'nyquist':
+            # nyquist_plot(self, *args, **kwargs)
+            raise NotImplementedError("Nyquist plots not yet supported")
+        else:
+            raise ValueError(f"unknown plot type '{plot_type}'")
+
 #
 # Bode plot
 #
+# This is the default method for plotting frequency responses.  There are
+# lots of options available for tuning the format of the plot, (hopefully)
+# covering most of the common use cases.
+#
 
 def bode_plot(
         data, omega=None, *fmt, ax=None, omega_limits=None, omega_num=None,
@@ -653,7 +685,6 @@ def _make_line_label(response, output_index, input_index):
         label += ", " if label != "" else ""
         label += f"{response.sysname}"
 
-        print(label)
         return label
 
     for index, response in enumerate(data):
@@ -1927,14 +1958,14 @@ def singular_values_response(
         svd_responses.append(
             FrequencyResponseData(
                 sigma_fresp, response.omega, _return_singvals=True,
-                outputs=[f'$\\sigma_{k}$' for k in range(sigma.shape[0])],
+                outputs=[f'$\\sigma_{{{k+1}}}$' for k in range(sigma.shape[0])],
                 inputs='inputs', dt=response.dt, plot_phase=False,
-                sysname=response.sysname,
+                sysname=response.sysname, plot_type='svplot',
                 title=f"Singular values for {response.sysname}"))
 
     # Return the responses in the same form that we received the systems
     if isinstance(sys, (list, tuple)):
-        return svd_responses
+        return FrequencyResponseList(svd_responses)
     else:
         return svd_responses[0]
 
@@ -2017,6 +2048,11 @@ def singular_values_plot(
     else:
         plot = True
 
+    # Warn the user if we got past something that is not real-valued
+    if any([not np.allclose(np.imag(response.fresp[:, 0, :]), 0)
+            for response in responses]):
+        warnings.warn("data has non-zero imaginary component")
+
     # Extract the data we need for plotting
     sigmas = [np.real(response.fresp[:, 0, :]) for response in responses]
     omegas = [response.omega for response in responses]
diff --git a/control/lti.py b/control/lti.py
@@ -119,8 +119,8 @@ def frequency_response(self, omega=None, squeeze=None):
         # Return the data as a frequency response data object
         response = self(s)
         return FrequencyResponseData(
-            response, omega, return_magphase=True, squeeze=squeeze, dt=self.dt,
-            sysname=self.name)
+            response, omega, return_magphase=True, squeeze=squeeze,
+            dt=self.dt, sysname=self.name, plot_type='bode')
 
     def dcgain(self):
         """Return the zero-frequency gain"""
@@ -470,8 +470,12 @@ def frequency_response(
 
         # Compute the frequency response
         responses.append(sys_.frequency_response(omega_sys, squeeze=squeeze))
-        
-    return responses if isinstance(sys, (list, tuple)) else responses[0]
+
+    if isinstance(sys, (list, tuple)):
+        from .freqplot import FrequencyResponseList
+        return FrequencyResponseList(responses)
+    else:
+        return responses[0]
 
 # Alternative name (legacy)
 def freqresp(sys, omega):
