pz -> pole_zero + add loci plotting

murrayrm · murrayrm · commit 1451f3a35bae · 2023-12-28T12:41:09.000-08:00
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -22,7 +22,7 @@
 from .freqplot import _freqplot_defaults, _get_line_labels
 from . import config
 
-__all__ = ['pzmap_response', 'pzmap_plot', 'pzmap']
+__all__ = ['pole_zero_map', 'root_locus_map', 'pole_zero_plot', 'pzmap']
 
 
 # Define default parameter values for this module
@@ -34,18 +34,21 @@
 
 
 # Classes for keeping track of pzmap plots
-class PoleZeroResponseList(list):
+class RootLocusList(list):
     def plot(self, *args, **kwargs):
-        return pzmap_plot(self, *args, **kwargs)
+        return pole_zero_plot(self, *args, **kwargs)
 
 
-class PoleZeroResponseData:
+class RootLocusData:
     def __init__(
-            self, poles, zeros, gains=None, loci=None, dt=None, sysname=None):
+            self, poles, zeros, gains=None, loci=None, xlim=None, ylim=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
 
@@ -54,38 +57,78 @@ def __iter__(self):
         return iter((self.poles, self.zeros))
 
     def plot(self, *args, **kwargs):
-        return pzmap_plot(self, *args, **kwargs)
+        return pole_zero_plot(self, *args, **kwargs)
 
 
-# pzmap response funciton
-def pzmap_response(sysdata):
+# Pole/zero map
+def pole_zero_map(sysdata):
     # Convert the first argument to a list
     syslist = sysdata if isinstance(sysdata, (list, tuple)) else [sysdata]
 
     responses = []
     for idx, sys in enumerate(syslist):
         responses.append(
-            PoleZeroResponseData(
+            RootLocusData(
                 sys.poles(), sys.zeros(), dt=sys.dt, sysname=sys.name))
 
     if isinstance(sysdata, (list, tuple)):
-        return PoleZeroResponseList(responses)
+        return RootLocusList(responses)
+    else:
+        return responses[0]
+
+
+# Root locus map
+def root_locus_map(sysdata, gains=None, xlim=None, ylim=None):
+    # Convert the first argument to a list
+    syslist = sysdata if isinstance(sysdata, (list, tuple)) else [sysdata]
+
+    responses = []
+    for idx, sys in enumerate(syslist):
+        from .rlocus import _systopoly1d, _default_gains
+        from .rlocus import _RLFindRoots, _RLSortRoots
+
+        if not sys.issiso():
+            raise ControlMIMONotImplemented(
+                "sys must be single-input single-output (SISO)")
+
+        # 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)
+        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))
+
+    if isinstance(sysdata, (list, tuple)):
+        return RootLocusList(responses)
     else:
         return responses[0]
 
 
 # TODO: Implement more elegant cross-style axes. See:
 #    https://matplotlib.org/2.0.2/examples/axes_grid/demo_axisline_style.html
 #    https://matplotlib.org/2.0.2/examples/axes_grid/demo_curvelinear_grid.html
