update label processing to provide common functionality + unit tests

murrayrm · murrayrm · commit cec9e7053dfb · 2024-08-08T14:54:49.000-07:00
diff --git a/control/ctrlplot.py b/control/ctrlplot.py
@@ -384,7 +384,7 @@ def _process_ax_keyword(
 
 # Turn label keyword into array indexed by trace, output, input
 # TODO: move to ctrlutil.py and update parameter names to reflect general use
-def _process_line_labels(label, ntraces, ninputs=0, noutputs=0):
+def _process_line_labels(label, ntraces=1, ninputs=0, noutputs=0):
     if label is None:
         return None
 
@@ -447,7 +447,7 @@ def _make_legend_labels(labels, ignore_common=False):
     if last_space < 0 or ignore_common:
         common_prefix = ''
     elif last_space > 0:
-        common_prefix = common_prefix[:last_space]
+        common_prefix = common_prefix[:last_space + 2]
     prefix_len = len(common_prefix)
 
     # Look for a common suffix (up to a space)
diff --git a/control/descfcn.py b/control/descfcn.py
@@ -380,7 +380,7 @@ def _cost(x):
 
 
 def describing_function_plot(
-        *sysdata, label="%5.2g @ %-5.2g", **kwargs):
+        *sysdata, point_label="%5.2g @ %-5.2g", label=None, **kwargs):
     """describing_function_plot(data, *args, **kwargs)
 
     Plot a Nyquist plot with a describing function for a nonlinear system.
@@ -420,7 +420,7 @@ def describing_function_plot(
         If True (default), refine the location of the intersection of the
         Nyquist curve for the linear system and the describing function to
         determine the intersection point
-    label : str, optional
+    point_label : str, optional
         Formatting string used to label intersection points on the Nyquist
         plot.  Defaults to "%5.2g @ %-5.2g".  Set to `None` to omit labels.
 
@@ -453,6 +453,8 @@ def describing_function_plot(
     # Process keywords
     warn_nyquist = config._process_legacy_keyword(
         kwargs, 'warn', 'warn_nyquist', kwargs.pop('warn_nyquist', None))
+    point_label = config._process_legacy_keyword(
+        kwargs, 'label', 'point_label', point_label)
 
     # TODO: update to be consistent with ctrlplot use of `label`
     if label not in (False, None) and not isinstance(label, str):
@@ -484,10 +486,10 @@ def describing_function_plot(
     lines[1] = plt.plot(dfresp.N_vals.real, dfresp.N_vals.imag)
 
     # Label the intersection points
-    if label:
+    if point_label:
         for pos, (a, omega) in zip(dfresp.positions, dfresp.intersections):
             # Add labels to the intersection points
-            plt.text(pos.real, pos.imag, label % (a, omega))
+            plt.text(pos.real, pos.imag, point_label % (a, omega))
 
     return ControlPlot(lines, cfig.axes, cfig.figure)
 
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -2078,7 +2078,8 @@ def gangof4_response(
 
     return FrequencyResponseData(
         data, omega, outputs=['y', 'u'], inputs=['r', 'd'],
-        title=f"Gang of Four for P={P.name}, C={C.name}", plot_phase=False)
+        title=f"Gang of Four for P={P.name}, C={C.name}",
+        sysname=f"P={P.name}, C={C.name}", plot_phase=False)
 
 
 def gangof4_plot(
diff --git a/control/grid.py b/control/grid.py
@@ -104,8 +104,7 @@ def sgrid(subplot=(1, 1, 1), scaling=None):
     ax = SubplotHost(fig, *subplot, grid_helper=grid_helper)
 
     # make ticklabels of right invisible, and top axis visible.
-    visible = True
-    ax.axis[:].major_ticklabels.set_visible(visible)
+    ax.axis[:].major_ticklabels.set_visible(True)
     ax.axis[:].major_ticks.set_visible(False)
     ax.axis[:].invert_ticklabel_direction()
     ax.axis[:].major_ticklabels.set_color('gray')
diff --git a/control/nichols.py b/control/nichols.py
@@ -19,7 +19,7 @@
 
 from . import config
 from .ctrlplot import ControlPlot, _get_line_labels, _process_ax_keyword, \
-    suptitle
+    _process_line_labels, suptitle
 from .ctrlutil import unwrap
 from .freqplot import _default_frequency_range, _freqplot_defaults
 from .lti import frequency_response
@@ -36,7 +36,7 @@
 
 def nichols_plot(
         data, omega=None, *fmt, grid=None, title=None, ax=None,
-        legend_loc='upper left', **kwargs):
+        legend_loc='upper left', label=None, **kwargs):
     """Nichols plot for a system.
 
     Plots a Nichols plot for the system over a (optional) frequency range.
