Bugfixes with ax, fig in sgrid, zgrid + ylim on rlocus

billtubbs · billtubbs · commit c5fa4e66d740 · 2021-08-14T22:43:28.000-04:00
diff --git a/control/grid.py b/control/grid.py
@@ -1,5 +1,5 @@
 import numpy as np
-from numpy import cos, sin, sqrt, linspace, pi, exp
+from numpy import cos, sin, sqrt, pi, exp
 import matplotlib.pyplot as plt
 from mpl_toolkits.axisartist import SubplotHost
 from mpl_toolkits.axisartist.grid_helper_curvelinear \
@@ -19,10 +19,9 @@ def __call__(self, direction, factor, values):
 
 
 class ModifiedExtremeFinderCycle(angle_helper.ExtremeFinderCycle):
-    '''Changed to allow only left hand-side polar grid
-
-    https://matplotlib.org/_modules/mpl_toolkits/axisartist/angle_helper.html#ExtremeFinderCycle.__call__
-    '''
+    # Changed to allow only left hand-side polar grid
+    # https://matplotlib.org/_modules/mpl_toolkits/axisartist/angle_helper.html#ExtremeFinderCycle.__call__
+    
     def __call__(self, transform_xy, x1, y1, x2, y2):
         x, y = np.meshgrid(
             np.linspace(x1, x2, self.nx), np.linspace(y1, y2, self.ny))
@@ -65,11 +64,16 @@ def __call__(self, transform_xy, x1, y1, x2, y2):
         return lon_min, lon_max, lat_min, lat_max
 
 
-def sgrid():
+def sgrid(ax=None):
+    """Draws continuous time damping and frequency grid"""
     # From matplotlib demos:
     # https://matplotlib.org/gallery/axisartist/demo_curvelinear_grid.html
     # https://matplotlib.org/gallery/axisartist/demo_floating_axis.html
 
+    if ax is None:
+        ax = plt.gca()
+    fig = ax.figure
+
     # PolarAxes.PolarTransform takes radian. However, we want our coordinate
     # system in degree
     tr = Affine2D().scale(np.pi/180., 1.) + PolarAxes.PolarTransform()
@@ -144,10 +148,11 @@ def _final_setup(ax):
     ax.set_ylabel('Imaginary')
     ax.axhline(y=0, color='black', lw=1)
     ax.axvline(x=0, color='black', lw=1)
-    plt.axis('equal')
+    ax.axis('equal')
 
 
 def nogrid():
+    # TODO: this doesn't look safe. convert to take ax argument
     f = plt.gcf()
     ax = plt.axes()
 
@@ -156,19 +161,19 @@ def nogrid():
 
 
 def zgrid(zetas=None, wns=None, ax=None):
-    '''Draws discrete damping and frequency grid'''
+    '''Draws discrete time damping and frequency grid'''
 
-    fig = plt.gcf()
     if ax is None:
-        ax = fig.gca()
+        ax = plt.gca()
+    fig = ax.figure
 
     # Constant damping lines
     if zetas is None:
-        zetas = linspace(0, 0.9, 10)
+        zetas = np.linspace(0, 0.9, 10)
     for zeta in zetas:
         # Calculate in polar coordinates
         factor = zeta/sqrt(1-zeta**2)
-        x = linspace(0, sqrt(1-zeta**2), 200)
+        x = np.linspace(0, sqrt(1-zeta**2), 200)
         ang = pi*x
         mag = exp(-pi*factor*x)
         # Draw upper part in retangular coordinates
@@ -188,10 +193,10 @@ def zgrid(zetas=None, wns=None, ax=None):
 
     # Constant natural frequency lines
     if wns is None:
-        wns = linspace(0, 1, 10)
+        wns = np.linspace(0, 1, 10)
     for a in wns:
         # Calculate in polar coordinates
-        x = linspace(-pi/2, pi/2, 200)
+        x = np.linspace(-pi/2, pi/2, 200)
         ang = pi*a*sin(x)
         mag = exp(-pi*a*cos(x))
         # Draw in retangular coordinates
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -67,15 +67,15 @@
 # Default values for module parameters
 _rlocus_defaults = {
     'rlocus.grid': True,
-    'rlocus.plotstr': 'b' if int(mpl.__version__[0]) == 1 else 'C0',
+    'rlocus.plotstr': '',
     'rlocus.print_gain': True,
     'rlocus.plot': True
 }
 
 
 # Main function: compute a root locus diagram
-def root_locus(sys, kvect=None, xlim=None, ylim=None,
-               plotstr=None, plot=True, print_gain=None, grid=None, ax=None,
+def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr=None, 
+               plot=True, print_gain=None, grid=None, ax=None,
                **kwargs):
 
     """Root locus plot
@@ -104,7 +104,7 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None,
         If True (default), report mouse clicks when close to the root locus
         branches, calculate gain, damping and print.
     grid : bool
-        If True plot omega-damping grid.  Default is False.
+        If True (default) plot omega-damping grid.
     ax : :class:`matplotlib.axes.Axes`
         Axes on which to create root locus plot
 
@@ -143,11 +143,12 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None,
     (nump, denp) = _systopoly1d(sys_loop)
 
     # if discrete-time system and if xlim and ylim are not given,
-    #  that we a view of the unit circle
-    if xlim is None and isdtime(sys, strict=True):
-        xlim = (-1.2, 1.2)
-    if ylim is None and isdtime(sys, strict=True):
-        xlim = (-1.3, 1.3)
+    # make sure that we a view of the unit circle
+    if isdtime(sys, strict=True):
+        if xlim is None:
+            xlim = (-1.2, 1.2)
+        if ylim is None:
+            ylim = (-1.3, 1.3)
 
     if kvect is None:
         start_mat = _RLFindRoots(nump, denp, [1])
@@ -168,8 +169,8 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None,
         else:
             if ax is None:
                 ax = plt.gca()
-                fig = ax.figure
-                ax.set_title('Root Locus')
+            fig = ax.figure
+            ax.set_title('Root Locus')
 
         if print_gain and not sisotool:
             fig.canvas.mpl_connect(
@@ -642,9 +643,7 @@ def _removeLine(label, ax):
 def _sgrid_func(fig=None, zeta=None, wn=None):
     if fig is None:
         fig = plt.gcf()
-        ax = fig.gca()
-    else:
-        ax = fig.axes[1]
+    ax = fig.gca()
 
     # Get locator function for x-axis, y-axis tick marks
     xlocator = ax.get_xaxis().get_major_locator()
@@ -748,7 +747,7 @@ def _default_zetas(xlim, ylim):
 
 
 def _default_wn(xloc, yloc, max_lines=7):
-    """Return default wn for root locus plot
+    """Return default natural frequencies for root locus plot
 
     This function computes a list of natural frequencies based on the grid
     parameters of the graph.
@@ -769,7 +768,7 @@ def _default_wn(xloc, yloc, max_lines=7):
 
     """
     sep = xloc[1]-xloc[0]       # separation between x-ticks
-    
+
     # Decide whether to use the x or y axis for determining wn
     if yloc[-1] / sep > max_lines*10:
         # y-axis scale >> x-axis scale
