python-control · murrayrm · Dec 28, 2018 · Feb 19, 2018 · Feb 19, 2018 · Feb 19, 2018
diff --git a/control/grid.py b/control/grid.py
@@ -0,0 +1,182 @@
+import numpy as np
+from numpy import cos, sin, sqrt, linspace, pi, exp
+import matplotlib.pyplot as plt
+from mpl_toolkits.axisartist import SubplotHost
+from mpl_toolkits.axisartist.grid_helper_curvelinear import GridHelperCurveLinear
+import mpl_toolkits.axisartist.angle_helper as angle_helper
+from matplotlib.projections import PolarAxes
+from matplotlib.transforms import Affine2D
+
+class FormatterDMS(object):
+    '''Transforms angle ticks to damping ratios'''
+    def __call__(self,direction,factor,values):
+        angles_deg = values/factor
+        damping_ratios = np.cos((180-angles_deg)*np.pi/180)
+        ret = ["%.2f"%val for val in damping_ratios]
+        return ret
+
+class ModifiedExtremeFinderCycle(angle_helper.ExtremeFinderCycle):
+    '''Changed to allow only left hand-side polar grid'''
+    def __call__(self, transform_xy, x1, y1, x2, y2):
+        x_, y_ = np.linspace(x1, x2, self.nx), np.linspace(y1, y2, self.ny)
+        x, y = np.meshgrid(x_, y_)
+        lon, lat = transform_xy(np.ravel(x), np.ravel(y))
+
+        with np.errstate(invalid='ignore'):
+            if self.lon_cycle is not None:
+                lon0 = np.nanmin(lon)
+                lon -= 360. * ((lon - lon0) > 360.) # Changed from 180 to 360 to be able to span only 90-270 (left hand side)
+            if self.lat_cycle is not None:
+                lat0 = np.nanmin(lat)
+                lat -= 360. * ((lat - lat0) > 360.) # Changed from 180 to 360 to be able to span only 90-270 (left hand side)
+
+        lon_min, lon_max = np.nanmin(lon), np.nanmax(lon)
+        lat_min, lat_max = np.nanmin(lat), np.nanmax(lat)
+
+        lon_min, lon_max, lat_min, lat_max = \
+            self._adjust_extremes(lon_min, lon_max, lat_min, lat_max)
+
+        return lon_min, lon_max, lat_min, lat_max
+
+def sgrid():
+    # From matplotlib demos:
+    # https://matplotlib.org/gallery/axisartist/demo_curvelinear_grid.html
+    # https://matplotlib.org/gallery/axisartist/demo_floating_axis.html
+
+    # PolarAxes.PolarTransform takes radian. However, we want our coordinate
+    # system in degree
+    tr = Affine2D().scale(np.pi/180., 1.) + PolarAxes.PolarTransform()
+    # polar projection, which involves cycle, and also has limits in
+    # its coordinates, needs a special method to find the extremes
+    # (min, max of the coordinate within the view).
+
+    # 20, 20 : number of sampling points along x, y direction
+    sampling_points = 20
+    extreme_finder = ModifiedExtremeFinderCycle(sampling_points, sampling_points,
+                                                     lon_cycle=360,
+                                                     lat_cycle=None,
+                                                     lon_minmax=(90,270),
+                                                     lat_minmax=(0, np.inf),)
+
+    grid_locator1 = angle_helper.LocatorDMS(15)
+    tick_formatter1 = FormatterDMS()
+    grid_helper = GridHelperCurveLinear(tr,
+                                        extreme_finder=extreme_finder,
+                                        grid_locator1=grid_locator1,
+                                        tick_formatter1=tick_formatter1
+                                        )
+
+    fig = plt.figure()
+    ax = SubplotHost(fig, 1, 1, 1, grid_helper=grid_helper)
+
+    # make ticklabels of right invisible, and top axis visible.
+    visible = True
+    ax.axis[:].major_ticklabels.set_visible(visible)
+    ax.axis[:].major_ticks.set_visible(False)
+    ax.axis[:].invert_ticklabel_direction()
+
+    ax.axis["wnxneg"] = axis = ax.new_floating_axis(0, 180)
+    axis.set_ticklabel_direction("-")
+    axis.label.set_visible(False)
+    ax.axis["wnxpos"] = axis = ax.new_floating_axis(0, 0)
+    axis.label.set_visible(False)
+    ax.axis["wnypos"] = axis = ax.new_floating_axis(0, 90)
+    axis.label.set_visible(False)
+    axis.set_axis_direction("left")
+    ax.axis["wnyneg"] = axis = ax.new_floating_axis(0, 270)
+    axis.label.set_visible(False)
+    axis.set_axis_direction("left")
+    axis.invert_ticklabel_direction()
+    axis.set_ticklabel_direction("-")
+
+    # let left axis shows ticklabels for 1st coordinate (angle)
+    ax.axis["left"].get_helper().nth_coord_ticks = 0
+    ax.axis["right"].get_helper().nth_coord_ticks = 0
+    ax.axis["left"].get_helper().nth_coord_ticks = 0
+    ax.axis["bottom"].get_helper().nth_coord_ticks = 0
+
+    fig.add_subplot(ax)
+
+    ### RECTANGULAR X Y AXES WITH SCALE
+    #par2 = ax.twiny()
+    #par2.axis["top"].toggle(all=False)
+    #par2.axis["right"].toggle(all=False)
+    #new_fixed_axis = par2.get_grid_helper().new_fixed_axis
+    #par2.axis["left"] = new_fixed_axis(loc="left",
+    #                                   axes=par2,
+    #                                   offset=(0, 0))
+    #par2.axis["bottom"] = new_fixed_axis(loc="bottom",
+    #                                     axes=par2,
+    #                                     offset=(0, 0))
+    ### FINISH RECTANGULAR
+
+    ax.grid(True, zorder=0,linestyle='dotted')
+
+    _final_setup(ax)
+    return ax, fig
+
+def _final_setup(ax):
+    ax.set_xlabel('Real')
+    ax.set_ylabel('Imaginary')
+    ax.axhline(y=0, color='black', lw=1)
+    ax.axvline(x=0, color='black', lw=1)
+    plt.axis('equal')
+
+def nogrid():
+    f = plt.figure()
+    ax = plt.axes()
+
+    _final_setup(ax)
+    return ax, f
+
+def zgrid(zetas=None, wns=None):
+    '''Draws discrete damping and frequency grid'''
+
+    fig = plt.figure()
+    ax = fig.gca()
+
+    # Constant damping lines
+    if zetas is None:
+        zetas = linspace(0, 0.9, 10)
+    for zeta in zetas:
+        # Calculate in polar coordinates
+        factor = zeta/sqrt(1-zeta**2)
+        x = linspace(0, sqrt(1-zeta**2),200)
+        ang = pi*x
+        mag = exp(-pi*factor*x)
+        # Draw upper part in retangular coordinates
+        xret = mag*cos(ang)
+        yret = mag*sin(ang)
+        ax.plot(xret,yret, 'k:', lw=1)
+        # Draw lower part in retangular coordinates
+        xret = mag*cos(-ang)
+        yret = mag*sin(-ang)
+        ax.plot(xret,yret,'k:', lw=1)
+        # Annotation
+        an_i = int(len(xret)/2.5)
+        an_x = xret[an_i]
+        an_y = yret[an_i]
+        ax.annotate(str(round(zeta,2)), xy=(an_x, an_y), xytext=(an_x, an_y), size=7)
+
+    # Constant natural frequency lines
+    if wns is None:
+        wns = linspace(0, 1, 10)
+    for a in wns:
+        # Calculate in polar coordinates
+        x = linspace(-pi/2,pi/2,200)
+        ang = pi*a*sin(x)
+        mag = exp(-pi*a*cos(x))
+        # Draw in retangular coordinates
+        xret = mag*cos(ang)
+        yret = mag*sin(ang)
+        ax.plot(xret,yret,'k:', lw=1)
+        # Annotation
+        an_i = -1
+        an_x = xret[an_i]
+        an_y = yret[an_i]
+        num = '{:1.1f}'.format(a)
+        ax.annotate("$\\frac{"+num+"\pi}{T}$", xy=(an_x, an_y), xytext=(an_x, an_y), size=9)
+
+    _final_setup(ax)
+    return ax, fig
+
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -40,15 +40,17 @@
 #
 # $Id:pzmap.py 819 2009-05-29 21:28:07Z murray $
 
