python-control/control/pzmap.py at main · python-control/python-control

676 lines (567 loc) · 25.4 KB
# pzmap.py - computations involving poles and zeros
# Initial author: Richard M. Murray
# Creation date: 7 Sep 2009
"""Computations involving poles and zeros.
This module contains functions that compute poles, zeros and related
quantities for a linear system, as well as the main functions for
storing and plotting pole/zero and root locus diagrams.  (The actual
computation of root locus diagrams is in rlocus.py.)
import itertools
import warnings
import matplotlib.pyplot as plt
import numpy as np
from numpy import imag, real
from . import config
from .config import _process_legacy_keyword
from .ctrlplot import ControlPlot, _get_color, _get_color_offset, \
    _get_line_labels, _process_ax_keyword, _process_legend_keywords, \
    _process_line_labels, _update_plot_title
from .grid import nogrid, sgrid, zgrid
from .iosys import isctime, isdtime
from .statesp import StateSpace
from .xferfcn import TransferFunction
__all__ = ['pole_zero_map', 'pole_zero_plot', 'pzmap', 'PoleZeroData',
           'PoleZeroList']
# Define default parameter values for this module
_pzmap_defaults = {
    'pzmap.grid': False,                 # Plot omega-damping grid
    'pzmap.marker_size': 6,             # Size of the markers
    'pzmap.marker_width': 1.5,          # Width of the markers
    'pzmap.expansion_factor': 1.8,      # Amount to scale plots beyond features
    'pzmap.buffer_factor': 1.05,        # Buffer to leave around plot peaks
# Classes for keeping track of pzmap plots
# The PoleZeroData class keeps track of the information that is on a
# pole/zero plot.
# In addition to the locations of poles and zeros, you can also save a set
# of gains and loci for use in generating a root locus plot.  The gain
# variable is a 1D array consisting of a list of increasing gains.  The
# loci variable is a 2D array indexed by [gain_idx, root_idx] that can be
# plotted using the `pole_zero_plot` function.
# The PoleZeroList class is used to return a list of pole/zero plots.  It
# is a lightweight wrapper on the built-in list class that includes a
# `plot` method, allowing plotting a set of root locus diagrams.
class PoleZeroData:
    """Pole/zero data object.
    This class is used as the return type for computing pole/zero responses
    and root locus diagrams.  It contains information on the location of
    system poles and zeros, as well as the gains and loci for root locus
    diagrams.
    Parameters
    ----------
    poles : ndarray
        1D array of system poles.
    zeros : ndarray
        1D array of system zeros.
    gains : ndarray, optional
        1D array of gains for root locus plots.
    loci : ndarray, optional
        2D array of poles, with each row corresponding to a gain.
    sysname : str, optional
        System name.
    sys : `StateSpace` or `TransferFunction`, optional
        System corresponding to the data.
    dt : None, True or float, optional
        System timebase (used for showing stability boundary).
    sort_loci : bool, optional
        Set to False to turn off sorting of loci into unique branches.
    def __init__(
            self, poles, zeros, gains=None, loci=None, dt=None, sysname=None,
            sys=None, sort_loci=True):
        from .rlocus import _RLSortRoots
        self.poles = poles
        self.zeros = zeros
        self.gains = gains
        if loci is not None and sort_loci:
            self.loci = _RLSortRoots(loci)
        else:
            self.loci = loci
        self.dt = dt
        self.sysname = sysname
        self.sys = sys
    # Implement functions to allow legacy assignment to tuple
    def __iter__(self):
        return iter((self.poles, self.zeros))
    def plot(self, *args, **kwargs):
        """Plot the pole/zero data.
        See `pole_zero_plot` for description of arguments and keywords.
        """
        return pole_zero_plot(self, *args, **kwargs)
class PoleZeroList(list):
    """List of PoleZeroData objects with plotting capability."""
    def plot(self, *args, **kwargs):
        """Plot pole/zero data.
        See `pole_zero_plot` for description of arguments and keywords.
        """
        return pole_zero_plot(self, *args, **kwargs)
    def replot(self, cplt: ControlPlot):
        """Update the pole/zero loci of an existing plot.
        Parameters
        ----------
        cplt : `ControlPlot`
            Graphics handles of the existing plot.
        """
        pole_zero_replot(self, cplt)
