python-control · murrayrm · Jan 5, 2018 · Mar 21, 2017 · Mar 21, 2017 · Mar 21, 2017
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -56,9 +56,10 @@
 
 __all__ = ['root_locus', 'rlocus']
 
+
 # Main function: compute a root locus diagram
 def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
-               PrintGain=True):
+               PrintGain=True, grid=False):
     """Root locus plot
 
     Calculate the root locus by finding the roots of 1+k*TF(s)
@@ -76,10 +77,12 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
     ylim : tuple or list, optional
         control of y-axis range
     Plot : boolean, optional (default = True)
-        If True, plot magnitude and phase
+        If True, plot root locus diagram.
     PrintGain: boolean (default = True)
         If True, report mouse clicks when close to the root-locus branches,
         calculate gain, damping and print
+    grid: boolean (default = False)
+        If True plot s-plane grid. 
 
     Returns
     -------
@@ -93,15 +96,22 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
     (nump, denp) = _systopoly1d(sys)
 
     if kvect is None:
-        kvect = _default_gains(sys)
-
-    # Compute out the loci
-    mymat = _RLFindRoots(sys, kvect)
-    mymat = _RLSortRoots(sys, mymat)
+        kvect, mymat, xlim, ylim = _default_gains(nump, denp, xlim, ylim)
+    else:
+        mymat = _RLFindRoots(nump, denp, kvect)
+        mymat = _RLSortRoots(mymat)
+
+    # Create the Plot
+    if Plot:
+        figure_number = pylab.get_fignums()
+        figure_title = [pylab.figure(numb).canvas.get_window_title() for numb in figure_number]
+        new_figure_name = "Root Locus"
+        rloc_num = 1
+        while new_figure_name in figure_title:
+            new_figure_name = "Root Locus " + str(rloc_num)
+            rloc_num += 1
+        f = pylab.figure(new_figure_name)
 
-    # Create the plot
-    if (Plot):
-        f = pylab.figure()
         if PrintGain:
             f.canvas.mpl_connect(
                 'button_release_event', partial(_RLFeedbackClicks, sys=sys))
@@ -113,7 +123,7 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
 
         # plot open loop zeros
         zeros = array(nump.r)
-        if zeros.any():
+        if zeros.size > 0:
             ax.plot(real(zeros), imag(zeros), 'o')
 
         # Now plot the loci
@@ -127,14 +137,139 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
             ax.set_ylim(ylim)
         ax.set_xlabel('Real')
         ax.set_ylabel('Imaginary')
-
+        if grid:
+            _sgrid_func()
     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 _default_gains(num, den, xlim, ylim):
+    """Unsupervised gains calculation for root locus plot. 
+
+    References:
+     Ogata, K. (2002). Modern control engineering (4th ed.). Upper Saddle River, NJ : New Delhi: Prentice Hall.."""
+
+    k_break, real_break = _break_points(num, den)
+    kmax = _k_max(num, den, real_break, k_break)
+    kvect = np.hstack((np.linspace(0, kmax, 50), np.real(k_break)))
+    kvect.sort()
+    mymat = _RLFindRoots(num, den, kvect)
+    mymat = _RLSortRoots(mymat)
+    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(open_loop_zeros,
+                                       np.ones(open_loop_poles.size - open_loop_zeros.size) * open_loop_zeros[-1])
+        mymat_xl = np.append(mymat, open_loop_zeros_xl)
+    else:
+        mymat_xl = mymat
+    singular_points = np.concatenate((num.roots, den.roots), axis=0)
+    important_points = np.concatenate((singular_points, real_break), axis=0)
+    important_points = np.concatenate((important_points, np.zeros(2)), axis=0)
+    mymat_xl = np.append(mymat_xl, important_points)
+    false_gain = den.coeffs[0] / num.coeffs[0]
+    if false_gain < 0 and not den.order > num.order:
+        raise ValueError("Not implemented support for 0 degrees root "
+                         "locus with equal order of numerator and denominator.")
+
+    if xlim is None and false_gain > 0:
+        x_tolerance = 0.05 * (np.max(np.real(mymat_xl)) - np.min(np.real(mymat_xl)))
+        xlim = _ax_lim(mymat_xl)
+    elif xlim is None and false_gain < 0:
+        axmin = np.min(np.real(important_points))-(np.max(np.real(important_points))-np.min(np.real(important_points)))
+        axmin = np.min(np.array([axmin, np.min(np.real(mymat_xl))]))
+        axmax = np.max(np.real(important_points))+np.max(np.real(important_points))-np.min(np.real(important_points))
+        axmax = np.max(np.array([axmax, np.max(np.real(mymat_xl))]))
+        xlim = [axmin, axmax]
+        x_tolerance = 0.05 * (axmax - axmin)
+    else:
+        x_tolerance = 0.05 * (xlim[1] - xlim[0])
+
+    if ylim is None:
+        y_tolerance = 0.05 * (np.max(np.imag(mymat_xl)) - np.min(np.imag(mymat_xl)))
+        ylim = _ax_lim(mymat_xl * 1j)
+    else:
+        y_tolerance = 0.05 * (ylim[1] - ylim[0])
+
+    tolerance = np.max([x_tolerance, y_tolerance])
+    distance_points = np.abs(np.diff(mymat, axis=0))
+    indexes_too_far = np.where(distance_points > tolerance)
+
+    while (indexes_too_far[0].size > 0) and (kvect.size < 5000):
+        for index in indexes_too_far[0]:
+            new_gains = np.linspace(kvect[index], kvect[index+1], 5)
+            new_points = _RLFindRoots(num, den, new_gains[1:4])
+            kvect = np.insert(kvect, index+1, new_gains[1:4])
+            mymat = np.insert(mymat, index+1, new_points, axis=0)
+
+        mymat = _RLSortRoots(mymat)
+        distance_points = np.abs(np.diff(mymat, axis=0)) > tolerance  # distance between points
+        indexes_too_far = np.where(distance_points)
+
+    new_gains = kvect[-1] * np.hstack((np.logspace(0, 3, 4)))
+    new_points = _RLFindRoots(num, den, new_gains[1:4])
+    kvect = np.append(kvect, new_gains[1:4])
+    mymat = np.concatenate((mymat, new_points), axis=0)
+    mymat = _RLSortRoots(mymat)
+    return kvect, mymat, xlim, ylim
+
+
+def _break_points(num, den):
+    """Extract break points over real axis and the gains give these location"""
+    # type: (np.poly1d, np.poly1d) -> (np.array, np.array)
+    dnum = num.deriv(m=1)
+    dden = den.deriv(m=1)
+    polynom = den * dnum - num * dden
+    real_break_pts = polynom.r
+    real_break_pts = real_break_pts[num(real_break_pts) != 0]  # don't care about infinite break points
+    k_break = -den(real_break_pts) / num(real_break_pts)
+    idx = k_break >= 0   # only positives gains
+    k_break = k_break[idx]
+    real_break_pts = real_break_pts[idx]
+    if len(k_break) == 0:
+        k_break = [0]
+        real_break_pts = den.roots
+    return k_break, real_break_pts
+
+
+def _ax_lim(mymat):
+    """Utility to get the axis limits"""
+    axmin = np.min(np.real(mymat))
+    axmax = np.max(np.real(mymat))
+    if axmax != axmin:
+        deltax = (axmax - axmin) * 0.02
+    else:
+        deltax = np.max([1., axmax / 2])
+    axlim = [axmin - deltax, axmax + deltax]
+    return axlim
+
+
+def _k_max(num, den, real_break_points, k_break_points):
+    """" Calculate the maximum gain for the root locus shown in the figure"""
+    asymp_number = den.order - num.order
+    singular_points = np.concatenate((num.roots, den.roots), axis=0)
+    important_points = np.concatenate((singular_points, real_break_points), axis=0)
+    false_gain = den.coeffs[0] / num.coeffs[0]
+
+    if asymp_number > 0:
+        asymp_center = (np.sum(den.roots) - np.sum(num.roots))/asymp_number
+        distance_max = 4 * np.max(np.abs(important_points - asymp_center))
+        asymp_angles = (2 * np.arange(0, asymp_number)-1) * np.pi / asymp_number
+        if false_gain > 0:
+            farthest_points = asymp_center + distance_max * np.exp(asymp_angles * 1j)  # farthest points over asymptotes
+        else:
+            asymp_angles = asymp_angles + np.pi
+            farthest_points = asymp_center + distance_max * np.exp(asymp_angles * 1j)  # farthest points over asymptotes
+        kmax_asymp = np.real(np.abs(den(farthest_points) / num(farthest_points)))
+    else:
+        kmax_asymp = np.abs([np.abs(den.coeffs[0]) / np.abs(num.coeffs[0]) * 3])
+
+    kmax = np.max(np.concatenate((np.real(kmax_asymp), np.real(k_break_points)), axis=0))
+    if np.abs(false_gain) > kmax:
+        kmax = np.abs(false_gain)
+    return kmax
+
+
 def _systopoly1d(sys):
     """Extract numerator and denominator polynomails for a system"""
     # Allow inputs from the signal processing toolbox