-from numpy import real, imag
-from .lti import LTI
+from numpy import real, imag, linspace, exp, cos, sin, sqrt
+from math import pi
+from .lti import LTI, isdtime, isctime
+from .grid import sgrid, zgrid, nogrid
 
 __all__ = ['pzmap']
 
 # TODO: Implement more elegant cross-style axes. See:
 #    http://matplotlib.sourceforge.net/examples/axes_grid/demo_axisline_style.html
 #    http://matplotlib.sourceforge.net/examples/axes_grid/demo_curvelinear_grid.html
-def pzmap(sys, Plot=True, title='Pole Zero Map'):
+def pzmap(sys, Plot=True, grid=False, title='Pole Zero Map'):
     """
     Plot a pole/zero map for a linear system.
 
@@ -59,6 +61,8 @@ def pzmap(sys, Plot=True, title='Pole Zero Map'):
     Plot: bool
         If ``True`` a graph is generated with Matplotlib,
         otherwise the poles and zeros are only computed and returned.
+    grid: boolean (default = False)
+        If True plot omega-damping grid.
 
     Returns
     -------
@@ -75,18 +79,22 @@ def pzmap(sys, Plot=True, title='Pole Zero Map'):
 
     if (Plot):
         import matplotlib.pyplot as plt
+
+        if grid:
+            if isdtime(sys, strict=True):
+                ax, fig = zgrid()
+            else:
+                ax, fig = sgrid()
+        else:
+            ax, fig = nogrid()
+
         # Plot the locations of the poles and zeros
         if len(poles) > 0:
-            plt.scatter(real(poles), imag(poles), s=50, marker='x')
+            ax.scatter(real(poles), imag(poles), s=50, marker='x', facecolors='k')
         if len(zeros) > 0:
-            plt.scatter(real(zeros), imag(zeros), s=50, marker='o',
-                        facecolors='none')
-        # Add axes
-        #Somewhat silly workaround
-        plt.axhline(y=0, color='black')
-        plt.axvline(x=0, color='black')
-        plt.xlabel('Re')
-        plt.ylabel('Im')
+            ax.scatter(real(zeros), imag(zeros), s=50, marker='o',
+                        facecolors='none', edgecolors='k')
+
 
         plt.title(title)
 

diff --git a/control/rlocus.py b/control/rlocus.py
@@ -47,12 +47,15 @@
 
 # Packages used by this module
 import numpy as np
-from scipy import array, poly1d, row_stack, zeros_like, real, imag
+from scipy import array, poly1d, row_stack, zeros_like, real, imag, exp, sin, cos, linspace, sqrt
+from math import pi
 import scipy.signal             # signal processing toolbox
 import pylab                    # plotting routines
 from .xferfcn import _convertToTransferFunction
 from .exception import ControlMIMONotImplemented
 from functools import partial
+from .lti import isdtime
+from .grid import sgrid, zgrid, nogrid
 
 __all__ = ['root_locus', 'rlocus']
 
@@ -82,7 +85,7 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=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. 
+        If True plot omega-damping grid.
 
     Returns
     -------
@@ -110,12 +113,18 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
         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 grid:
+            if isdtime(sys, strict=True):
+                ax, f = zgrid()
+            else:
+                ax, f = sgrid()
+        else:
+            ax, f = nogrid()
+        pylab.title(new_figure_name)
 
         if PrintGain:
             f.canvas.mpl_connect(
                 'button_release_event', partial(_RLFeedbackClicks, sys=sys))
-        ax = pylab.axes()
 
         # plot open loop poles
         poles = array(denp.r)
@@ -128,17 +137,13 @@ def root_locus(sys, kvect=None, xlim=None, ylim=None, plotstr='-', Plot=True,
 
         # Now plot the loci
         for col in mymat.T:
-            ax.plot(real(col), imag(col), plotstr)
+            ax.plot(real(col), imag(col), plotstr, lw=3)
 
         # Set up plot axes and labels
         if xlim:
             ax.set_xlim(xlim)
         if ylim:
             ax.set_ylim(ylim)
-        ax.set_xlabel('Real')
-        ax.set_ylabel('Imaginary')
-        if grid:
-            _sgrid_func()
     return mymat, kvect
 
 
@@ -350,88 +355,4 @@ def _RLFeedbackClicks(event, sys):
         print("Clicked at %10.4g%+10.4gj gain %10.4g damp %10.4g" %
               (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