add legacy processing for root_locus

murrayrm · murrayrm · commit 0137056dace0 · 2023-12-28T12:41:09.000-08:00
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -75,7 +75,7 @@ def root_locus_map(sysdata, gains=None):
 
 
 def root_locus_plot(
-        sysdata, kvect=None, grid=None, plot=True, **kwargs):
+        sysdata, kvect=None, grid=None, plot=None, **kwargs):
     """Root locus plot.
 
     Calculate the root locus by finding the roots of 1 + k * G(s) where G
@@ -131,12 +131,29 @@ def root_locus_plot(
     grid = config._get_param('rlocus', 'grid', grid, _rlocus_defaults)
 
     responses = root_locus_map(sysdata, gains=kvect)
-    # TODO: update to include legacy keyword processing (use pole_zero_plot?)
-    if plot:
-        responses.plot(grid=grid, **kwargs)
 
-    # TODO: legacy return value; update
-    return responses.loci, responses.gains
+    #
+    # Process `plot` keyword
+    #
+    # See bode_plot for a description of how this keyword is handled to
+    # support legacy implementatoins of root_locus.
+    #
+    if plot is not None:
+        warnings.warn(
+            "`root_locus` return values of loci, gains is deprecated; "
+            "use root_locus_map()", DeprecationWarning)
+
+    if plot is False:
+        return responses.loci, responses.gains
+
+    # Plot the root loci
+    out = responses.plot(grid=grid, **kwargs)
+
+    # Legacy processing: return locations of poles and zeros as a tuple
+    if plot is True:
+        return responses.loci, responses.gains
+
+    return out
 
 
 # TODO: get rid of zoom functionality?
diff --git a/control/tests/rlocus_test.py b/control/tests/rlocus_test.py
@@ -55,7 +55,7 @@ def testRootLocus(self, sys):
         self.check_cl_poles(sys, roots, klist)
 
         # now check with plotting
-        roots, k_out = root_locus(sys, klist)
+        roots, k_out = root_locus(sys, klist, plot=True)
         np.testing.assert_equal(len(roots), len(klist))
         np.testing.assert_allclose(klist, k_out)
         self.check_cl_poles(sys, roots, klist)
@@ -68,7 +68,7 @@ def test_without_gains(self, sys):
     @pytest.mark.slow
     @pytest.mark.parametrize('grid', [None, True, False])
     def test_root_locus_plot_grid(self, sys, grid):
-        rlist, klist = root_locus(sys, grid=grid)
+        rlist, klist = root_locus(sys, plot=True, grid=grid)
         ax = plt.gca()
         n_gridlines = sum([int(line.get_linestyle() in [':', 'dotted',
                                                         '--', 'dashed'])
@@ -82,7 +82,7 @@ def test_root_locus_plot_grid(self, sys, grid):
     def test_root_locus_neg_false_gain_nonproper(self):
         """ Non proper TranferFunction with negative gain: Not implemented"""
         with pytest.raises(ValueError, match="with equal order"):
-            root_locus(TransferFunction([-1, 2], [1, 2]))
+            root_locus(TransferFunction([-1, 2], [1, 2]), plot=True)
 
     # TODO: cover and validate negative false_gain branch in _default_gains()
 
@@ -93,7 +93,7 @@ def test_root_locus_zoom(self):
         """Check the zooming functionality of the Root locus plot"""
         system = TransferFunction([1000], [1, 25, 100, 0])
         plt.figure()
-        root_locus(system)
+        root_locus(system, plot=True)
         fig = plt.gcf()
         ax_rlocus = fig.axes[0]
 
@@ -135,7 +135,7 @@ def test_rlocus_default_wn(self):
         sys = ct.tf(*sp.signal.zpk2tf(
             [-1e-2, 1-1e7j, 1+1e7j], [0, -1e7j, 1e7j], 1))
 
-        ct.root_locus(sys)
+        ct.root_locus(sys, plot=True)
 
 
 # TODO: add additional test cases
