adjust root locus spacing and add scaling keyword

murrayrm · murrayrm · commit ee46ac6a2fa1 · 2023-12-28T12:41:09.000-08:00
diff --git a/control/grid.py b/control/grid.py
@@ -1,6 +1,14 @@
+# grid.py - code to add gridlines to root locus and pole-zero diagrams
+#
+# This code generates grids for pole-zero diagrams (including root locus
+# diagrams).  Rather than just draw a grid in place, it uses the AxisArtist
+# package to generate a custom grid that will scale with the figure.
+#
+
 import numpy as np
 from numpy import cos, sin, sqrt, linspace, pi, exp
 import matplotlib.pyplot as plt
+
 from mpl_toolkits.axisartist import SubplotHost
 from mpl_toolkits.axisartist.grid_helper_curvelinear \
     import GridHelperCurveLinear
@@ -67,14 +75,15 @@ def __call__(self, transform_xy, x1, y1, x2, y2):
         return lon_min, lon_max, lat_min, lat_max
 
 
-def sgrid():
+def sgrid(scaling=None):
     # From matplotlib demos:
     # https://matplotlib.org/gallery/axisartist/demo_curvelinear_grid.html
     # https://matplotlib.org/gallery/axisartist/demo_floating_axis.html
 
     # PolarAxes.PolarTransform takes radian. However, we want our coordinate
-    # system in degree
+    # system in degrees
     tr = Affine2D().scale(np.pi/180., 1.) + PolarAxes.PolarTransform()
+
     # polar projection, which involves cycle, and also has limits in
     # its coordinates, needs a special method to find the extremes
     # (min, max of the coordinate within the view).
@@ -91,6 +100,7 @@ def sgrid():
         tr, extreme_finder=extreme_finder, grid_locator1=grid_locator1,
         tick_formatter1=tick_formatter1)
 
+    # Set up an axes with a specialized grid helper
     fig = plt.gcf()
     ax = SubplotHost(fig, 1, 1, 1, grid_helper=grid_helper)
 
@@ -101,6 +111,7 @@ def sgrid():
     ax.axis[:].invert_ticklabel_direction()
     ax.axis[:].major_ticklabels.set_color('gray')
 
+    # Set up internal tickmarks and labels along the real/imag axes
     ax.axis["wnxneg"] = axis = ax.new_floating_axis(0, 180)
     axis.set_ticklabel_direction("-")
     axis.label.set_visible(False)
@@ -125,35 +136,26 @@ def sgrid():
     ax.axis["bottom"].get_helper().nth_coord_ticks = 0
 
     fig.add_subplot(ax)
-
-    # RECTANGULAR X Y AXES WITH SCALE
-    # par2 = ax.twiny()
-    # par2.axis["top"].toggle(all=False)
-    # par2.axis["right"].toggle(all=False)
-    # new_fixed_axis = par2.get_grid_helper().new_fixed_axis
-    # par2.axis["left"] = new_fixed_axis(loc="left",
-    #                                   axes=par2,
-    #                                   offset=(0, 0))
-    # par2.axis["bottom"] = new_fixed_axis(loc="bottom",
-    #                                     axes=par2,
-    #                                     offset=(0, 0))
-    # FINISH RECTANGULAR
-
     ax.grid(True, zorder=0, linestyle='dotted')
 
-    _final_setup(ax)
+    _final_setup(ax, scaling=scaling)
     return ax, fig
 
 
-def _final_setup(ax):
+# Utility function used by all grid code
+def _final_setup(ax, scaling=None):
     ax.set_xlabel('Real')
     ax.set_ylabel('Imaginary')
     ax.axhline(y=0, color='black', lw=0.5)
     ax.axvline(x=0, color='black', lw=0.5)
-    plt.axis('equal')
+
+    # Set up the scaling for the axes
+    scaling = 'equal' if scaling is None else scaling
+    plt.axis(scaling)
 
 
-def nogrid(dt=None, ax=None):
+# If not grid is given, at least separate stable/unstable regions
+def nogrid(dt=None, ax=None, scaling=None):
     fig = plt.gcf()
     if ax is None:
         ax = fig.gca()
@@ -163,11 +165,12 @@ def nogrid(dt=None, ax=None):
         s = np.linspace(0, 2*pi, 100)
         ax.plot(np.cos(s), np.sin(s), 'k--', lw=0.5, dashes=(5, 5))
 
-    _final_setup(ax)
+    _final_setup(ax, scaling=scaling)
     return ax, fig
 
-
-def zgrid(zetas=None, wns=None, ax=None):
+# Grid for discrete time system (drawn, not rendered by AxisArtist)
+# TODO (at some point): think about using customized grid generator?
+def zgrid(zetas=None, wns=None, ax=None, scaling=None):
     """Draws discrete damping and frequency grid"""
 
     fig = plt.gcf()
@@ -218,5 +221,9 @@ def zgrid(zetas=None, wns=None, ax=None):
         ax.annotate(r"$\frac{"+num+r"\pi}{T}$", xy=(an_x, an_y),
                     xytext=(an_x, an_y), size=9)
 
