updated unit test for coverage + some fixes

murrayrm · Jan 12, 2024 · Aug 3, 2023 · Aug 9, 2023 · Aug 12, 2023 · Aug 21, 2023
commit 52f81271f49a653cc2bab570ed8c7ac7ee07b5a3
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -205,11 +205,11 @@ def pole_zero_plot(
 
     Returns
     -------
-    lines : List of Line2D
+    lines : array of list of Line2D
         Array of Line2D objects for each set of markers in the plot. The
         shape of the array is given by (nsys, 2) where nsys is the number
-        of systems or Nyquist responses passed to the function.  The second
-        index specifies the pzmap object type:
+        of systems or responses passed to the function.  The second index
+        specifies the pzmap object type:
 
         * lines[idx, 0]: poles
         * lines[idx, 1]: zeros
@@ -276,8 +276,11 @@ def pole_zero_plot(
     else:
         raise TypeError("unknown system data type")
 
+    # Decide whether we are plotting any root loci
+    rlocus_plot = any([resp.loci is not None for resp in pzmap_responses])
+
     # Turn on interactive mode by default, if allowed
-    if interactive is None and len(pzmap_responses) == 1 \
+    if interactive is None and rlocus_plot and len(pzmap_responses) == 1 \
        and pzmap_responses[0].sys is not None:
         interactive = True
 
@@ -318,19 +321,26 @@ def pole_zero_plot(
             elif all([isdtime(dt=response.dt) for response in data]):
                 ax, fig = zgrid(scaling=scaling)
             else:
-                ValueError(
-                    "incompatible time responses; don't know how to grid")
+                raise ValueError(
+                    "incompatible time bases; don't know how to grid")
+            # Store the limits for later use
+            xlim, ylim = ax.get_xlim(), ax.get_ylim()
         elif len(axs) == 0:
             if grid == 'empty':
                 # Leave off grid entirely
                 ax = plt.axes()
+                xlim = ylim = [0, 0]    # use data to set limits
             else:
                 # draw stability boundary; use first response timebase
                 ax, fig = nogrid(data[0].dt, scaling=scaling)
+                xlim, ylim = ax.get_xlim(), ax.get_ylim()
         else:
             # Use the existing axes and any grid that is there
             ax = axs[0]
 
+            # Store the limits for later use
+            xlim, ylim = ax.get_xlim(), ax.get_ylim()
+
             # Issue a warning if the user tried to set the grid type
             if grid:
                 warnings.warn("axis already exists; grid keyword ignored")
@@ -345,13 +355,12 @@ def pole_zero_plot(
             color_offset = color_cycle.index(last_color) + 1
 
     # Create a list of lines for the output
-    out = np.empty((len(pzmap_responses), 3), dtype=object)
+    out = np.empty(
+        (len(pzmap_responses), 3 if rlocus_plot else 2), dtype=object)
     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()
-    loci_count = 0
     for idx, response in enumerate(pzmap_responses):
         poles = response.poles
         zeros = response.zeros
@@ -397,7 +406,7 @@ def pole_zero_plot(
             # TODO: add arrows to root loci (reuse Nyquist arrow code?)
 
     # Set the axis limits to something reasonable
-    if any([response.loci is not None for response in pzmap_responses]):
+    if rlocus_plot:
         # Set up the limits for the plot using information from loci
         ax.set_xlim(xlim if xlim_user is None else xlim_user)
         ax.set_ylim(ylim if ylim_user is None else ylim_user)

diff --git a/control/rlocus.py b/control/rlocus.py
@@ -133,14 +133,15 @@ def root_locus_plot(
 
     Returns
     -------
-    lines : List of Line2D
+    lines : array of list of Line2D
         Array of Line2D objects for each set of markers in the plot. The
-        shape of the array is given by (nsys, 2) where nsys is the number
-        of systems or Nyquist responses passed to the function.  The second
-        index specifies the pzmap object type:
+        shape of the array is given by (nsys, 3) where nsys is the number
+        of systems or responses passed to the function.  The second index
+        specifies the object type:
 
         * lines[idx, 0]: poles
         * lines[idx, 1]: zeros
+        * lines[idx, 2]: loci
 
     roots, gains : ndarray
         (legacy) If the `plot` keyword is given, returns the

diff --git a/control/tests/pzmap_test.py b/control/tests/pzmap_test.py
@@ -12,6 +12,7 @@
 from matplotlib import pyplot as plt
 from mpl_toolkits.axisartist import Axes as mpltAxes
 
+import control as ct
 from control import TransferFunction, config, pzmap
 
 
@@ -82,7 +83,47 @@ def test_pzmap(kwargs, setdefaults, dt, editsdefaults, mplcleanup):
         assert not plt.get_fignums()
 
 
+def test_polezerodata():
+    sys = ct.rss(4, 1, 1)
+    pzdata = ct.pole_zero_map(sys)
+    np.testing.assert_equal(pzdata.poles, sys.poles())
+    np.testing.assert_equal(pzdata.zeros, sys.zeros())
+
+    # Extract data from PoleZeroData
+    poles, zeros = pzdata
+    np.testing.assert_equal(poles, sys.poles())
+    np.testing.assert_equal(zeros, sys.zeros())
+
+    # Legacy return format
+    with pytest.warns(DeprecationWarning, match=".* values .* deprecated"):
+        poles, zeros = ct.pole_zero_plot(pzdata, plot=False)
+    np.testing.assert_equal(poles, sys.poles())
+    np.testing.assert_equal(zeros, sys.zeros())
+
+    with pytest.warns(DeprecationWarning, match=".* values .* deprecated"):
+        poles, zeros = ct.pole_zero_plot(pzdata, plot=True)
+    np.testing.assert_equal(poles, sys.poles())
+    np.testing.assert_equal(zeros, sys.zeros())
+
+
 def test_pzmap_raises():
     with pytest.raises(TypeError):
         # not an LTI system
-        pzmap(([1], [1,2]))
+        pzmap(([1], [1, 2]))
+
+    sys1 = ct.rss(2, 1, 1)
+    sys2 = sys1.sample(0.1)
+    with pytest.raises(ValueError, match="incompatible time bases"):
+        pzdata = ct.pole_zero_plot([sys1, sys2], grid=True)
+
+    with pytest.warns(UserWarning, match="axis already exists"):
+        fig, ax = plt.figure(), plt.axes()
+        ct.pole_zero_plot(sys1, ax=ax, grid='empty')
+
+
+def test_pzmap_limits():
+    sys = ct.tf([1, 2], [1, 2, 3])
+    out = ct.pole_zero_plot(sys, xlim=[-1, 1], ylim=[-1, 1])
+    ax = ct.get_plot_axes(out)[0, 0]
+    assert ax.get_xlim() == (-1, 1)
+    assert ax.get_ylim() == (-1, 1)
diff --git a/control/tests/rlocus_test.py b/control/tests/rlocus_test.py
@@ -67,12 +67,22 @@ def test_without_gains(self, sys):
         self.check_cl_poles(sys, roots, kvect)
 
     @pytest.mark.parametrize("grid", [None, True, False, 'empty'])
-    def test_root_locus_plot_grid(self, sys, grid):
+    @pytest.mark.parametrize("method", ['plot', 'map', 'response', 'pzmap'])
+    def test_root_locus_plot_grid(self, sys, grid, method):
         import mpl_toolkits.axisartist as AA
 
         # Generate the root locus plot
         plt.clf()
-        ct.root_locus_plot(sys, grid=grid)
+        if method == 'plot':
+            ct.root_locus_plot(sys, grid=grid)
+        elif method == 'map':
+            ct.root_locus_map(sys).plot(grid=grid)
+        elif method == 'response':
+            response = ct.root_locus_map(sys)
+            ct.root_locus_plot(response, grid=grid)
+        elif method == 'pzmap':
+            response = ct.root_locus_map(sys)
+            ct.pole_zero_plot(response, grid=grid)
 
         # Count the number of dotted/dashed lines in the plot
         ax = plt.gca()
@@ -85,7 +95,7 @@ def test_root_locus_plot_grid(self, sys, grid):
         if grid == 'empty':
             assert n_gridlines == 0
             assert not isinstance(ax, AA.Axes)
-        elif grid is False:
+        elif grid is False or method == 'pzmap' and grid is None:
             assert n_gridlines == 2 if sys.isctime() else 3
             assert not isinstance(ax, AA.Axes)
         elif sys.isdtime(strict=True):