add label keyword to singular_value_plot and nyquist_plot

murrayrm · murrayrm · commit 0b6348f44819 · 2024-06-26T08:50:50.000-07:00
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -1501,7 +1501,7 @@ def nyquist_response(
 
 
 def nyquist_plot(
-        data, omega=None, plot=None, label_freq=0, color=None,
+        data, omega=None, plot=None, label_freq=0, color=None, label=None,
         return_contour=None, title=None, legend_loc='upper right', **kwargs):
     """Nyquist plot for a system.
 
@@ -1590,6 +1590,11 @@ def nyquist_plot(
         imaginary axis. Portions of the Nyquist plot corresponding to indented
         portions of the contour are plotted using a different line style.
 
+    label : str or array-like of str
+        If present, replace automatically generated label(s) with the given
+        label(s).  If sysdata is a list, strings should be specified for each
+        system.
+
     label_freq : int, optiona
         Label every nth frequency on the plot.  If not specified, no labels
         are generated.
@@ -1739,6 +1744,9 @@ def _parse_linestyle(style_name, allow_false=False):
     if not isinstance(data, (list, tuple)):
         data = [data]
 
+    # Process label keyword
+    line_labels = _process_line_labels(label, len(data))
+
     # If we are passed a list of systems, compute response first
     if all([isinstance(
             sys, (StateSpace, TransferFunction, FrequencyResponseData))
@@ -1804,12 +1812,14 @@ def _parse_linestyle(style_name, allow_false=False):
             reg_mask, abs(resp) > max_curve_magnitude)
         resp[rescale] *= max_curve_magnitude / abs(resp[rescale])
 
+        # Get the label to use for the line
+        label = response.sysname if line_labels is None else line_labels[idx]
+
         # Plot the regular portions of the curve (and grab the color)
         x_reg = np.ma.masked_where(reg_mask, resp.real)
         y_reg = np.ma.masked_where(reg_mask, resp.imag)
         p = plt.plot(
-            x_reg, y_reg, primary_style[0], color=color,
-            label=response.sysname, **kwargs)
+            x_reg, y_reg, primary_style[0], color=color, label=label, **kwargs)
         c = p[0].get_color()
         out[idx] += p
 
@@ -2211,7 +2221,7 @@ def singular_values_response(
 
 def singular_values_plot(
         data, omega=None, *fmt, plot=None, omega_limits=None, omega_num=None,
-        title=None, legend_loc='center right', **kwargs):
+        label=None, title=None, legend_loc='center right', **kwargs):
     """Plot the singular values for a system.
 
     Plot the singular values as a function of frequency for a system or
@@ -2257,6 +2267,10 @@ def singular_values_plot(
     grid : bool
         If True, plot grid lines on gain and phase plots.  Default is set by
         `config.defaults['freqplot.grid']`.
+    label : str or array-like of str
+        If present, replace automatically generated label(s) with the given
+        label(s).  If sysdata is a list, strings should be specified for each
+        system.
     omega_limits : array_like of two values
         Set limits for plotted frequency range. If Hz=True the limits
         are in Hz otherwise in rad/s.
@@ -2306,6 +2320,9 @@ def singular_values_plot(
 
         responses = data
 
+    # Process label keyword
+    line_labels = _process_line_labels(label, len(data))
+
     # Process (legacy) plot keyword
     if plot is not None:
         warnings.warn(
@@ -2385,11 +2402,14 @@ def singular_values_plot(
             with plt.rc_context(freqplot_rcParams):
                 out[idx_sys] = ax_sigma.semilogx(
                     omega, 20 * np.log10(sigma), *fmt,
-                    label=sysname, **color_arg, **kwargs)
+                    label=label, **color_arg, **kwargs)
         else:
             with plt.rc_context(freqplot_rcParams):
                 out[idx_sys] = ax_sigma.loglog(
-                    omega, sigma, label=sysname, *fmt, **color_arg, **kwargs)
+                    omega, sigma, label=label, *fmt, **color_arg, **kwargs)
+
+        # Get the label to use for the line
+        label = sysname if line_labels is None else line_labels[idx]
 
         # Plot the Nyquist frequency
         if nyq_freq is not None:
@@ -2653,7 +2673,7 @@ def _get_line_labels(ax, use_color=True):
 
 
 # Turn label keyword into array indexed by trace, output, input
-def _process_line_labels(label, nsys, ninput, noutput):
+def _process_line_labels(label, nsys, ninputs=0, noutputs=0):
     if label is None:
         return None
 
@@ -2669,7 +2689,8 @@ def _process_line_labels(label, nsys, ninput, noutput):
     # Turn the data into a 3D array of appropriate shape
     # TODO: allow more sophisticated broadcasting
     try:
-        line_labels = line_labels.reshape(nsys, ninput, noutput)
+        if ninputs > 0 and noutputs > 0:
+            line_labels = line_labels.reshape(nsys, ninputs, noutputs)
     except:
         if line_labels.shape[0] != nsys:
             raise ValueError("number of labels must match number of traces")
diff --git a/control/tests/freqplot_test.py b/control/tests/freqplot_test.py
@@ -346,57 +346,60 @@ def _get_visible_limits(ax):
         _get_visible_limits(ax.reshape(-1)[0]), np.array([1, 100]))
 
 
-def test_freqplot_trace_labels():
+@pytest.mark.parametrize(
+    "plt_fcn", [ct.bode_plot, ct.singular_values_plot, ct.nyquist_plot])
+def test_bode_trace_labels(plt_fcn):
     sys1 = ct.rss(2, 1, 1, name='sys1')
     sys2 = ct.rss(3, 1, 1, name='sys2')
 
     # Make sure default labels are as expected
-    out = ct.bode_plot([sys1, sys2])
+    out = ct.plt_fcn([sys1, sys2])
     axs = ct.get_plot_axes(out)
     legend = axs[0, 0].get_legend().get_texts()
     assert legend[0].get_text() == 'sys1'
     assert legend[1].get_text() == 'sys2'
     plt.close()
 
     # Override labels all at once
-    out = ct.bode_plot([sys1, sys2], label=['line1', 'line2'])
+    out = ct.plt_fcn([sys1, sys2], label=['line1', 'line2'])
     axs = ct.get_plot_axes(out)
     legend = axs[0, 0].get_legend().get_texts()
     assert legend[0].get_text() == 'line1'
     assert legend[1].get_text() == 'line2'
     plt.close()
 
     # Override labels one at a time
-    out = ct.bode_plot(sys1, label='line1')
-    out = ct.bode_plot(sys2, label='line2')
+    out = ct.plt_fcn(sys1, label='line1')
+    out = ct.plt_fcn(sys2, label='line2')
     axs = ct.get_plot_axes(out)
     legend = axs[0, 0].get_legend().get_texts()
     assert legend[0].get_text() == 'line1'
     assert legend[1].get_text() == 'line2'
     plt.close()
 
-    # Multi-dimensional data
-    sys1 = ct.rss(2, 2, 2, name='sys1')
-    sys2 = ct.rss(3, 2, 2, name='sys2')
-
-    # Check out some errors first
-    with pytest.raises(ValueError, match="number of labels must match"):
-        ct.bode_plot([sys1, sys2], label=['line1'])
-    with pytest.raises(ValueError, match="labels must be given for each"):
-        ct.bode_plot(sys1, label=['line1'])
-
-    # Now do things that should work
-    out = ct.bode_plot(
-        [sys1, sys2],
-        label=[
-            [['line1', 'line1'], ['line1', 'line1']],
-            [['line2', 'line2'], ['line2', 'line2']],
-        ])
-    axs = ct.get_plot_axes(out)
-    legend = axs[0, -1].get_legend().get_texts()
-    assert legend[0].get_text() == 'line1'
-    assert legend[1].get_text() == 'line2'
-    plt.close()
+    if plt_fcn == ct.bode_plot:
+        # Multi-dimensional data
+        sys1 = ct.rss(2, 2, 2, name='sys1')
+        sys2 = ct.rss(3, 2, 2, name='sys2')
+
+        # Check out some errors first
+        with pytest.raises(ValueError, match="number of labels must match"):
+            ct.bode_plot([sys1, sys2], label=['line1'])
+        with pytest.raises(ValueError, match="labels must be given for each"):
+            ct.bode_plot(sys1, label=['line1'])
+
+        # Now do things that should work
+        out = ct.bode_plot(
+            [sys1, sys2],
+            label=[
+                [['line1', 'line1'], ['line1', 'line1']],
+                [['line2', 'line2'], ['line2', 'line2']],
+            ])
+        axs = ct.get_plot_axes(out)
+        legend = axs[0, -1].get_legend().get_texts()
+        assert legend[0].get_text() == 'line1'
+        assert legend[1].get_text() == 'line2'
+        plt.close()
 
 
 @pytest.mark.parametrize("plt_fcn", [ct.bode_plot, ct.singular_values_plot])
