Fixed axes ratio 'equal' when grid=None for discrete rlocus

billtubbs · billtubbs · commit decc8b871b75 · 2021-08-14T23:36:44.000-04:00
diff --git a/control/grid.py b/control/grid.py
@@ -151,13 +151,12 @@ def _final_setup(ax):
     ax.axis('equal')
 
 
-def nogrid():
-    # TODO: this doesn't look safe. convert to take ax argument
-    f = plt.gcf()
-    ax = plt.axes()
-
+def nogrid(ax=None):
+    if ax is None:
+        ax = plt.gca()
+    fig = ax.figure
     _final_setup(ax)
-    return ax, f
+    return ax, fig
 
 
 def zgrid(zetas=None, wns=None, ax=None):
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -59,7 +59,7 @@
 from .xferfcn import _convert_to_transfer_function
 from .exception import ControlMIMONotImplemented
 from .sisotool import _SisotoolUpdate
-from .grid import sgrid, zgrid
+from .grid import sgrid, zgrid, nogrid
 from . import config
 
 __all__ = ['root_locus', 'rlocus']
@@ -220,27 +220,31 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr=None,
         for index, col in enumerate(mymat.T):
             ax.plot(real(col), imag(col), plotstr, label='rootlocus')
 
-        # Set up plot axes and labels
-        ax.set_xlabel('Real')
-        ax.set_ylabel('Imaginary')
-
         # Draw the grid
         if grid:
             if isdtime(sys, strict=True):
                 zgrid(ax=ax)
             else:
                 _sgrid_func(fig=fig if sisotool else None)
         else:
+            # TODO: could some of the settings below be done with nogrid:
+            #nogrid(ax=ax)
+            # nogrid calls grid._final_setup() which has slightly different
+            # settings for axhline and axvline linewidths and no zorder
+            # setting. For now I am leaving the original commands below:
             ax.axhline(0., linestyle=':', color='k', linewidth=.75, zorder=-20)
             ax.axvline(0., linestyle=':', color='k', linewidth=.75, zorder=-20)
             if isdtime(sys, strict=True):
                 ax.add_patch(plt.Circle(
                     (0, 0), radius=1.0, linestyle=':', edgecolor='k',
                     linewidth=0.75, fill=False, zorder=-20))
+            ax.set_xlabel('Real')
+            ax.set_ylabel('Imaginary')
+            ax.axis('equal')
 
         # Set the axes limits
-        # Note: I moved this to after computing grid lines above because
-        # something in above was changing the xlim for discrete zgrid.
+        # Note: I moved this to after the grid lines code above because
+        # something in above was changing the xlim for discrete zgrids
         if xlim:
             ax.set_xlim(xlim)
         if ylim:
