update xlim, ylim handling in pzmap + other supporting changes

murrayrm · murrayrm · commit 6e335981ec98 · 2023-12-28T12:41:09.000-08:00
diff --git a/control/grid.py b/control/grid.py
@@ -8,6 +8,8 @@
 from matplotlib.projections import PolarAxes
 from matplotlib.transforms import Affine2D
 
+from .iosys import isdtime
+
 
 class FormatterDMS(object):
     '''Transforms angle ticks to damping ratios'''
@@ -142,17 +144,22 @@ def sgrid():
 def _final_setup(ax):
     ax.set_xlabel('Real')
     ax.set_ylabel('Imaginary')
-    ax.axhline(y=0, color='black', lw=1)
-    ax.axvline(x=0, color='black', lw=1)
+    ax.axhline(y=0, color='black', lw=0.5)
+    ax.axvline(x=0, color='black', lw=0.5)
     plt.axis('equal')
 
 
-def nogrid():
-    f = plt.gcf()
+def nogrid(dt=None):
+    fig = plt.gcf()
     ax = plt.axes()
 
+    # Draw the unit circle for discrete time systems
+    if isdtime(dt=dt, strict=True):
+        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)
-    return ax, f
+    return ax, fig
 
 
 def zgrid(zetas=None, wns=None, ax=None):
diff --git a/control/iosys.py b/control/iosys.py
@@ -503,42 +503,64 @@ def common_timebase(dt1, dt2):
         raise ValueError("Systems have incompatible timebases")
 
 # Check to see if a system is a discrete time system
-def isdtime(sys, strict=False):
+def isdtime(sys=None, dt=None, strict=False):
     """
     Check to see if a system is a discrete time system.
 
     Parameters
     ----------
-    sys : I/O or LTI system
-        System to be checked
+    sys : I/O system, optional
+        System to be checked.
+    dt : None or number, optional
+        Timebase to be checked.
     strict: bool (default = False)
-        If strict is True, make sure that timebase is not None
+        If strict is True, make sure that timebase is not None.
     """
 
-    # Check to see if this is a constant
+    # See if we were passed a timebase instead of a system
+    if sys is None:
+        if dt is None:
+            return True if not strict else False
+        else:
+            return dt > 0
+    elif dt is not None:
+        raise TypeError("passing both system and timebase not allowed")
+
+    # Check timebase of the system
     if isinstance(sys, (int, float, complex, np.number)):
-        # OK as long as strict checking is off
+        # Constants OK as long as strict checking is off
         return True if not strict else False
     else:
         return sys.isdtime(strict)
 
 
 # Check to see if a system is a continuous time system
-def isctime(sys, strict=False):
+def isctime(sys=None, dt=None, strict=False):
     """
     Check to see if a system is a continuous-time system.
 
     Parameters
     ----------
-    sys : I/O or LTI system
-        System to be checked
+    sys : I/O system, optional
+        System to be checked.
+    dt : None or number, optional
+        Timebase to be checked.
     strict: bool (default = False)
-        If strict is True, make sure that timebase is not None
+        If strict is True, make sure that timebase is not None.
     """
 
-    # Check to see if this is a constant
+    # See if we were passed a timebase instead of a system
+    if sys is None:
+        if dt is None:
+            return True if not strict else False
+        else:
+            return dt == 0
+    elif dt is not None:
+        raise TypeError("passing both system and timebase not allowed")
+
+    # Check timebase of the system
     if isinstance(sys, (int, float, complex, np.number)):
-        # OK as long as strict checking is off
+        # Constants OK as long as strict checking is off
         return True if not strict else False
     else:
         return sys.isctime(strict)
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -4,7 +4,9 @@
 # Date: 7 Sep 2009
 #
 # This file contains functions that compute poles, zeros and related
-# quantities for a linear system.
+# quantities for a linear system, as well as the main functions for
+# storing and plotting pole/zero and root locus diagrams.  (The actual
+# computation of root locus diagrams is in rlocus.py.)
 #
 
 import numpy as np
@@ -41,14 +43,11 @@ def plot(self, *args, **kwargs):
 
 class RootLocusData:
     def __init__(
-            self, poles, zeros, gains=None, loci=None, xlim=None, ylim=None,
-            dt=None, sysname=None):
+            self, poles, zeros, gains=None, loci=None, dt=None, sysname=None):
         self.poles = poles
         self.zeros = zeros
         self.gains = gains
         self.loci = loci
-        self.xlim = xlim
-        self.ylim = ylim
         self.dt = dt
         self.sysname = sysname
 
@@ -78,7 +77,8 @@ def pole_zero_map(sysdata):
 
 
 # Root locus map
