turn off title update if ax is given

murrayrm · murrayrm · commit 8abb6189f0ee · 2024-08-08T14:54:49.000-07:00
diff --git a/control/ctrlplot.py b/control/ctrlplot.py
@@ -60,12 +60,12 @@
 #                 legend_array[i, j] = ax.legend(
 #                     lines, labels, loc=legend_map[i, j])
 #
-#     # Update the plot title
+#     # Update the plot title (only if ax was not given)
 #     sysnames = [response.sysname for response in data]
-#     if title is None:
+#     if ax == None and title is None:
 #         title = "Name plot for " + ", ".join(sysnames)
 #         _update_plot_title(title, fig, rcParams=rcParams)
-#     else:
+#     elif ax == None:
 #         _update_plot_title(title, fig, rcParams=rcParams, use_existing=False)
 #
 #     # Legacy processing of plot keyword
@@ -362,8 +362,9 @@ def _process_ax_keyword(
                     text.set_visible(False)     # turn off the text
                     del text                    # get rid of it completely
     else:
+        axs = np.atleast_1d(axs)
         try:
-            axs = np.asarray(axs).reshape(shape)
+            axs = axs.reshape(shape)
         except ValueError:
             raise ValueError(
                 "specified axes are not the right shape; "
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -957,14 +957,14 @@ def gen_zero_centered_series(val_min, val_max, period):
     sysnames = [response.sysname for response in data \
                 if not (response.sysname in seen or seen.add(response.sysname))]
 
-    if title is None:
+    if ax is None and 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_plot_title(title, fig, rcParams=rcParams, frame=suptitle_frame)
-    else:
+    elif ax is None:
         _update_plot_title(
             title, fig=fig, rcParams=rcParams, frame=suptitle_frame,
             use_existing=False)
@@ -1679,6 +1679,7 @@ def nyquist_plot(
     arrow_size = config._get_param(
         'nyquist', 'arrow_size', kwargs, _nyquist_defaults, pop=True)
     arrow_style = config._get_param('nyquist', 'arrow_style', kwargs, None)
+    ax_user = ax
     max_curve_magnitude = config._get_param(
         'nyquist', 'max_curve_magnitude', kwargs, _nyquist_defaults, pop=True)
     max_curve_offset = config._get_param(
@@ -1773,7 +1774,7 @@ def _parse_linestyle(style_name, allow_false=False):
         return (counts, contours) if return_contour else counts
 
     fig, ax = _process_ax_keyword(
-        ax, shape=(1, 1), squeeze=True, rcParams=rcParams)
+        ax_user, shape=(1, 1), squeeze=True, rcParams=rcParams)
 
     # Create a list of lines for the output
     out = np.empty(len(nyquist_responses), dtype=object)
@@ -1956,11 +1957,12 @@ def _parse_linestyle(style_name, allow_false=False):
         legend=None
 
     # Add the title
-    if title is None:
-        title = "Nyquist plot for " + ", ".join(labels)
-    _update_plot_title(
-        title, fig=fig, rcParams=rcParams, frame=suptitle_frame,
-        use_existing=False)
+    if ax_user is None:
+        if title is None:
+            title = "Nyquist plot for " + ", ".join(labels)
+        _update_plot_title(
+            title, fig=fig, rcParams=rcParams, frame=suptitle_frame,
+            use_existing=False)
 
     # Legacy return pocessing
     if plot is True or return_contour is not None:
@@ -2418,11 +2420,12 @@ def singular_values_plot(
         legend = None
 
     # Add the title
-    if title is None:
-        title = "Singular values for " + ", ".join(labels)
-    _update_plot_title(
-        title, fig=fig, rcParams=rcParams, frame=suptitle_frame,
-        use_existing=False)
+    if ax is None:
+        if title is None:
+            title = "Singular values for " + ", ".join(labels)
+        _update_plot_title(
+            title, fig=fig, rcParams=rcParams, frame=suptitle_frame,
+            use_existing=False)
 
     # Legacy return processing
     if plot is not None:
diff --git a/control/nichols.py b/control/nichols.py
@@ -150,9 +150,11 @@ def nichols_plot(
         legend = None
 
     # Add the title
-    if title is None:
-        title = "Nichols plot for " + ", ".join(labels)
-    _update_plot_title(title, fig=fig, rcParams=rcParams, use_existing=False)
+    if ax is None:
+        if title is None:
+            title = "Nichols plot for " + ", ".join(labels)
+        _update_plot_title(
+            title, fig=fig, rcParams=rcParams, use_existing=False)
 
     return ControlPlot(out, ax_nichols, fig, legend=legend)
 
diff --git a/control/phaseplot.py b/control/phaseplot.py
@@ -219,14 +219,13 @@ 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):
-            if title is None:
-                title = f"Phase portrait for {sys.name}"
-            _update_plot_title(title, use_existing=False)
-            ax.set_xlabel(sys.state_labels[0])
-            ax.set_ylabel(sys.state_labels[1])
+        if title is None:
+            title = f"Phase portrait for {sys.name}"
+        _update_plot_title(title, use_existing=False)
+        ax.set_xlabel(sys.state_labels[0])
+        ax.set_ylabel(sys.state_labels[1])
+        plt.tight_layout()
 
-    plt.tight_layout()
     return ControlPlot(out, ax, fig)
 
 
diff --git a/control/tests/ctrlplot_test.py b/control/tests/ctrlplot_test.py
@@ -126,7 +126,7 @@ def test_plot_ax_processing(resp_fcn, plot_fcn):
     if resp_fcn is not None:
         cplt3 = resp.plot(**kwargs, **meth_kwargs, ax=cplt1.axes)
     else:
-        cplt3 = plot_fcn(*args, **kwargs, **meth_kwargs)
+        cplt3 = plot_fcn(*args, **kwargs, **meth_kwargs, ax=cplt1.axes)
     assert cplt3.figure == cplt1.figure
 
     # Plot should have landed on top of previous plot, in different colors
@@ -136,6 +136,51 @@ 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)
 
+    #
+    # Plot on a user-contructed figure
+    #
+
+    # Store modified properties from previous figure
+    cplt_titlesize = cplt3.figure._suptitle.get_fontsize()
+    cplt_labelsize = \
+        cplt3.axes.reshape(-1)[0].get_yticklabels()[0].get_fontsize()
+
+    # Set up some axes with a known title
+    fig, axs = plt.subplots(2, 3)
+    title = "User-constructed figure"
+    plt.suptitle(title)
+    titlesize = fig._suptitle.get_fontsize()
+    assert titlesize != cplt_titlesize
+    labelsize = axs[0, 0].get_yticklabels()[0].get_fontsize()
+    assert labelsize != cplt_labelsize
+
+    # Figure out what to pass as the ax keyword
+    match resp_fcn, plot_fcn:
+        case _, ct.bode_plot:
+            ax = [axs[0, 1], axs[1, 1]]
+
+        case ct.gangof4_response, _:
+            ax = [axs[0, 1], axs[0, 2], axs[1, 1], axs[1, 2]]
+
+        case (ct.forced_response | ct.input_output_response, _):
+            ax = [axs[0, 1], axs[1, 1]]
+
+        case _, _:
+            ax = [axs[0, 1]]
+
+    # Call the plotting function, passing the axes
+    if resp_fcn is not None:
+        resp = resp_fcn(*args, **kwargs)
+        cplt4 = resp.plot(**kwargs, **meth_kwargs, ax=ax)
+    else:
+        # No response function available; just plot the data
+        cplt4 = plot_fcn(*args, **kwargs, **meth_kwargs, ax=ax)
+
+    # Check to make sure original settings did not change
+    assert fig._suptitle.get_text() == title
+    assert fig._suptitle.get_fontsize() == titlesize
+    assert ax[0].get_yticklabels()[0].get_fontsize() == labelsize
+
 
 @pytest.mark.parametrize("resp_fcn, plot_fcn", resp_plot_fcns)
 @pytest.mark.usefixtures('mplcleanup')
@@ -338,13 +383,13 @@ def test_plot_title_processing(resp_fcn, plot_fcn):
         case _:
             raise RuntimeError(f"didn't recognize {resp_fnc}, {plot_fnc}")
 
-    # Generate the first plot, with default labels
+    # Generate the first plot, with default title
     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
+        # Generate second plot with default title
         cplt2 = plot_fcn(*args2, **kwargs, **plot_fcn_kwargs)
         assert cplt1.figure._suptitle._text == title_prefix + default_title
         plt.close()
@@ -372,6 +417,11 @@ def test_plot_title_processing(resp_fcn, plot_fcn):
     assert cplt2.figure._suptitle._text == "Another title"
     plt.close()
 
+    # Generate the plots with no title
+    cplt = plot_fcn(
+        *args1, **kwargs, **plot_fcn_kwargs, title=False)
+    assert cplt.figure._suptitle == None
+
 
 @pytest.mark.usefixtures('mplcleanup')
 def test_rcParams():
diff --git a/control/timeplot.py b/control/timeplot.py
@@ -175,6 +175,7 @@ def time_response_plot(
     # Process keywords and set defaults
     #
     # Set up defaults
+    ax_user = ax
     time_label = config._get_param(
         'timeplot', 'time_label', kwargs, _timeplot_defaults, pop=True)
     rcParams = config._get_param(
@@ -657,10 +658,10 @@ def _make_line_label(signal_index, signal_labels, trace_index):
     # list of systems (e.g., "Step response for sys[1], sys[2]").
     #
 
-    if title is None:
+    if ax_user is None and title is None:
         title = data.title if title == None else title
         _update_plot_title(title, fig, rcParams=rcParams)
-    else:
+    elif ax_user is None:
         _update_plot_title(title, fig, rcParams=rcParams, use_existing=False)
 
     return ControlPlot(out, ax_array, fig, legend=legend_map)
