update title processing to be uniform across _plot functions

murrayrm · Aug 9, 2024 · Jun 28, 2024 · Jun 27, 2024 · Jun 29, 2024 · Jul 3, 2024
commit 89cda3d0c22d87f43d23a062e2586cc43aeee717
diff --git a/control/ctrlplot.py b/control/ctrlplot.py
@@ -64,7 +64,9 @@
 #     sysnames = [response.sysname for response in data]
 #     if title is None:
 #         title = "Name plot for " + ", ".join(sysnames)
-#     _update_suptitle(fig, title, rcParams=rcParams)
+#         _update_suptitle(title, fig, rcParams=rcParams)
+#     else
+#         suptitle(title, fig, rcParams=rcParams)
 #
 #     # Legacy processing of plot keyword
 #     if plot is True:
@@ -159,6 +161,9 @@ def __setitem__(self, item, val):
     def reshape(self, *args):
         return self.lines.reshape(*args)
 
+    def set_plot_title(self, title, frame='axes'):
+        suptitle(title, fig=self.figure, frame=frame)
+
 
 #
 # User functions
@@ -467,7 +472,7 @@ def _make_legend_labels(labels, ignore_common=False):
     return labels
 
 
-def _update_suptitle(fig, title, rcParams=None, frame='axes'):
+def _update_suptitle(title, fig=None, frame='axes', rcParams=None):
     if fig is not None and isinstance(title, str):
         # Get the current title, if it exists
         old_title = None if fig._suptitle is None else fig._suptitle._text