-def root_locus_map(sysdata, gains=None, xlim=None, ylim=None):
+# TODO: use rlocus.py computation instead
+def root_locus_map(sysdata, gains=None):
     # Convert the first argument to a list
     syslist = sysdata if isinstance(sysdata, (list, tuple)) else [sysdata]
 
@@ -94,22 +94,16 @@ def root_locus_map(sysdata, gains=None, xlim=None, ylim=None):
         # Convert numerator and denominator to polynomials if they aren't
         nump, denp = _systopoly1d(sys[0, 0])
 
-        if xlim is None and sys.isdtime(strict=True):
-            xlim = (-1.2, 1.2)
-        if ylim is None and sys.isdtime(strict=True):
-            xlim = (-1.3, 1.3)
-
         if gains is None:
-            kvect, root_array, xlim, ylim = _default_gains(
-                nump, denp, xlim, ylim)
+            kvect, root_array, _, _ = _default_gains(nump, denp, None, None)
         else:
             kvect = np.atleast_1d(gains)
             root_array = _RLFindRoots(nump, denp, kvect)
             root_array = _RLSortRoots(root_array)
 
         responses.append(RootLocusData(
             sys.poles(), sys.zeros(), kvect, root_array,
-            dt=sys.dt, sysname=sys.name, xlim=xlim, ylim=ylim))
+            dt=sys.dt, sysname=sys.name))
 
     if isinstance(sysdata, (list, tuple)):
         return RootLocusList(responses)
