murrayrm · Jan 5, 2018 · Mar 21, 2017 · Mar 21, 2017 · Mar 21, 2017 · Mar 21, 2017
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -59,7 +59,7 @@
 
 # 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)
@@ -77,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
     -------
@@ -99,9 +101,17 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
         mymat = _RLFindRoots(nump, denp, kvect)
         mymat = _RLSortRoots(mymat)
 
-    # Create the plot
-    if (Plot):
-        f = pylab.figure()
+    # 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)
+
         if PrintGain:
             f.canvas.mpl_connect(
                 'button_release_event', partial(_RLFeedbackClicks, sys=sys))
@@ -127,6 +137,8 @@ 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(f)
 
     return mymat, kvect
 
@@ -150,6 +162,7 @@ def _default_gains(num, den, xlim, ylim):
         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((singular_points, np.zeros(2)), axis=0)
     mymat_xl = np.append(mymat_xl, important_points)
 
     if xlim is None:
@@ -305,6 +318,58 @@ 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):
+    ax = fig.gca()
+    ylocator = ax.get_yaxis().get_major_locator()
+    xlocator = ax.get_yaxis().get_major_locator()
+
+    angules = np.arange(-90, 80, 15)*np.pi/180
+
+    # zeta-constant lines
+    y_over_x = np.tan(angules[1::])*ylocator()[-1]/xlocator()[-1]
+    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
+    index = 0
+    zeta = np.sin(np.pi/2-angules[1::])
+
+    for yp in y_over_x:
+        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)
+        elif 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
+                ytext_pos = - ytext_pos
+            else:
+                ytext_pos = ylim[0]
+                xtext_pos = -xtext_pos
+            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
+    for xt in xlocator():
+        if xt < 0:
+            yp = np.sin(angules)*np.abs(xt)
+            xp = -np.cos(angules)*np.abs(xt)
+            ax.plot(xp, yp, color='gray',
+                    linestyle='dashed', linewidth=0.5)
+            an = "%.2f" % -xt
+            ax.annotate(an, textcoords='data', xy=[xt, 0], fontsize=8)
 
 rlocus = root_locus