implement pole_zero_subplots

murrayrm · murrayrm · commit 4f3d618f6f0d · 2024-08-08T14:54:49.000-07:00
diff --git a/control/ctrlplot.py b/control/ctrlplot.py
@@ -72,6 +72,7 @@
 #
 #     return ControlPlot(lines, ax_array, fig, legend=legend_map)
 
+import itertools
 import warnings
 from os.path import commonprefix
 
@@ -81,7 +82,7 @@
 
 from . import config
 
-__all__ = ['ControlPlot', 'suptitle', 'get_plot_axes']
+__all__ = ['ControlPlot', 'suptitle', 'get_plot_axes', 'pole_zero_subplots']
 
 #
 # Style parameters
@@ -241,14 +242,76 @@ def get_plot_axes(line_array):
     else:
         return _get_axes(line_array)
 
+
+def pole_zero_subplots(
+        nrows, ncols, grid=None, dt=None, fig=None, scaling=None,
+        rcParams=None):
+    """Create axes for pole/zero plot.
+
+    Parameters
+    ----------
+    nrows, ncols : int
+        Number of rows and columns
+    grid : True, False, or 'empty', optional
+        Grid style to use.  Can also be a list, in which case each subplot
+        will have a different style (columns then rows).
+    dt : timebase, option
+        Timebase for each subplot (or a list of timebases).
+    scaling : 'auto', 'equal', or None
+        Scaling to apply to the subplots.
+    fig : :class:`matplotlib.figure.Figure`
+        Figure to use for creating subplots.
+
+    Returns
+    -------
+    ax_array : array
+        2D array of axes
+
+    """
+    from .grid import sgrid, zgrid, nogrid
+    from .iosys import isctime
+
+    if rcParams is None:
+        rcParams = _ctrlplot_rcParams
+
+    if fig is None:
+        fig = plt.gcf()
+
+    if not isinstance(grid, list):
+        grid = [grid] * nrows * ncols
+    if not isinstance(dt, list):
+        dt = [dt] * nrows * ncols
+
+    ax_array = np.full((nrows, ncols), None)
+    index = 0
+    with plt.rc_context(rcParams):
+        for row, col in itertools.product(range(nrows), range(ncols)):
+            match grid[index], isctime(dt=dt[index]):
+                case 'empty', _:        # empty grid
+                    ax_array[row, col] = fig.add_subplot(nrows, ncols, index+1)
+
+                case True, True:        # continuous time grid
+                    ax_array[row, col], _ = sgrid(
+                        (nrows, ncols, index+1), scaling=scaling)
+
+                case True, False:       # discrete time grid
+                    ax_array[row, col] = fig.add_subplot(nrows, ncols, index+1)
+                    zgrid(ax=ax_array[row, col], scaling=scaling)
+
+                case False | None, _:   # no grid (just stability boundaries)
+                    ax_array[row, col] = fig.add_subplot(nrows, ncols, index+1)
+                    nogrid(
+                        ax=ax_array[row, col], dt=dt[index], scaling=scaling)
+            index += 1
+    return ax_array
+
 #
 # Utility functions
 #
 # These functions are used by plotting routines to provide a consistent way
 # of processing and displaying information.
 #
 
-
 def _process_ax_keyword(
         axs, shape=(1, 1), rcParams=None, squeeze=False, clear_text=False,
         create_axes=True):
diff --git a/control/grid.py b/control/grid.py
@@ -74,7 +74,7 @@ def __call__(self, transform_xy, x1, y1, x2, y2):
         return lon_min, lon_max, lat_min, lat_max
 
 
-def sgrid(scaling=None):
+def sgrid(subplot=(1, 1, 1), scaling=None):
     # From matplotlib demos:
     # https://matplotlib.org/gallery/axisartist/demo_curvelinear_grid.html
     # https://matplotlib.org/gallery/axisartist/demo_floating_axis.html
@@ -101,7 +101,7 @@ def sgrid(scaling=None):
 
     # Set up an axes with a specialized grid helper
     fig = plt.gcf()
-    ax = SubplotHost(fig, 1, 1, 1, grid_helper=grid_helper)
+    ax = SubplotHost(fig, *subplot, grid_helper=grid_helper)
 
     # make ticklabels of right invisible, and top axis visible.
     visible = True
diff --git a/control/sisotool.py b/control/sisotool.py
@@ -136,7 +136,7 @@ def sisotool(sys, initial_gain=None, xlim_rlocus=None, ylim_rlocus=None,
     #     ax=fig.axes[1])
     ax_rlocus = fig.axes[1]
     root_locus_map(sys[0, 0]).plot(
-        xlim=xlim_rlocus, ylim=ylim_rlocus, grid=rlocus_grid,
+        xlim=xlim_rlocus, ylim=ylim_rlocus,
         initial_gain=initial_gain, ax=ax_rlocus)
     if rlocus_grid is False:
         # Need to generate grid manually, since root_locus_plot() won't
diff --git a/examples/plot_gallery.py b/examples/plot_gallery.py
@@ -120,13 +120,13 @@ def invpend_update(t, x, u, params):
 
 # root locus
 with create_figure("Root locus plot") as fig:
-    ax1, ax2 = fig.subplots(2, 1)
+    ax_array = ct.pole_zero_subplots(2, 1, grid=[True, False])
+    ax1, ax2 = ax_array[:, 0]
     sys1 = ct.tf([1, 2], [1, 2, 3], name='sys1')
     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")
+    ct.root_locus_plot([sys1, sys2], ax=ax1)
+    ct.root_locus_plot([sys1, sys2], ax=ax2)
+    plt.suptitle("Root locus plots (w/ specified axes)", fontsize='medium')
 
 # sisotool
 with create_figure("sisotool"):