# Pole/zero map
def pole_zero_map(sysdata):
    """Compute the pole/zero map for an LTI system.
    Parameters
    ----------
    sysdata : `StateSpace` or `TransferFunction`
        Linear system for which poles and zeros are computed.
    Returns
    -------
    pzmap_data : `PoleZeroMap`
        Pole/zero map containing the poles and zeros of the system.  Use
        ``pzmap_data.plot()`` or ``pole_zero_plot(pzmap_data)`` to plot the
        pole/zero map.
    # Convert the first argument to a list
    syslist = sysdata if isinstance(sysdata, (list, tuple)) else [sysdata]
    responses = []
    for idx, sys in enumerate(syslist):
        responses.append(
            PoleZeroData(
                sys.poles(), sys.zeros(), dt=sys.dt, sysname=sys.name))
    if isinstance(sysdata, (list, tuple)):
        return PoleZeroList(responses)
        return responses[0]
# TODO: Implement more elegant cross-style axes. See:
#   https://matplotlib.org/2.0.2/examples/axes_grid/demo_axisline_style.html
#   https://matplotlib.org/2.0.2/examples/axes_grid/demo_curvelinear_grid.html
def pole_zero_plot(
        data, plot=None, grid=None, title=None, color=None, marker_size=None,
        marker_width=None, xlim=None, ylim=None, interactive=None, ax=None,
        scaling=None, initial_gain=None, label=None, **kwargs):
    """Plot a pole/zero map for a linear system.
    If the system data include root loci, a root locus diagram for the
    system is plotted.  When the root locus for a single system is plotted,
    clicking on a location on the root locus will mark the gain on all
    branches of the diagram and show the system gain and damping for the
    given pole in the axes title.  Set to False to turn off this behavior.
    Parameters
    ----------
    data : List of `PoleZeroData` objects or `LTI` systems
        List of pole/zero response data objects generated by pzmap_response()
        or root_locus_map() that are to be plotted.  If a list of systems
        is given, the poles and zeros of those systems will be plotted.
    grid : bool or str, optional
        If True plot omega-damping grid, if False show imaginary
        axis for continuous-time systems, unit circle for discrete-time
        systems.  If 'empty', do not draw any additional lines.  Default
        value is set by `config.defaults['pzmap.grid']` or
        `config.defaults['rlocus.grid']`.
    plot : bool, optional
        (legacy) If True a graph is generated with matplotlib,
        otherwise the poles and zeros are only computed and returned.
        If this argument is present, the legacy value of poles and
        zeros is returned.
    Returns
    -------
    cplt : `ControlPlot` object
        Object containing the data that were plotted.  See `ControlPlot`
        for more detailed information.
    cplt.lines : array of list of `matplotlib.lines.Line2D`
        The shape of the array is given by (nsys, 2) where nsys is the number
        of systems or responses passed to the function.  The second index
        specifies the pzmap object type:
            - lines[idx, 0]: poles
            - lines[idx, 1]: zeros
    cplt.axes : 2D array of `matplotlib.axes.Axes`
        Axes for each subplot.
    cplt.figure : `matplotlib.figure.Figure`
        Figure containing the plot.
    cplt.legend : 2D array of `matplotlib.legend.Legend`
        Legend object(s) contained in the plot.
    poles, zeros : list of arrays
        (legacy) If the `plot` keyword is given, the system poles and zeros
        are returned.
    Other Parameters
    ----------------
    ax : `matplotlib.axes.Axes`, optional
        The matplotlib axes to draw the figure on.  If not specified and
        the current figure has a single axes, that axes is used.
        Otherwise, a new figure is created.
    color : matplotlib color spec, optional
        Specify the color of the markers and lines.
    initial_gain : float, optional
        If given, the specified system gain will be marked on the plot.
    interactive : bool, optional
        Turn off interactive mode for root locus plots.
    label : str or array_like of str, optional
        If present, replace automatically generated label(s) with given
        label(s).  If data is a list, strings should be specified for each
        system.
    legend_loc : int or str, optional
        Include a legend in the given location. Default is 'upper right',
        with no legend for a single response.  Use False to suppress legend.
    marker_color : str, optional
        Set the color of the markers used for poles and zeros.
    marker_size : int, optional
        Set the size of the markers used for poles and zeros.
    marker_width : int, optional
        Set the line width of the markers used for poles and zeros.
    rcParams : dict
        Override the default parameters used for generating plots.
        Default is set by `config.defaults['ctrlplot.rcParams']`.
    scaling : str or list, optional
        Set the type of axis scaling.  Can be 'equal' (default), 'auto', or
        a list of the form [xmin, xmax, ymin, ymax].
    show_legend : bool, optional
        Force legend to be shown if True or hidden if False.  If
        None, then show legend when there is more than one line on the
        plot or `legend_loc` has been specified.
    title : str, optional
        Set the title of the plot.  Defaults to plot type and system name(s).
    xlim : list, optional
        Set the limits for the x axis.
    ylim : list, optional
        Set the limits for the y axis.
    By default, the pzmap function calls matplotlib.pyplot.axis('equal'),
    which means that trying to reset the axis limits may not behave as
    expected.  To change the axis limits, use the `scaling` keyword of use
    matplotlib.pyplot.gca().axis('auto') and then set the axis limits to
    the desired values.
    Pole/zero plots that use the continuous-time omega-damping grid do not
    work with the `ax` keyword argument, due to the way that axes grids
    are implemented.  The `grid` argument must be set to False or
    'empty' when using the `ax` keyword argument.
    The limits of the pole/zero plot are set based on the location features
    in the plot, including the location of poles, zeros, and local maxima
    of root locus curves.  The locations of local maxima are expanded by a
    buffer factor set by `config.defaults['phaseplot.buffer_factor']` that is
    applied to the locations of the local maxima.  The final axis limits
    are set to by the largest features in the plot multiplied by an
    expansion factor set by `config.defaults['phaseplot.expansion_factor']`.
    The default value for the buffer factor is 1.05 (5% buffer around local
    maxima) and the default value for the expansion factor is 1.8 (80%
    increase in limits around the most distant features).
    # Get parameter values
    label = _process_line_labels(label)
    marker_size = config._get_param('pzmap', 'marker_size', marker_size, 6)
    marker_width = config._get_param('pzmap', 'marker_width', marker_width, 1.5)
    user_color = _process_legacy_keyword(kwargs, 'marker_color', 'color', color)
    rcParams = config._get_param('ctrlplot', 'rcParams', kwargs, pop=True)
    user_ax = ax
    xlim_user, ylim_user = xlim, ylim
    # If argument was a singleton, turn it into a tuple
    if not isinstance(data, (list, tuple)):
        data = [data]
    # If we are passed a list of systems, compute response first
    if all([isinstance(
            sys, (StateSpace, TransferFunction)) for sys in data]):
        # Get the response, popping off keywords used there
        pzmap_responses = pole_zero_map(data)
    elif all([isinstance(d, PoleZeroData) for d in data]):
        pzmap_responses = data
        raise TypeError("unknown system data type")
    # Decide whether we are plotting any root loci
    rlocus_plot = any([resp.loci is not None for resp in pzmap_responses])
    # Turn on interactive mode by default, if allowed
    if interactive is None and rlocus_plot and len(pzmap_responses) == 1 \
       and pzmap_responses[0].sys is not None:
        interactive = True
    # Legacy return value processing
    if plot is not None:
        warnings.warn(
            "pole_zero_plot() return value of poles, zeros is deprecated; "
            "use pole_zero_map()", FutureWarning)
        # Extract out the values that we will eventually return
        poles = [response.poles for response in pzmap_responses]
        zeros = [response.zeros for response in pzmap_responses]
    if plot is False:
        if len(data) == 1:
            return poles[0], zeros[0]
        else:
            return poles, zeros
    # Initialize the figure
    fig, ax = _process_ax_keyword(
        user_ax, rcParams=rcParams, squeeze=True, create_axes=False)
    legend_loc, _, show_legend = _process_legend_keywords(
        kwargs, None, 'upper right')
    # Make sure there are no remaining keyword arguments
    if kwargs:
        raise TypeError("unrecognized keywords: ", str(kwargs))
    if ax is None:
        # Determine what type of grid to use
        if rlocus_plot:
            from .rlocus import _rlocus_defaults
            grid = config._get_param('rlocus', 'grid', grid, _rlocus_defaults)
        else:
            grid = config._get_param('pzmap', 'grid', grid, _pzmap_defaults)
        # Create the axes with the appropriate grid
        with plt.rc_context(rcParams):
            if grid and grid != 'empty':
                if all([isctime(dt=response.dt) for response in data]):
                    ax, fig = sgrid(scaling=scaling)
                elif all([isdtime(dt=response.dt) for response in data]):
                    ax, fig = zgrid(scaling=scaling)
                else:
                    raise ValueError(
                        "incompatible time bases; don't know how to grid")
                # Store the limits for later use
                xlim, ylim = ax.get_xlim(), ax.get_ylim()
            elif grid == 'empty':
                ax = plt.axes()
                xlim = ylim = [np.inf, -np.inf] # use data to set limits
            else:
                ax, fig = nogrid(data[0].dt, scaling=scaling)
                xlim, ylim = ax.get_xlim(), ax.get_ylim()
        # Store the limits for later use
        xlim, ylim = ax.get_xlim(), ax.get_ylim()
        if grid is not None:
            warnings.warn("axis already exists; grid keyword ignored")
    # Get color offset for the next line to be drawn
    color_offset, color_cycle = _get_color_offset(ax)
    # Create a list of lines for the output
    out = np.empty(
        (len(pzmap_responses), 3 if rlocus_plot else 2), dtype=object)
    for i, j in itertools.product(range(out.shape[0]), range(out.shape[1])):
        out[i, j] = []          # unique list in each element
    # Plot the responses (and keep track of axes limits)
    for idx, response in enumerate(pzmap_responses):
        poles = response.poles
        zeros = response.zeros
        # Get the color to use for this response
        color = _get_color(user_color, offset=color_offset + idx)
        # Plot the locations of the poles and zeros
        if len(poles) > 0:
            if label is None:
                label_ = response.sysname if response.loci is None else None
            else:
                label_ = label[idx]
            out[idx, 0] = ax.plot(
                real(poles), imag(poles), marker='x', linestyle='',
                markeredgecolor=color, markerfacecolor=color,
                markersize=marker_size, markeredgewidth=marker_width,
                color=color, label=label_)
        if len(zeros) > 0:
            out[idx, 1] = ax.plot(
                real(zeros), imag(zeros), marker='o', linestyle='',
                markeredgecolor=color, markerfacecolor='none',
                markersize=marker_size, markeredgewidth=marker_width,
                color=color)
        # Plot the loci, if present
        if response.loci is not None:
            label_ = response.sysname if label is None else label[idx]
            for locus in response.loci.transpose():
                out[idx, 2] += ax.plot(
                    real(locus), imag(locus), color=color, label=label_)
            # Compute the axis limits to use based on the response
            resp_xlim, resp_ylim = _compute_root_locus_limits(response)
            # Keep track of the current limits
            xlim = [min(xlim[0], resp_xlim[0]), max(xlim[1], resp_xlim[1])]
            ylim = [min(ylim[0], resp_ylim[0]), max(ylim[1], resp_ylim[1])]
            # Plot the initial gain, if given
            if initial_gain is not None:
                _mark_root_locus_gain(ax, response.sys, initial_gain)
            # TODO: add arrows to root loci (reuse Nyquist arrow code?)
    # Set the axis limits to something reasonable
    if rlocus_plot:
        # Set up the limits for the plot using information from loci
        ax.set_xlim(xlim if xlim_user is None else xlim_user)
        ax.set_ylim(ylim if ylim_user is None else ylim_user)
        # No root loci => only set axis limits if users specified them
        if xlim_user is not None:
            ax.set_xlim(xlim_user)
        if ylim_user is not None:
            ax.set_ylim(ylim_user)
    # List of systems that are included in this plot
    lines, labels = _get_line_labels(ax)
    # Add legend if there is more than one system plotted
    if show_legend or len(labels) > 1 and show_legend != False:
        if response.loci is None:
            # Use "x o" for the system label, via matplotlib tuple handler
            from matplotlib.legend_handler import HandlerTuple
            from matplotlib.lines import Line2D
            line_tuples = []
            for pole_line in lines:
                zero_line = Line2D(
                    [0], [0], marker='o', linestyle='',
                    markeredgecolor=pole_line.get_markerfacecolor(),
                    markerfacecolor='none', markersize=marker_size,
                    markeredgewidth=marker_width)
                handle = (pole_line, zero_line)
                line_tuples.append(handle)
            with plt.rc_context(rcParams):
                legend = ax.legend(
                    line_tuples, labels, loc=legend_loc,
                    handler_map={tuple: HandlerTuple(ndivide=None)})
        else:
            # Regular legend, with lines
            with plt.rc_context(rcParams):
                legend = ax.legend(lines, labels, loc=legend_loc)
        legend = None
    # Add the title
    if title is None:
        title = ("Root locus plot for " if rlocus_plot
                 else "Pole/zero plot for ") + ", ".join(labels)
    if user_ax is None:
        _update_plot_title(
            title, fig, rcParams=rcParams, frame='figure',
            use_existing=False)
    # Add dispatcher to handle choosing a point on the diagram
    if interactive:
        if len(pzmap_responses) > 1:
            raise NotImplementedError(
                "interactive mode only allowed for single system")
        elif pzmap_responses[0].sys == None:
            raise SystemError("missing system information")
        else:
            sys = pzmap_responses[0].sys
        # Define function to handle mouse clicks
        def _click_dispatcher(event):
            # Find the gain corresponding to the clicked point
            K, s = _find_root_locus_gain(event, sys, ax)
            if K is not None:
                # Mark the gain on the root locus diagram
                _mark_root_locus_gain(ax, sys, K)
                # Display the parameters in the axes title
                with plt.rc_context(rcParams):
                    ax.set_title(_create_root_locus_label(sys, K, s))
            ax.figure.canvas.draw()
        fig.canvas.mpl_connect('button_release_event', _click_dispatcher)
    # Legacy processing: return locations of poles and zeros as a tuple
    if plot is True:
        if len(data) == 1:
            return poles, zeros
        else:
            TypeError("system lists not supported with legacy return values")
    return ControlPlot(out, ax, fig, legend=legend)