@@ -203,7 +197,7 @@ def pole_zero_plot(
             return poles, zeros
 
     # Initialize the figure
-    # TODO: turn into standard utility function
+    # TODO: turn into standard utility function (from plotutil.py?)
     fig = plt.gcf()
     axs = fig.get_axes()
     if len(axs) > 1:
@@ -213,21 +207,22 @@ def pole_zero_plot(
     with plt.rc_context(freqplot_rcParams):
         if grid:
             plt.clf()
-            if all([response.dt in [0, None] for response in data]):
+            if all([isctime(dt=response.dt) for response in data]):
                 ax, fig = sgrid()
-            elif all([response.dt > 0 for response in data]):
+            elif all([isdtime(dt=response.dt) for response in data]):
                 ax, fig = zgrid()
             else:
                 ValueError(
                     "incompatible time responses; don't know how to grid")
         elif len(axs) == 0:
-            ax, fig = nogrid()
+            ax, fig = nogrid(data[0].dt)        # use first response timebase
         else:
-            # Use the existing axes
+            # Use the existing axes and any grid that is there
+            # TODO: allow axis to be overriden via parameter
             ax = axs[0]
 
-    # Handle color cycle manually as all singular values
-    # of the same systems are expected to be of the same color
+    # Handle color cycle manually as all root locus segments
+    # of the same system are expected to be of the same color
     color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']
     color_offset = 0
     if len(ax.lines) > 0:
@@ -240,6 +235,7 @@ def pole_zero_plot(
     for i, j in itertools.product(range(out.shape[0]), range(out.shape[1])):
         out[i, j] = []          # unique list in each element
 
+    # Plot the responses (and keep track of axes limits)
     xlim, ylim = ax.get_xlim(), ax.get_ylim()
     for idx, response in enumerate(pzmap_responses):
         poles = response.poles
@@ -272,13 +268,14 @@ def pole_zero_plot(
                     real(locus), imag(locus), color=color,
                     label=response.sysname)
 
-        # Compute the axis limits to use
-        if response.xlim is not None:
-            xlim = (min(xlim[0], response.xlim[0]),
-                    max(xlim[1], response.xlim[1]))
-        if response.ylim is not None:
-            ylim = (min(ylim[0], response.ylim[0]),
-                    max(ylim[1], response.ylim[1]))
+            # Compute the axis limits to use based on the response
+            resp_xlim, resp_ylim = _compute_root_locus_limits(response.loci)
+
+            # Keep track of the current limits
+            xlim = [min(xlim[0], resp_xlim[0]), max(xlim[1], resp_xlim[1])]
+            ylim = [min(ylim[0], resp_ylim[0]), max(ylim[1], resp_ylim[1])]
+
+            # TODO: add arrows to root loci (reuse Nyquist arrow code?)
 
     # Set up the limits for the plot
     ax.set_xlim(xlim if xlim_user is None else xlim_user)
@@ -329,4 +326,31 @@ def pole_zero_plot(
     return out
 
 
+# Utility function to compute limits for root loci
+def _compute_root_locus_limits(loci):
+    # Go through each locus
+    xlim, ylim = [0, 0], 0
+    for locus in loci.transpose():
+        # Include all starting points
+        xlim = [min(xlim[0], locus[0].real), max(xlim[1], locus[0].real)]
+        ylim = max(ylim, locus[0].imag)
+
+        # 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))]
+
+        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
+    rho = 1.5
+    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]
+
+
 pzmap = pole_zero_plot
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -32,6 +32,7 @@
 __all__ = ['root_locus', 'rlocus']
 
 # Default values for module parameters
+# TODO: merge these with pzmap parameters (?)
 _rlocus_defaults = {
     'rlocus.grid': True,
     'rlocus.plotstr': 'b' if int(mpl.__version__[0]) == 1 else 'C0',
@@ -41,15 +42,16 @@
 
 
 # Main function: compute a root locus diagram
+# TODO: update to use pzmap data structures and plotting
 def root_locus(sys, kvect=None, xlim=None, ylim=None,
                plotstr=None, plot=True, print_gain=None, grid=None, ax=None,
                initial_gain=None, **kwargs):
 
     """Root locus plot.
 
-    Calculate the root locus by finding the roots of 1+k*TF(s)
-    where TF is self.num(s)/self.den(s) and each k is an element
-    of kvect.
+    Calculate the root locus by finding the roots of 1 + k * G(s) where G
+    is a linear system with transfer function num(s)/den(s) and each k is
+    an element of kvect.
 
     Parameters
     ----------
@@ -127,6 +129,7 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None,
         raise TypeError("unrecognized keywords: ", str(kwargs))
 
     # Create the Plot
+    # TODO: replace with pole_zero_plot and move additional functionality there
     if plot:
         if ax is None:
             ax = plt.gca()
@@ -139,6 +142,7 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None,
         else:
             start_roots = None
 
+        # TODO: don't rely on `start_roots` (sisotool holdover)
         if print_gain and start_roots is None:
             fig.canvas.mpl_connect(
                 'button_release_event',
@@ -148,7 +152,8 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None,
             ax.plot(
                 [root.real for root in start_roots],
                 [root.imag for root in start_roots],
-                marker='s', markersize=6, zorder=20, color='k', label='gain_point')
+                marker='s', markersize=6, zorder=20, color='k',
+                label='gain_point')
             s = start_roots[0][0]
             if isdtime(sys, strict=True):
                 zeta = -np.cos(np.angle(np.log(s)))
@@ -219,16 +224,23 @@ def _default_gains(num, den, xlim, ylim, zoom_xlim=None, zoom_ylim=None):
     Saddle River, NJ : New Delhi: Prentice Hall..
 
     """
+    # Compute the break points on the real axis for the root locus plot
     k_break, real_break = _break_points(num, den)
+
+    # Decide on the maximum gain to use and create the gain vector
     kmax = _k_max(num, den, real_break, k_break)
     kvect = np.hstack((np.linspace(0, kmax, 50), np.real(k_break)))
     kvect.sort()
 
+    # Find the roots for all of the gains and sort them
     root_array = _RLFindRoots(num, den, kvect)
     root_array = _RLSortRoots(root_array)
+
+    # Keep track of the open loop poles and zeros
     open_loop_poles = den.roots
     open_loop_zeros = num.roots
 
+    # ???
     if open_loop_zeros.size != 0 and \
        open_loop_zeros.size < open_loop_poles.size:
         open_loop_zeros_xl = np.append(
@@ -283,7 +295,8 @@ def _default_gains(num, den, xlim, ylim, zoom_xlim=None, zoom_ylim=None):
         tolerance = x_tolerance
     else:
         tolerance = np.min([x_tolerance, y_tolerance])
-    indexes_too_far = _indexes_filt(root_array, tolerance, zoom_xlim, zoom_ylim)
+    indexes_too_far = _indexes_filt(
+        root_array, tolerance, zoom_xlim, zoom_ylim)
 
     # Add more points into the root locus for points that are too far apart
     while len(indexes_too_far) > 0 and kvect.size < 5000:
@@ -295,7 +308,8 @@ def _default_gains(num, den, xlim, ylim, zoom_xlim=None, zoom_ylim=None):
             root_array = np.insert(root_array, index + 1, new_points, axis=0)
 
         root_array = _RLSortRoots(root_array)
-        indexes_too_far = _indexes_filt(root_array, tolerance, zoom_xlim, zoom_ylim)
+        indexes_too_far = _indexes_filt(
+            root_array, tolerance, zoom_xlim, zoom_ylim)
 
     new_gains = kvect[-1] * np.hstack((np.logspace(0, 3, 4)))
     new_points = _RLFindRoots(num, den, new_gains[1:4])
@@ -601,6 +615,7 @@ def _removeLine(label, ax):
             del line
 
 
+# TODO: remove and replace with sgid()?
 def _sgrid_func(ax, zeta=None, wn=None):
     # Get locator function for x-axis, y-axis tick marks
     xlocator = ax.get_xaxis().get_major_locator()
diff --git a/control/tests/rlocus_test.py b/control/tests/rlocus_test.py