@@ -157,29 +292,33 @@ def _systopoly1d(sys):
     # Check to see if num, den are already polynomials; otherwise convert
     if (not isinstance(nump, poly1d)):
         nump = poly1d(nump)
+
     if (not isinstance(denp, poly1d)):
         denp = poly1d(denp)
 
     return (nump, denp)
 
 
-def _RLFindRoots(sys, kvect):
+def _RLFindRoots(nump, denp, kvect):
     """Find the roots for the root locus."""
 
     # Convert numerator and denominator to polynomials if they aren't
-    (nump, denp) = _systopoly1d(sys)
-
     roots = []
     for k in kvect:
         curpoly = denp + k * nump
         curroots = curpoly.r
+        if len(curroots) < denp.order:
+            # if I have fewer poles than open loop, it is because i have one at infinity
+            curroots = np.insert(curroots, len(curroots), np.inf)
+
         curroots.sort()
         roots.append(curroots)
+
     mymat = row_stack(roots)
     return mymat
 
 
-def _RLSortRoots(sys, mymat):
+def _RLSortRoots(mymat):
     """Sort the roots from sys._RLFindRoots, so that the root
     locus doesn't show weird pseudo-branches as roots jump from
     one branch to another."""
@@ -209,6 +348,90 @@ def _RLFeedbackClicks(event, sys):
     K = -1./sys.horner(s)
     if abs(K.real) > 1e-8 and abs(K.imag/K.real) < 0.04:
         print("Clicked at %10.4g%+10.4gj gain %10.4g damp %10.4g" %
-              (s.real, s.imag, K.real, -1*s.real/abs(s)))
+              (s.real, s.imag, K.real, -1 * s.real / abs(s)))
+
+
+def _sgrid_func(fig=None, zeta=None, wn=None):
+    if fig is None:
+        fig = pylab.gcf()
+    ax = fig.gca()
+    xlocator = ax.get_xaxis().get_major_locator()
+
+    ylim = ax.get_ylim()
+    ytext_pos_lim = ylim[1] - (ylim[1] - ylim[0]) * 0.03
+    xlim = ax.get_xlim()
+    xtext_pos_lim = xlim[0] + (xlim[1] - xlim[0]) * 0.0
+
+    if zeta is None:
+        zeta = _default_zetas(xlim, ylim)
+
+    angules = []
+    for z in zeta:
+        if (z >= 1e-4) and (z <= 1):
+            angules.append(np.pi/2 + np.arcsin(z))
+        else:
+            zeta.remove(z)
+    y_over_x = np.tan(angules)
+
+    # zeta-constant lines
+
+    index = 0
+
+    for yp in y_over_x:
+        ax.plot([0, xlocator()[0]], [0, yp*xlocator()[0]], color='gray',
+                linestyle='dashed', linewidth=0.5)
+        ax.plot([0, xlocator()[0]], [0, -yp * xlocator()[0]], color='gray',
+                linestyle='dashed', linewidth=0.5)
+        an = "%.2f" % zeta[index]
+        if yp < 0:
+            xtext_pos = 1/yp * ylim[1]
+            ytext_pos = yp * xtext_pos_lim
+            if np.abs(xtext_pos) > np.abs(xtext_pos_lim):
+                xtext_pos = xtext_pos_lim
+            else:
+                ytext_pos = ytext_pos_lim
+            ax.annotate(an, textcoords='data', xy=[xtext_pos, ytext_pos], fontsize=8)
+        index += 1
+    ax.plot([0, 0], [ylim[0], ylim[1]], color='gray', linestyle='dashed', linewidth=0.5)
+
+    angules = np.linspace(-90, 90, 20)*np.pi/180
+    if wn is None:
+        wn = _default_wn(xlocator(), ylim)
+
+    for om in wn:
+        if om < 0:
+            yp = np.sin(angules)*np.abs(om)
+            xp = -np.cos(angules)*np.abs(om)
+            ax.plot(xp, yp, color='gray',
+                    linestyle='dashed', linewidth=0.5)
+            an = "%.2f" % -om
+            ax.annotate(an, textcoords='data', xy=[om, 0], fontsize=8)
+
+
+def _default_zetas(xlim, ylim):
+    """Return default list of dumps coefficients"""
+    sep1 = -xlim[0]/4
+    ang1 = [np.arctan((sep1*i)/ylim[1]) for i in np.arange(1, 4, 1)]
+    sep2 = ylim[1] / 3
+    ang2 = [np.arctan(-xlim[0]/(ylim[1]-sep2*i)) for i in np.arange(1, 3, 1)]
+
+    angules = np.concatenate((ang1, ang2))
+    angules = np.insert(angules, len(angules), np.pi/2)
+    zeta = np.sin(angules)
+    return zeta.tolist()
+
+
+def _default_wn(xloc, ylim):
+    """Return default wn for root locus plot"""
+
+    wn = xloc
+    sep = xloc[1]-xloc[0]
+    while np.abs(wn[0]) < ylim[1]:
+        wn = np.insert(wn, 0, wn[0]-sep)
+
+    while len(wn) > 7:
+        wn = wn[0:-1:2]
+
+    return wn
 
 rlocus = root_locus