add additional tests for freqplot label override + small fixes

murrayrm · murrayrm · commit 606cf482e43d · 2024-06-26T08:50:50.000-07:00
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -2397,6 +2397,9 @@ def singular_values_plot(
         sysname = response.sysname if response.sysname is not None \
             else f"Unknown-{idx_sys}"
 
+        # Get the label to use for the line
+        label = sysname if line_labels is None else line_labels[idx_sys]
+
         # Plot the data
         if dB:
             with plt.rc_context(freqplot_rcParams):
@@ -2408,9 +2411,6 @@ def singular_values_plot(
                 out[idx_sys] = ax_sigma.loglog(
                     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:
             ax_sigma.axvline(
@@ -2678,7 +2678,7 @@ def _process_line_labels(label, nsys, ninputs=0, noutputs=0):
         return None
 
     if isinstance(label, str):
-        label = [[[label]]]
+        label = [label]
 
     # Convert to an ndarray, if not done aleady
     try:
@@ -2687,10 +2687,13 @@ def _process_line_labels(label, nsys, ninputs=0, noutputs=0):
         raise ValueError("label must be a string or array_like")
 
     # Turn the data into a 3D array of appropriate shape
-    # TODO: allow more sophisticated broadcasting
+    # TODO: allow more sophisticated broadcasting (and error checking)
     try:
         if ninputs > 0 and noutputs > 0:
-            line_labels = line_labels.reshape(nsys, ninputs, noutputs)
+            if line_labels.ndim == 1:
+                line_labels = line_labels.reshape(nsys, 1, 1)
+            line_labels = np.broadcast_to(
+                line_labels,(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,33 +346,66 @@ def _get_visible_limits(ax):
         _get_visible_limits(ax.reshape(-1)[0]), np.array([1, 100]))
 
 
+def test_gangof4_trace_labels():
+    P1 = ct.rss(2, 1, 1, name='P1')
+    P2 = ct.rss(3, 1, 1, name='P2')
+    C = ct.rss(1, 1, 1, name='C')
+
+    # Make sure default labels are as expected
+    out = ct.gangof4_response(P1, C).plot()
+    out = ct.gangof4_response(P2, C).plot()
+    axs = ct.get_plot_axes(out)
+    legend = axs[0, 1].get_legend().get_texts()
+    assert legend[0].get_text() == 'None'
+    assert legend[1].get_text() == 'None'
+    plt.close()
+
+    # Override labels
+    out = ct.gangof4_response(P1, C).plot(label='line1')
+    out = ct.gangof4_response(P2, C).plot(label='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, ct.nyquist_plot])
-def test_bode_trace_labels(plt_fcn):
+def test_freqplot_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.plt_fcn([sys1, sys2])
+    out = plt_fcn([sys1, sys2])
     axs = ct.get_plot_axes(out)
-    legend = axs[0, 0].get_legend().get_texts()
+    if axs.ndim == 1:
+        legend = axs[0].get_legend().get_texts()
+    else:
+        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.plt_fcn([sys1, sys2], label=['line1', 'line2'])
+    out = plt_fcn([sys1, sys2], label=['line1', 'line2'])
     axs = ct.get_plot_axes(out)
-    legend = axs[0, 0].get_legend().get_texts()
+    if axs.ndim == 1:
+        legend = axs[0].get_legend().get_texts()
+    else:
+        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.plt_fcn(sys1, label='line1')
-    out = ct.plt_fcn(sys2, label='line2')
+    out = plt_fcn(sys1, label='line1')
+    out = plt_fcn(sys2, label='line2')
     axs = ct.get_plot_axes(out)
-    legend = axs[0, 0].get_legend().get_texts()
+    if axs.ndim == 1:
+        legend = axs[0].get_legend().get_texts()
+    else:
+        legend = axs[0, 0].get_legend().get_texts()
     assert legend[0].get_text() == 'line1'
     assert legend[1].get_text() == 'line2'
     plt.close()
@@ -385,8 +418,11 @@ def test_bode_trace_labels(plt_fcn):
         # 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'])
+
+        with pytest.xfail(reason="need better broadcast checking on labels"):
+            with pytest.raises(
+                    ValueError, match="labels must be given for each"):
+                ct.bode_plot(sys1, overlay_inputs=True, label=['line1'])
 
         # Now do things that should work
         out = ct.bode_plot(
diff --git a/control/tests/kwargs_test.py b/control/tests/kwargs_test.py
@@ -193,9 +193,9 @@ def test_matplotlib_kwargs(function, nsysargs, moreargs, kwargs, mplcleanup):
 def test_response_plot_kwargs(data_fcn, plot_fcn, mimo):
     # Create a system for testing
     if mimo:
-        response = data_fcn(control.rss(4, 2, 2))
+        response = data_fcn(control.rss(4, 2, 2, strictly_proper=True))
     else:
-        response = data_fcn(control.rss(4, 1, 1))
+        response = data_fcn(control.rss(4, 1, 1, strictly_proper=True))
 
     # Make sure that calling the data function with unknown keyword errs
     with pytest.raises(