diff --git a/control/freqplot.py b/control/freqplot.py
@@ -142,6 +142,12 @@ def bode_plot(
 
     Other Parameters
     ----------------
+    ax : array of Axes
+        The matplotlib Axes to draw the figure on.  If not specified, the
+        Axes for the current figure are used or, if there is no current
+        figure with the correct number and shape of Axes, a new figure is
+        created.  The shape of the array must match the shape of the
+        plotted data.
     grid : bool
         If True, plot grid lines on gain and phase plots.  Default is set by
         `config.defaults['freqplot.grid']`.
@@ -173,6 +179,8 @@ def bode_plot(
     rcParams : dict
         Override the default parameters used for generating plots.
         Default is set by config.default['freqplot.rcParams'].
+    title : str, optional
+        Set the title of the plot.  Defaults to plot type and system name(s).
     wrap_phase : bool or float
         If wrap_phase is `False` (default), then the phase will be unwrapped
         so that it is continuously increasing or decreasing.  If wrap_phase is
@@ -948,13 +956,16 @@ def gen_zero_centered_series(val_min, val_max, period):
     seen = set()
     sysnames = [response.sysname for response in data \
                 if not (response.sysname in seen or seen.add(response.sysname))]
+
     if title is None:
         if data[0].title is None:
             title = "Bode plot for " + ", ".join(sysnames)
         else:
+            # Allow data to set the title (used by gangof4)
             title = data[0].title
-
-    _update_suptitle(fig, title, rcParams=rcParams, frame=suptitle_frame)
+        _update_suptitle(title, fig, rcParams=rcParams, frame=suptitle_frame)
+    else:
+        suptitle(title, fig=fig, rcParams=rcParams, frame=suptitle_frame)
 
     #
     # Create legends
@@ -1603,6 +1614,8 @@ def nyquist_plot(
     start_marker_size : float, optional
         Start marker size (in display coordinates).  Default value is
         4 and can be set using config.defaults['nyquist.start_marker_size'].
+    title : str, optional
+        Set the title of the plot.  Defaults to plot type and system name(s).
     warn_nyquist : bool, optional
         If set to 'False', turn off warnings about frequencies above Nyquist.
     warn_encirclements : bool, optional
@@ -1870,7 +1883,7 @@ def _parse_linestyle(style_name, allow_false=False):
         # Display the unit circle, to read gain crossover frequency
         if unit_circle:
             plt.plot(cos, sin, **config.defaults['nyquist.circle_style'])
-        
+
         # Draw circles for given magnitudes of sensitivity
         if ms_circles is not None:
             for ms in ms_circles:
@@ -2243,6 +2256,8 @@ def singular_values_plot(
     rcParams : dict
         Override the default parameters used for generating plots.
         Default is set up config.default['freqplot.rcParams'].
+    title : str, optional
+        Set the title of the plot.  Defaults to plot type and system name(s).
 
     See Also
     --------

diff --git a/control/nichols.py b/control/nichols.py
@@ -60,6 +60,8 @@ def nichols_plot(
     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.
+    title : str, optional
+        Set the title of the plot.  Defaults to plot type and system name(s).
     **kwargs : :func:`matplotlib.pyplot.plot` keyword properties, optional
         Additional keywords passed to `matplotlib` to specify line properties.
 

diff --git a/control/phaseplot.py b/control/phaseplot.py
@@ -54,7 +54,8 @@
 def phase_plane_plot(
         sys, pointdata=None, timedata=None, gridtype=None, gridspec=None,
         plot_streamlines=True, plot_vectorfield=False, plot_equilpoints=True,
-        plot_separatrices=True, ax=None, suppress_warnings=False, **kwargs
+        plot_separatrices=True, ax=None, suppress_warnings=False, title=None,
+        **kwargs
 ):
     """Plot phase plane diagram.
 
@@ -135,6 +136,8 @@ def phase_plane_plot(
         in the dict as keywords to :func:`~control.phaseplot.separatrices`.
     suppress_warnings : bool, optional
         If set to `True`, suppress warning messages in generating trajectories.
+    title : str, optional
+        Set the title of the plot.  Defaults to plot type and system name(s).
 
     """
     # Process arguments
@@ -217,7 +220,9 @@ def _create_kwargs(global_kwargs, local_kwargs, **other_kwargs):
     # TODO: update to common code pattern
     if user_ax is None:
         with plt.rc_context(rcParams):
-            suptitle(f"Phase portrait for {sys.name}")
+            if title is None:
+                title = f"Phase portrait for {sys.name}"
+            suptitle(title)
             ax.set_xlabel(sys.state_labels[0])
             ax.set_ylabel(sys.state_labels[1])
 

diff --git a/control/pzmap.py b/control/pzmap.py
@@ -249,7 +249,7 @@ def pole_zero_plot(
     marker_width : int, optional
         Set the line width of the markers used for poles and zeros.
     title : str, optional
-        Set the title of the plot.  Defaults plot type and system name(s).
+        Set the title of the plot.  Defaults to plot type and system name(s).
     xlim : list, optional
         Set the limits for the x axis.
     ylim : list, optional

diff --git a/control/tests/ctrlplot_test.py b/control/tests/ctrlplot_test.py
@@ -128,7 +128,7 @@ def test_plot_ax_processing(resp_fcn, plot_fcn):
     else:
         cplt3 = plot_fcn(*args, **kwargs, **meth_kwargs)
     assert cplt3.figure == cplt1.figure
-    
+
     # Plot should have landed on top of previous plot, in different colors
     assert np.all(cplt3.axes == cplt1.axes)
     assert len(cplt3.lines[0]) == len(cplt1.lines[0])
@@ -255,6 +255,124 @@ def assert_legend(cplt, expected_texts):
         assert_legend(cplt, expected_labels)
 
 
+@pytest.mark.parametrize("resp_fcn, plot_fcn", resp_plot_fcns)
+@pytest.mark.usefixtures('mplcleanup')
+def test_plot_title_processing(resp_fcn, plot_fcn):
+    # 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(2, 1, 1, strictly_proper=True, name="sys[2]")
+
+    # Set up arguments
+    kwargs = meth_kwargs = plot_fcn_kwargs = {}
+    default_title = "sys[1], sys[2]"
+    expected_title = "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_title = "P=sys[1], C=sys[1]_C, P=sys[2], C=sys[1]_C"
+
+        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, )
+
+    # Store the expected title prefix
+    match resp_fcn, plot_fcn:
+        case _, ct.bode_plot:
+            title_prefix = "Bode plot for "
+        case _, ct.nichols_plot:
+            title_prefix = "Nichols plot for "
+        case _, ct.singular_values_plot:
+            title_prefix = "Singular values for "
+        case _, ct.gangof4_plot:
+            title_prefix = "Gang of Four for "
+        case _, ct.describing_function_plot:
+            title_prefix = "Nyquist plot for "
+        case _, ct.phase_plane_plot:
+            title_prefix = "Phase portrait for "
+        case _, ct.pole_zero_plot:
+            title_prefix = "Pole/zero plot for "
+        case _, ct.nyquist_plot:
+            title_prefix = "Nyquist plot for "
+        case _, ct.root_locus_plot:
+            title_prefix = "Root locus plot for "
+        case ct.initial_response, _:
+            title_prefix = "Initial response for "
+        case ct.step_response, _:
+            title_prefix = "Step response for "
+        case ct.impulse_response, _:
+            title_prefix = "Impulse response for "
+        case ct.forced_response, _:
+            title_prefix = "Forced response for "
+        case ct.input_output_response, _:
+            title_prefix = "Input/output response for "
+        case _:
+            raise RuntimeError(f"didn't recognize {resp_fnc}, {plot_fnc}")
+
+    # Generate the first plot, with default labels
+    cplt1 = plot_fcn(*args1, **kwargs, **plot_fcn_kwargs)
+    assert cplt1.figure._suptitle._text.startswith(title_prefix)
+
+    # Skip functions not intended for sequential calling
+    if plot_fcn not in nolabel_plot_fcns:
+        # Generate second plot with default labels
+        cplt2 = plot_fcn(*args2, **kwargs, **plot_fcn_kwargs)
+        assert cplt1.figure._suptitle._text == title_prefix + default_title
+        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 cplt.figure._suptitle._text == title_prefix + default_title
+        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 cplt.figure._suptitle._text == title_prefix + default_title
+        plt.close()
+
+    # Generate plots sequentially, with updated titles
+    cplt1 = plot_fcn(
+        *args1, **kwargs, **plot_fcn_kwargs, title="My first title")
+    cplt2 = plot_fcn(
+        *args2, **kwargs, **plot_fcn_kwargs, title="My new title")
+    assert cplt2.figure._suptitle._text == "My new title"
+    plt.close()
+
+    # Update using set_plot_title
+    cplt2.set_plot_title("Another title")
+    assert cplt2.figure._suptitle._text == "Another title"
+    plt.close()
+
+
 @pytest.mark.usefixtures('mplcleanup')
 def test_rcParams():
     sys = ct.rss(2, 2, 2)

diff --git a/control/timeplot.py b/control/timeplot.py
@@ -15,8 +15,8 @@
 import numpy as np
 
 from . import config
-from .ctrlplot import ControlPlot, _ctrlplot_rcParams, _make_legend_labels, \
-    _update_suptitle
+from .ctrlplot import ControlPlot, suptitle, _ctrlplot_rcParams, \
+    _make_legend_labels, _update_suptitle
 
 __all__ = ['time_response_plot', 'combine_time_responses']
 
@@ -139,6 +139,8 @@ def time_response_plot(
         axis or ``legend_loc`` or ``legend_map`` have been specified.
     time_label : str, optional
         Label to use for the time axis.
+    title : str, optional
+        Set the title of the plot.  Defaults to plot type and system name(s).
     trace_labels : list of str, optional
         Replace the default trace labels with the given labels.
     trace_props : array of dicts
@@ -196,9 +198,6 @@ def time_response_plot(
         'timeplot', 'trace_props', kwargs, _timeplot_defaults, pop=True)
     tprop_len = len(trace_props)
 
-    # Set the title for the data
-    title = data.title if title == None else title
-
     # Determine whether or not to plot the input data (and how)
     if plot_inputs is None:
         plot_inputs = data.plot_inputs
@@ -658,7 +657,11 @@ def _make_line_label(signal_index, signal_labels, trace_index):
     # list of systems (e.g., "Step response for sys[1], sys[2]").
     #
 
-    _update_suptitle(fig, title, rcParams=rcParams)
+    if title is None:
+        title = data.title if title == None else title
+        _update_suptitle(title, fig, rcParams=rcParams)
+    else:
+        suptitle(title, fig, rcParams=rcParams)
 
     return ControlPlot(out, ax_array, fig, legend=legend_map)
 
@@ -676,6 +679,9 @@ def combine_time_responses(response_list, trace_labels=None, title=None):
     trace_labels : list of str, optional
         List of labels for each trace.  If not specified, trace names are
         taken from the input data or set to None.
+    title : str, optional
+        Set the title to use when plotting.  Defaults to plot type and
+        system name(s).
 
     Returns
     -------

diff --git a/control/timeresp.py b/control/timeresp.py
@@ -192,6 +192,9 @@ class TimeResponseData:
         response.  If ntraces is 0 (default) then the data represents a
         single trace with the trace index surpressed in the data.
 
+    title : str, optional
+        Set the title to use when plotting.
+
     trace_labels : array of string, optional
         Labels to use for traces (set to sysname it ntraces is 0)