def pole_zero_replot(pzmap_responses, cplt):
    """Update the loci of a plot after zooming/panning.
    Parameters
    ----------
    pzmap_responses : PoleZeroMap list
        Responses to update.
    cplt : ControlPlot
        Collection of plot handles.
    for idx, response in enumerate(pzmap_responses):
        # remove the old data
        for l in cplt.lines[idx, 2]:
            l.set_data([], [])
        # update the line data
        if response.loci is not None:
            for il, locus in enumerate(response.loci.transpose()):
                    cplt.lines[idx,2][il].set_data(real(locus), imag(locus))
                except IndexError:
                    # not expected, but more lines apparently needed
                    cplt.lines[idx,2].append(cplt.ax[0,0].plot(
                        real(locus), imag(locus)))
# Utility function to find gain corresponding to a click event
def _find_root_locus_gain(event, sys, ax):
    # Get the current axis limits to set various thresholds
    xlim, ylim = ax.get_xlim(), ax.get_ylim()
    # Catch type error when event click is in the figure but not on curve
        s = complex(event.xdata, event.ydata)
        K = -1. / sys(s)
        K_xlim = -1. / sys(
            complex(event.xdata + 0.05 * abs(xlim[1] - xlim[0]), event.ydata))
        K_ylim = -1. / sys(
            complex(event.xdata, event.ydata + 0.05 * abs(ylim[1] - ylim[0])))
    except TypeError:
        K, s = float('inf'), None
        K_xlim = K_ylim = float('inf')
    # Compute tolerances for deciding if we clicked on the root locus
    # This is a bit of black magic that sets some limits for how close we
    # need to be to the root locus in order to consider it a click on the
    # actual curve.  Otherwise, we will just ignore the click.
    x_tolerance = 0.1 * abs((xlim[1] - xlim[0]))
    y_tolerance = 0.1 * abs((ylim[1] - ylim[0]))
    gain_tolerance = np.mean([x_tolerance, y_tolerance]) * 0.1 + \
        0.1 * max([abs(K_ylim.imag/K_ylim.real), abs(K_xlim.imag/K_xlim.real)])
    # Decide whether to pay attention to this event
    if abs(K.real) > 1e-8 and abs(K.imag / K.real) < gain_tolerance and \
       event.inaxes == ax.axes and K.real > 0.:
        return K.real, s
        return None, None
# Mark points corresponding to a given gain on root locus plot
def _mark_root_locus_gain(ax, sys, K):
    from .rlocus import _RLFindRoots, _systopoly1d
    # Remove any previous gain points
    for line in reversed(ax.lines):
        if line.get_label() == '_gain_point':
            line.remove()
            del line
    # Visualize clicked point, displaying all roots
    # TODO: allow marker parameters to be set
    nump, denp = _systopoly1d(sys)
    root_array = _RLFindRoots(nump, denp, K.real)
    ax.plot(
        [root.real for root in root_array], [root.imag for root in root_array],
        marker='s', markersize=6, zorder=20, label='_gain_point', color='k')
# Return a string identifying a clicked point
def _create_root_locus_label(sys, K, s):
    # Figure out the damping ratio
    if isdtime(sys, strict=True):
        zeta = -np.cos(np.angle(np.log(s)))
        zeta = -1 * s.real / abs(s)
    return "Clicked at: %.4g%+.4gj   gain = %.4g  damping = %.4g" % \
        (s.real, s.imag, K.real, zeta)
# Utility function to compute limits for root loci
def _compute_root_locus_limits(response):
    loci = response.loci
    # Start with information about zeros, if present
    if response.sys is not None and response.sys.zeros().size > 0:
        xlim = [
            min(0, np.min(response.sys.zeros().real)),
            max(0, np.max(response.sys.zeros().real))
        ylim = max(0, np.max(response.sys.zeros().imag))
        xlim, ylim = [np.inf, -np.inf], 0
    # Go through each locus and look for features
    rho = config._get_param('pzmap', 'buffer_factor')
    for locus in loci.transpose():
        # Include all starting points
        xlim = [min(xlim[0], locus[0].real), max(xlim[1], locus[0].real)]
        ylim = max(ylim, locus[0].imag)
        # Find the local maxima of root locus curve
        xpeaks = np.where(
            np.diff(np.abs(locus.real)) < 0, locus.real[0:-1], 0)
        if xpeaks.size > 0:
            xlim = [
                min(xlim[0], np.min(xpeaks) * rho),
                max(xlim[1], np.max(xpeaks) * rho)
        ypeaks = np.where(
            np.diff(np.abs(locus.imag)) < 0, locus.imag[0:-1], 0)
        if ypeaks.size > 0:
            ylim = max(ylim, np.max(ypeaks) * rho)
    if isctime(dt=response.dt):
        # Adjust the limits to include some space around features
        # TODO: use _k_max and project out to max k for all value?
        rho = config._get_param('pzmap', 'expansion_factor')
        xlim[0] = rho * xlim[0] if xlim[0] < 0 else 0
        xlim[1] = rho * xlim[1] if xlim[1] > 0 else 0
        ylim = rho * ylim if ylim > 0 else np.max(np.abs(xlim))
    # Make sure the limits make sense
    if xlim == [0, 0]:
        xlim = [-1, 1]
    if ylim == 0:
        ylim = 1
    return xlim, [-ylim, ylim]
pzmap = pole_zero_plot
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

pzmap.py

Latest commit

History

pzmap.py

File metadata and controls
