Fix suptitle for root locus plots

murrayrm · murrayrm · commit f8180343f20b · 2024-08-08T14:54:49.000-07:00
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -18,7 +18,8 @@
 from numpy import cos, exp, imag, linspace, real, sin, sqrt
 
 from . import config
-from .ctrlplot import ControlPlot, _get_line_labels
+from .ctrlplot import ControlPlot, suptitle, _get_line_labels, \
+    _process_ax_keyword
 from .freqplot import _freqplot_defaults
 from .grid import nogrid, sgrid, zgrid
 from .iosys import isctime, isdtime
@@ -366,6 +367,7 @@ def pole_zero_plot(
 
     # Handle color cycle manually as all root locus segments
     # of the same system are expected to be of the same color
+    # TODO: replace with common function?
     color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']
     color_offset = 0
     if len(ax.lines) > 0:
@@ -470,10 +472,10 @@ def pole_zero_plot(
 
     # Add the title
     if title is None:
-        title = "Pole/zero plot for " + ", ".join(labels)
+        title = ("Root locus plot for " if rlocus_plot
+                 else "Pole/zero plot for ") + ", ".join(labels)
     if user_ax is None:
-        with plt.rc_context(freqplot_rcParams):
-            fig.suptitle(title)
+        suptitle(title)
 
     # Add dispather to handle choosing a point on the diagram
     if interactive:
diff --git a/control/tests/rlocus_test.py b/control/tests/rlocus_test.py
@@ -278,6 +278,9 @@ def test_root_locus_documentation(savefigs=False):
         plt.figure()
         test_root_locus_plots(
             sys, grid=grid, xlim=xlim, ylim=ylim, interactive=interactive)
+        ct.suptitle(
+            f"sys={sys.name}, {grid=}, {xlim=}, {ylim=}, {interactive=}",
+            frame='figure')
 
     # Run tests that generate plots for the documentation
     test_root_locus_documentation(savefigs=True)
diff --git a/examples/plot_gallery.py b/examples/plot_gallery.py
@@ -125,6 +125,7 @@ def invpend_update(t, x, u, params):
     sys2 = ct.tf([1, 0.2], [1, 1, 3, 1, 1], name='sys2')
     ct.root_locus_plot([sys1, sys2], grid=True, ax=ax1)
     ct.root_locus_plot([sys1, sys2], grid=False, ax=ax2)
+    ct.suptitle("Root locus plots (w/ specified axes)")
     print("  -- BUG: should have 2 x 1 array of plots")
 
 # sisotool