@@ -53,6 +53,10 @@ def nichols_plot(
         The `omega` parameter must be present (use omega=None if needed).
     grid : boolean, optional
         True if the plot should include a Nichols-chart grid. Default is True.
+    label : str or array-like of str
+        If present, replace automatically generated label(s) with given
+        label(s).  If sysdata is a list, strings should be specified for each
+        system.
     legend_loc : str, optional
         For plots with multiple lines, a legend will be included in the
         given location.  Default is 'upper left'.  Use False to supress.
@@ -61,7 +65,7 @@ def nichols_plot(
 
     Returns
     -------
-  cplt : :class:`ControlPlot` object
+    cplt : :class:`ControlPlot` object
         Object containing the data that were plotted:
 
           * cplt.lines: 1D array of :class:`matplotlib.lines.Line2D` objects.
@@ -81,6 +85,7 @@ def nichols_plot(
     """
     # Get parameter values
     grid = config._get_param('nichols', 'grid', grid, True)
+    label = _process_line_labels(label)
     rcParams = config._get_param(
         'freqplot', 'rcParams', kwargs, _freqplot_defaults, pop=True)
 
@@ -113,12 +118,13 @@ def nichols_plot(
         x = unwrap(np.degrees(phase), 360)
         y = 20*np.log10(mag)
 
-        # Decide on the system name
+        # Decide on the system name and label
         sysname = response.sysname if response.sysname is not None \
             else f"Unknown-{idx_sys}"
+        label_ = sysname if label is None else label[idx]
 
         # Generate the plot
-        out[idx] = ax_nichols.plot(x, y, *fmt, label=sysname, **kwargs)
+        out[idx] = ax_nichols.plot(x, y, *fmt, label=label_, **kwargs)
 
     # Label the plot axes
     plt.xlabel('Phase [deg]')
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -19,7 +19,7 @@
 
 from . import config
 from .ctrlplot import ControlPlot, suptitle, _get_line_labels, \
-    _process_ax_keyword
+    _process_ax_keyword, _process_line_labels
 from .freqplot import _freqplot_defaults
 from .grid import nogrid, sgrid, zgrid
 from .iosys import isctime, isdtime
@@ -173,7 +173,7 @@ def pole_zero_plot(
         data, plot=None, grid=None, title=None, marker_color=None,
         marker_size=None, marker_width=None, legend_loc='upper right',
         xlim=None, ylim=None, interactive=None, ax=None, scaling=None,
-        initial_gain=None, **kwargs):
+        initial_gain=None, label=None, **kwargs):
     """Plot a pole/zero map for a linear system.
 
     If the system data include root loci, a root locus diagram for the
@@ -230,25 +230,30 @@ def pole_zero_plot(
     scaling : str or list, optional
         Set the type of axis scaling.  Can be 'equal' (default), 'auto', or
         a list of the form [xmin, xmax, ymin, ymax].
-    title : str, optional
-        Set the title of the plot.  Defaults plot type and system name(s).
+    initial_gain : float, optional
+        If given, the specified system gain will be marked on the plot.
+
+    interactive : bool, optional
+        Turn off interactive mode for root locus plots.
+    label : str or array-like of str
+        If present, replace automatically generated label(s) with given
+        label(s).  If data is a list, strings should be specified for each
+        system.
+    legend_loc : str, optional
+        For plots with multiple lines, a legend will be included in the
+        given location.  Default is 'center right'.  Use False to supress.
     marker_color : str, optional
         Set the color of the markers used for poles and zeros.
     marker_size : int, optional
         Set the size of the markers used for poles and zeros.
     marker_width : int, optional
         Set the line width of the markers used for poles and zeros.
-    legend_loc : str, optional
-        For plots with multiple lines, a legend will be included in the
-        given location.  Default is 'center right'.  Use False to supress.
+    title : str, optional
+        Set the title of the plot.  Defaults plot type and system name(s).
     xlim : list, optional
         Set the limits for the x axis.
     ylim : list, optional
         Set the limits for the y axis.
-    interactive : bool, optional
-        Turn off interactive mode for root locus plots.
-    initial_gain : float, optional
-        If given, the specified system gain will be marked on the plot.
 
     Notes
     -----
@@ -260,13 +265,14 @@ def pole_zero_plot(
 
     """
     # Get parameter values
+    label = _process_line_labels(label)
     marker_size = config._get_param('pzmap', 'marker_size', marker_size, 6)
     marker_width = config._get_param('pzmap', 'marker_width', marker_width, 1.5)
-    xlim_user, ylim_user = xlim, ylim
     rcParams = config._get_param(
         'freqplot', 'rcParams', kwargs, _freqplot_defaults,
         pop=True, last=True)
     user_ax = ax
+    xlim_user, ylim_user = xlim, ylim
 
     # If argument was a singleton, turn it into a tuple
     if not isinstance(data, (list, tuple)):
@@ -371,12 +377,15 @@ def pole_zero_plot(
 
         # Plot the locations of the poles and zeros
         if len(poles) > 0:
-            label = response.sysname if response.loci is None else None
+            if label is None:
+                label_ = response.sysname if response.loci is None else None
+            else:
+                label_ = label[idx]
             out[idx, 0] = ax.plot(
                 real(poles), imag(poles), marker='x', linestyle='',
                 markeredgecolor=color, markerfacecolor=color,
                 markersize=marker_size, markeredgewidth=marker_width,
-                color=color, label=label)
+                color=color, label=label_)
         if len(zeros) > 0:
             out[idx, 1] = ax.plot(
                 real(zeros), imag(zeros), marker='o', linestyle='',
@@ -386,10 +395,10 @@ def pole_zero_plot(
 
         # Plot the loci, if present
         if response.loci is not None:
+            label_ = response.sysname if label is None else label[idx]
             for locus in response.loci.transpose():
                 out[idx, 2] += ax.plot(
-                    real(locus), imag(locus), color=color,
-                    label=response.sysname)
+                    real(locus), imag(locus), color=color, label=label_)
 
             # Compute the axis limits to use based on the response
             resp_xlim, resp_ylim = _compute_root_locus_limits(response)
diff --git a/control/tests/ctrlplot_test.py b/control/tests/ctrlplot_test.py
@@ -29,6 +29,7 @@
     (ct.input_output_response,          ct.time_response_plot),
 ]
 
+nolabel_plot_fcns = [ct.describing_function_plot, ct.phase_plane_plot]
 deprecated_fcns = [ct.phase_plot]
 
 # Make sure we didn't miss any plotting functions
@@ -135,6 +136,125 @@ def test_plot_ax_processing(resp_fcn, plot_fcn):
         assert get_line_color(cplt3) != get_line_color(cplt1)
         assert get_line_color(cplt3) != get_line_color(cplt2)
 
+
+@pytest.mark.parametrize("resp_fcn, plot_fcn", resp_plot_fcns)
+@pytest.mark.usefixtures('mplcleanup')
+def test_plot_label_processing(resp_fcn, plot_fcn):
+    # Utility function to make sure legends are OK
+    def assert_legend(cplt, expected_texts):
+        # Check to make sure the labels are OK in legend
+        legend = None
+        for ax in cplt.axes.flatten():
+            legend = ax.get_legend()
+            if legend is not None:
+                break
+        if expected_texts is None:
+            assert legend is None
+        else:
+            assert legend is not None
+            legend_texts = [entry.get_text() for entry in legend.get_texts()]
+            assert legend_texts == expected_texts
+
+    # Create some systems to use
+    sys1 = ct.rss(2, 1, 1, strictly_proper=True, name="sys[1]")
+    sys1c = ct.rss(4, 1, 1, strictly_proper=True, name="sys[1]_C")
+    sys2 = ct.rss(4, 1, 1, strictly_proper=True, name="sys[2]")
+
+    # Set up arguments
+    kwargs = meth_kwargs = plot_fcn_kwargs = {}
+    default_labels = ["sys[1]", "sys[2]"]
+    expected_labels = ["sys1_", "sys2_"]
+    match resp_fcn, plot_fcn:
+        case ct.describing_function_response, _:
+            F = ct.descfcn.saturation_nonlinearity(1)
+            amp = np.linspace(1, 4, 10)
+            args1 = (sys1, F, amp)
+            args2 = (sys2, F, amp)
+
+        case ct.gangof4_response, _:
+            args1 = (sys1, sys1c)
+            args2 = (sys2, sys1c)
+            default_labels = ["P=sys[1]", "P=sys[2]"]
+
+        case ct.frequency_response, ct.nichols_plot:
+            args1 = (sys1, )
+            args2 = (sys2, )
+            meth_kwargs = {'plot_type': 'nichols'}
+
+        case ct.root_locus_map, ct.root_locus_plot:
+            args1 = (sys1, )
+            args2 = (sys2, )
+            meth_kwargs = plot_fcn_kwargs = {'interactive': False}
+
+        case (ct.forced_response | ct.input_output_response, _):
+            timepts = np.linspace(1, 10)
+            U = np.sin(timepts)
+            args1 = (resp_fcn(sys1, timepts, U), )
+            args2 = (resp_fcn(sys2, timepts, U), )
+            argsc = (resp_fcn([sys1, sys2], timepts, U), )
+
+        case (ct.impulse_response | ct.initial_response | ct.step_response, _):
+            args1 = (resp_fcn(sys1), )
+            args2 = (resp_fcn(sys2), )
+            argsc = (resp_fcn([sys1, sys2]), )
+
+        case _, _:
+            args1 = (sys1, )
+            args2 = (sys2, )
+
+    if plot_fcn in nolabel_plot_fcns:
+        pytest.skip(f"labels not implemented for {plot_fcn}")
+
+    # Generate the first plot, with default labels
+    cplt1 = plot_fcn(*args1, **kwargs, **plot_fcn_kwargs)
+    assert isinstance(cplt1, ct.ControlPlot)
+    assert_legend(cplt1, None)
+
+    # Generate second plot with default labels
+    cplt2 = plot_fcn(*args2, **kwargs, **plot_fcn_kwargs)
+    assert isinstance(cplt2, ct.ControlPlot)
+    assert_legend(cplt2, default_labels)
+    plt.close()
+
+    # Generate both plots at the same time
+    if len(args1) == 1 and plot_fcn != ct.time_response_plot:
+        cplt = plot_fcn([*args1, *args2], **kwargs, **plot_fcn_kwargs)
+        assert isinstance(cplt, ct.ControlPlot)
+        assert_legend(cplt, default_labels)
+    elif len(args1) == 1 and plot_fcn == ct.time_response_plot:
+        # Use TimeResponseList.plot() to generate combined response
+        cplt = argsc[0].plot(**kwargs, **plot_fcn_kwargs)
+        assert isinstance(cplt, ct.ControlPlot)
+        assert_legend(cplt, default_labels)
+    plt.close()
+
+    # Generate plots sequentially, with updated labels
+    cplt1 = plot_fcn(
+        *args1, **kwargs, **plot_fcn_kwargs, label=expected_labels[0])
+    assert isinstance(cplt1, ct.ControlPlot)
+    assert_legend(cplt1, None)
+
+    cplt2 = plot_fcn(
+        *args2, **kwargs, **plot_fcn_kwargs, label=expected_labels[1])
+    assert isinstance(cplt2, ct.ControlPlot)
+    assert_legend(cplt2, expected_labels)
+    plt.close()
+
+    # Generate both plots at the same time, with updated labels
+    if len(args1) == 1 and plot_fcn != ct.time_response_plot:
+        cplt = plot_fcn(
+            [*args1, *args2], **kwargs, **plot_fcn_kwargs,
+            label=expected_labels)
+        assert isinstance(cplt, ct.ControlPlot)
+        assert_legend(cplt, expected_labels)
+    elif len(args1) == 1 and plot_fcn == ct.time_response_plot:
+        # Use TimeResponseList.plot() to generate combined response
+        cplt = argsc[0].plot(
+            **kwargs, **plot_fcn_kwargs, label=expected_labels)
+        assert isinstance(cplt, ct.ControlPlot)
+        assert_legend(cplt, expected_labels)
+
+
 @pytest.mark.usefixtures('mplcleanup')
 def test_rcParams():
     sys = ct.rss(2, 2, 2)
diff --git a/control/tests/freqplot_test.py b/control/tests/freqplot_test.py
diff --git a/control/timeplot.py b/control/timeplot.py
