Fix python-control#45 - make gain vector optional in root_locus

cwrowley · cwrowley · commit 24eb89bca40b · 2015-04-05T11:39:59.000-04:00
diff --git a/control/matlab.py b/control/matlab.py
@@ -1110,11 +1110,8 @@ def rlocus(sys, klist = None, **keywords):
     ----------
     sys: StateSpace or TransferFunction
         Linear system
-    klist:
+    klist: iterable, optional
         optional list of gains
-
-    Keyword parameters
-    ------------------
     xlim : control of x-axis range, normally with tuple, for
         other options, see matplotlib.axes
     ylim : control of y-axis range
@@ -1132,12 +1129,8 @@ def rlocus(sys, klist = None, **keywords):
         list of gains used to compute roots
     """
     from .rlocus import root_locus
-    #! TODO: update with a smart calculation of the gains using sys poles/zeros
-    if klist == None:
-        klist = logspace(-3, 3)
 
-    rlist = root_locus(sys, klist, **keywords)
-    return rlist, klist
+    return root_locus(sys, klist, **keywords)
 
 def margin(*args):
     """Calculate gain and phase margins and associated crossover frequencies
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -46,43 +46,51 @@
 # $Id$
 
 # Packages used by this module
+import numpy as np
 from scipy import array, poly1d, row_stack, zeros_like, real, imag
 import scipy.signal             # signal processing toolbox
 import pylab                    # plotting routines
 from . import xferfcn
 from .exception import ControlMIMONotImplemented
 from functools import partial
 
-
 # Main function: compute a root locus diagram
-def root_locus(sys, kvect, xlim=None, ylim=None, plotstr='-', Plot=True,
+def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
                PrintGain=True):
     """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.
 
     Parameters
     ----------
-    sys : linsys
+    sys : LTI object
         Linear input/output systems (SISO only, for now)
-    kvect : gain_range (default = None)
+    kvect : list or ndarray, optional
         List of gains to use in computing diagram
-    xlim : control of x-axis range, normally with tuple, for
-        other options, see matplotlib.axes
-    ylim : control of y-axis range
-    Plot : boolean (default = True)
+    xlim : tuple or list, optional
+        control of x-axis range, normally with tuple (see matplotlib.axes)
+    ylim : tuple or list, optional
+        control of y-axis range
+    Plot : boolean, optional (default = True)
         If True, plot magnitude and phase
     PrintGain: boolean (default = True)
         If True, report mouse clicks when close to the root-locus branches,
         calculate gain, damping and print
-    Return values
-    -------------
-    rlist : list of computed root locations
+
+    Returns
+    -------
+    rlist : ndarray
+        Computed root locations, given as a 2d array
+    klist : ndarray or list
+        Gains used.  Same as klist keyword argument if provided.
     """
 
     # Convert numerator and denominator to polynomials if they aren't
     (nump, denp) = _systopoly1d(sys)
 
+    if kvect is None:
+        kvect = _default_gains(sys)
+
     # Compute out the loci
     mymat = _RLFindRoots(sys, kvect)
     mymat = _RLSortRoots(sys, mymat)
@@ -116,8 +124,11 @@ def root_locus(sys, kvect, xlim=None, ylim=None, plotstr='-', Plot=True,
         ax.set_xlabel('Real')
         ax.set_ylabel('Imaginary')
 
-    return mymat
+    return mymat, kvect
 
+def _default_gains(sys):
+    # TODO: update with a smart calculation of the gains using sys poles/zeros
+    return np.logspace(-3, 3)
 
 # Utility function to extract numerator and denominator polynomials
 def _systopoly1d(sys):
diff --git a/control/tests/matlab_test.py b/control/tests/matlab_test.py
@@ -309,7 +309,10 @@ def testRlocus(self):
         rlocus(self.siso_ss1)
         rlocus(self.siso_tf1)
         rlocus(self.siso_tf2)
-        rlist, klist = rlocus(self.siso_tf2, klist=[1, 10, 100], Plot=False)
+        klist = [1, 10, 100]
+        rlist, klist_out = rlocus(self.siso_tf2, klist=klist, Plot=False)
+        np.testing.assert_equal(len(rlist), len(klist))
+        np.testing.assert_array_equal(klist, klist_out)
 
     def testNyquist(self):
         nyquist(self.siso_ss1)
@@ -618,7 +621,7 @@ def testCombi01(self):
 #        for i in range(len(tfdata)):
 #            np.testing.assert_array_almost_equal(tfdata_1[i], tfdata_2[i])
 
-def suite():
+def test_suite():
    return unittest.TestLoader().loadTestsFromTestCase(TestMatlab)
 
 if __name__ == '__main__':
diff --git a/control/tests/rlocus_test.py b/control/tests/rlocus_test.py
@@ -17,21 +17,32 @@ def setUp(self):
         """This contains some random LTI systems and scalars for testing."""
 
         # Two random SISO systems.
-        self.sys1 = TransferFunction([1, 2], [1, 2, 3])
-        self.sys2 = StateSpace([[1., 4.], [3., 2.]], [[1.], [-4.]],
+        sys1 = TransferFunction([1, 2], [1, 2, 3])
+        sys2 = StateSpace([[1., 4.], [3., 2.]], [[1.], [-4.]],
             [[1., 0.]], [[0.]])
+        self.systems = (sys1, sys2)
+
+    def check_cl_poles(self, sys, pole_list, k_list):
+        for k, poles in zip(k_list, pole_list):
+            poles_expected = np.sort(feedback(sys, k).pole())
+            poles = np.sort(poles)
+            np.testing.assert_array_almost_equal(poles, poles_expected)
 
     def testRootLocus(self):
         """Basic root locus plot"""
         klist = [-1, 0, 1]
-        rlist = root_locus(self.sys1, [-1, 0, 1], Plot=False)
-
-        for k in klist:
-            np.testing.assert_array_almost_equal(
-                np.sort(rlist[k]), 
-                np.sort(feedback(self.sys1, klist[k]).pole()))
-
-def suite():
+        for sys in self.systems:
+            roots, k_out = root_locus(sys, klist, Plot=False)
+            np.testing.assert_equal(len(roots), len(klist))
+            np.testing.assert_array_equal(klist, k_out)
+            self.check_cl_poles(sys, roots, klist)
+
+    def test_without_gains(self):
+        for sys in self.systems:
+            roots, kvect = root_locus(sys, Plot=False)
+            self.check_cl_poles(sys, roots, kvect)
+
+def test_suite():
     return unittest.TestLoader().loadTestsFromTestCase(TestRootLocus)
 
 if __name__ == "__main__":