-def pzmap_plot(
+def pole_zero_plot(
         data, plot=None, grid=None, title=None, marker_color=None,
         marker_size=None, marker_width=None, legend_loc='upper right',
-        **kwargs):
+        xlim=None, ylim=None, **kwargs):
     """Plot a pole/zero map for a linear system.
 
     Parameters
     ----------
-    sysdata: List of PoleZeroResponseData objects or LTI systems
+    sysdata: List of RootLocusData objects or LTI systems
         List of pole/zero response data objects generated by pzmap_response
         or rootlocus_response() that are to be plotted.  If a list of systems
         is given, the poles and zeros of those systems will be plotted.
@@ -124,6 +167,7 @@ def pzmap_plot(
     grid = config._get_param('pzmap', 'grid', grid, False)
     marker_size = config._get_param('pzmap', 'marker_size', marker_size, 6)
     marker_width = config._get_param('pzmap', 'marker_width', marker_width, 1.5)
+    xlim_user, ylim_user = xlim, ylim
     freqplot_rcParams = config._get_param(
         'freqplot', 'rcParams', kwargs, _freqplot_defaults,
         pop=True, last=True)
@@ -136,17 +180,17 @@ def pzmap_plot(
     if all([isinstance(
             sys, (StateSpace, TransferFunction)) for sys in data]):
         # Get the response, popping off keywords used there
-        pzmap_responses = pzmap_response(data)
-    elif all([isinstance(d, PoleZeroResponseData) for d in data]):
+        pzmap_responses = pole_zero_map(data)
+    elif all([isinstance(d, RootLocusData) for d in data]):
         pzmap_responses = data
     else:
         raise TypeError("unknown system data type")
 
     # Legacy return value processing
     if plot is not None:
         warnings.warn(
-            "`pzmap_plot` return values of poles, zeros is deprecated; "
-            "use pzmap_response()", DeprecationWarning)
+            "`pole_zero_plot` return values of poles, zeros is deprecated; "
+            "use pole_zero_map()", DeprecationWarning)
 
         # Extract out the values that we will eventually return
         poles = [response.poles for response in pzmap_responses]
@@ -169,9 +213,9 @@ def pzmap_plot(
     with plt.rc_context(freqplot_rcParams):
         if grid:
             plt.clf()
-            if all([response.isctime() for response in data]):
+            if all([response.dt in [0, None] for response in data]):
                 ax, fig = sgrid()
-            elif all([response.isdtime() for response in data]):
+            elif all([response.dt > 0 for response in data]):
                 ax, fig = zgrid()
             else:
                 ValueError(
@@ -192,10 +236,11 @@ def pzmap_plot(
             color_offset = color_cycle.index(last_color) + 1
 
     # Create a list of lines for the output
-    out = np.empty((len(pzmap_responses), 2), dtype=object)
+    out = np.empty((len(pzmap_responses), 3), dtype=object)
     for i, j in itertools.product(range(out.shape[0]), range(out.shape[1])):
         out[i, j] = []          # unique list in each element
 
+    xlim, ylim = ax.get_xlim(), ax.get_ylim()
     for idx, response in enumerate(pzmap_responses):
         poles = response.poles
         zeros = response.zeros
@@ -208,44 +253,65 @@ def pzmap_plot(
 
         # Plot the locations of the poles and zeros
         if len(poles) > 0:
+            label = response.sysname if response.loci is None else None
             out[idx, 0] = ax.plot(
                 real(poles), imag(poles), marker='x', linestyle='',
                 markeredgecolor=color, markerfacecolor=color,
                 markersize=marker_size, markeredgewidth=marker_width,
-                label=response.sysname)
+                label=label)
         if len(zeros) > 0:
             out[idx, 1] = ax.plot(
                 real(zeros), imag(zeros), marker='o', linestyle='',
                 markeredgecolor=color, markerfacecolor='none',
                 markersize=marker_size, markeredgewidth=marker_width)
 
+        # Plot the loci, if present
+        if response.loci is not None:
+            for locus in response.loci.transpose():
+                out[idx, 2] += ax.plot(
+                    real(locus), imag(locus), color=color,
+                    label=response.sysname)
+
+        # Compute the axis limits to use
+        xlim = (min(xlim[0], response.xlim[0]), max(xlim[1], response.xlim[1]))
+        ylim = (min(ylim[0], response.ylim[0]), max(ylim[1], response.ylim[1]))
+
+    # Set up the limits for the plot
+    ax.set_xlim(xlim if xlim_user is None else xlim_user)
+    ax.set_ylim(ylim if ylim_user is None else ylim_user)
+
     # List of systems that are included in this plot
     lines, labels = _get_line_labels(ax)
 
-    # Update the lines to use tuples for poles and zeros
-    from matplotlib.lines import Line2D
-    from matplotlib.legend_handler import HandlerTuple
-    line_tuples = []
-    for pole_line in lines:
-        zero_line = Line2D(
-            [0], [0], marker='o', linestyle='',
-            markeredgecolor=pole_line.get_markerfacecolor(),
-            markerfacecolor='none', markersize=marker_size,
-            markeredgewidth=marker_width)
-        handle = (pole_line, zero_line)
-        line_tuples.append(handle)
-    print(line_tuples)
-
     # Add legend if there is more than one system plotted
     if len(labels) > 1 and legend_loc is not False:
-        with plt.rc_context(freqplot_rcParams):
-            ax.legend(
-                line_tuples, labels, loc=legend_loc,
-                handler_map={tuple: HandlerTuple(ndivide=None)})
+        if response.loci is None:
+            # Use "x o" for the system label, via matplotlib tuple handler
+            from matplotlib.lines import Line2D
+            from matplotlib.legend_handler import HandlerTuple
+
+            line_tuples = []
+            for pole_line in lines:
+                zero_line = Line2D(
+                    [0], [0], marker='o', linestyle='',
+                    markeredgecolor=pole_line.get_markerfacecolor(),
+                    markerfacecolor='none', markersize=marker_size,
+                    markeredgewidth=marker_width)
+            handle = (pole_line, zero_line)
+            line_tuples.append(handle)
+
+            with plt.rc_context(freqplot_rcParams):
+                ax.legend(
+                    line_tuples, labels, loc=legend_loc,
+                    handler_map={tuple: HandlerTuple(ndivide=None)})
+        else:
+            # Regular legend, with lines
+            with plt.rc_context(freqplot_rcParams):
+                ax.legend(lines, labels, loc=legend_loc)
 
     # Add the title
     if title is None:
-        title = "Pole/zero map for " + ", ".join(labels)
+        title = "Pole/zero plot for " + ", ".join(labels)
     with plt.rc_context(freqplot_rcParams):
         fig.suptitle(title)
 
@@ -259,4 +325,4 @@ def pzmap_plot(
     return out
 
 
-pzmap = pzmap_plot
+pzmap = pole_zero_plot
diff --git a/control/tests/kwargs_test.py b/control/tests/kwargs_test.py
@@ -241,7 +241,7 @@ def test_response_plot_kwargs(data_fcn, plot_fcn, mimo):
     'nyquist_response': test_response_plot_kwargs,
     'nyquist_plot': test_matplotlib_kwargs,
     'pzmap': test_unrecognized_kwargs,
-    'pzmap_plot': test_unrecognized_kwargs,
+    'pole_zero_plot': test_unrecognized_kwargs,
     'rlocus': test_unrecognized_kwargs,
     'root_locus': test_unrecognized_kwargs,
     'rss': test_unrecognized_kwargs,