-    _final_setup(ax)
+    # Set default axes to allow some room around the unit circle
+    ax.set_xlim([-1.1, 1.1])
+    ax.set_ylim([-1.1, 1.1])
+
+    _final_setup(ax, scaling=scaling)
     return ax, fig
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -42,7 +42,8 @@
     'pzmap.grid': False,                # Plot omega-damping grid
     'pzmap.marker_size': 6,             # Size of the markers
     'pzmap.marker_width': 1.5,          # Width of the markers
-    'pzmap.expansion_factor': 2,        # Amount to scale plots beyond features
+    'pzmap.expansion_factor': 1.8,      # Amount to scale plots beyond features
+    'pzmap.buffer_factor': 1.05,        # Buffer to leave around plot peaks
 }
 
 #
@@ -110,7 +111,7 @@ def pole_zero_map(sysdata):
 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=False, ax=None,
+        xlim=None, ylim=None, interactive=False, ax=None, scaling=None,
         initial_gain=None, **kwargs):
     # TODO: update docstring (see other response/plot functions for style)
     """Plot a pole/zero map for a linear system.
@@ -144,7 +145,7 @@ def pole_zero_plot(
         (legacy) If the `plot` keyword is given, the system poles and zeros
         are returned.
 
-    Notes (TODO: update)
+    Notes (TODO: update, including scaling)
     -----
     The pzmap function calls matplotlib.pyplot.axis('equal'), which means
     that trying to reset the axis limits may not behave as expected.  To
@@ -209,14 +210,15 @@ def pole_zero_plot(
         if grid:
             plt.clf()
             if all([isctime(dt=response.dt) for response in data]):
-                ax, fig = sgrid()
+                ax, fig = sgrid(scaling=scaling)
             elif all([isdtime(dt=response.dt) for response in data]):
-                ax, fig = zgrid()
+                ax, fig = zgrid(scaling=scaling)
             else:
                 ValueError(
                     "incompatible time responses; don't know how to grid")
         elif len(axs) == 0:
-            ax, fig = nogrid(data[0].dt)        # use first response timebase
+            # use first response timebase
+            ax, fig = nogrid(data[0].dt, scaling=scaling)
         else:
             # Use the existing axes and any grid that is there
             # TODO: allow axis to be overriden via parameter
@@ -270,7 +272,7 @@ def pole_zero_plot(
                     label=response.sysname)
 
             # Compute the axis limits to use based on the response
-            resp_xlim, resp_ylim = _compute_root_locus_limits(response.loci)
+            resp_xlim, resp_ylim = _compute_root_locus_limits(response)
 
             # Keep track of the current limits
             xlim = [min(xlim[0], resp_xlim[0]), max(xlim[1], resp_xlim[1])]
@@ -433,11 +435,22 @@ def _create_root_locus_label(sys, K, s):
 
 
 # Utility function to compute limits for root loci
-# TODO: compare to old code and recapture functionality (especially asymptotes)
 # TODO: (note that sys is now available => code here may not be needed)
-def _compute_root_locus_limits(loci):
-    # Go through each locus
-    xlim, ylim = [0, 0], 0
+def _compute_root_locus_limits(response):
+    loci = response.loci
+
+    # Start with information about zeros, if present
+    if response.sys is not None and response.sys.zeros().size > 0:
+        xlim = [
+            min(0, np.min(response.sys.zeros().real)),
+            max(0, np.max(response.sys.zeros().real))
+        ]
+        ylim = max(0, np.max(response.sys.zeros().imag))
+    else:
+        xlim, ylim = [0, 0], 0
+
+    # Go through each locus and look for features
+    rho = config._get_param('pzmap', 'buffer_factor')
     for locus in loci.transpose():
         # Include all starting points
         xlim = [min(xlim[0], locus[0].real), max(xlim[1], locus[0].real)]
@@ -446,18 +459,22 @@ def _compute_root_locus_limits(loci):
         # Find the local maxima of root locus curve
         xpeaks = np.where(
             np.diff(np.abs(locus.real)) < 0, locus.real[0:-1], 0)
-        xlim = [min(xlim[0], np.min(xpeaks)), max(xlim[1], np.max(xpeaks))]
+        xlim = [
+            min(xlim[0], np.min(xpeaks) * rho),
+            max(xlim[1], np.max(xpeaks) * rho)
+        ]
 
         ypeaks = np.where(
             np.diff(np.abs(locus.imag)) < 0, locus.imag[0:-1], 0)
-        ylim = max(ylim, np.max(ypeaks))
-
-    # Adjust the limits to include some space around features
-    # TODO: use _k_max and project out to max k for all value?
-    rho = config._get_param('pzmap', 'expansion_factor')
-    xlim[0] = rho * xlim[0] if xlim[0] < 0 else 0
-    xlim[1] = rho * xlim[1] if xlim[1] > 0 else 0
-    ylim = rho * ylim if ylim > 0 else np.max(np.abs(xlim))
+        ylim = max(ylim, np.max(ypeaks) * rho)
+
+    if isctime(dt=response.dt):
+        # Adjust the limits to include some space around features
+        # TODO: use _k_max and project out to max k for all value?
+        rho = config._get_param('pzmap', 'expansion_factor')
+        xlim[0] = rho * xlim[0] if xlim[0] < 0 else 0
+        xlim[1] = rho * xlim[1] if xlim[1] > 0 else 0
+        ylim = rho * ylim if ylim > 0 else np.max(np.abs(xlim))
 
     return xlim, [-ylim, ylim]
 
