10312443515

Committed 09 Aug 2024 02:07AM UTC coverage: 94.694% (+0.04%) from 94.65%

Build # 10312443515

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #1034 from murrayrm/ctrlplot_updates-27Jun2024

Control plot refactoring for consistent functionality

Run Details

9137 of 9649 relevant lines covered (94.69%)

8.27 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.44

control/phaseplot.py

# phaseplot.py - generate 2D phase portraits
#
# Author: Richard M. Murray
# Date:   23 Mar 2024 (legacy version information below)
#
# TODO
# * Allow multiple timepoints (and change timespec name to T?)
# * Update linestyles (color -> linestyle?)
# * Check for keyword compatibility with other plot routines
# * Set up configuration parameters (nyquist --> phaseplot)

"""Module for generating 2D phase plane plots.

The :mod:`control.phaseplot` module contains functions for generating 2D
phase plots. The base function for creating phase plane portraits is
:func:`~control.phase_plane_plot`, which generates a phase plane portrait
for a 2 state I/O system (with no inputs).  In addition, several other
functions are available to create customized phase plane plots:

* boxgrid: Generate a list of points along the edge of a box
* circlegrid: Generate list of points around a circle
* equilpoints: Plot equilibrium points in the phase plane
* meshgrid: Generate a list of points forming a mesh
* separatrices: Plot separatrices in the phase plane
* streamlines: Plot stream lines in the phase plane
* vectorfield: Plot a vector field in the phase plane

"""

import math
import warnings

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import odeint

from . import config
from .ctrlplot import ControlPlot, _add_arrows_to_line2D, _get_color, \
    _process_ax_keyword, _update_plot_title
from .exception import ControlNotImplemented
from .nlsys import NonlinearIOSystem, find_eqpt, input_output_response

__all__ = ['phase_plane_plot', 'phase_plot', 'box_grid']

# Default values for module parameter variables
_phaseplot_defaults = {
    'phaseplot.arrows': 2,                  # number of arrows around curve
    'phaseplot.arrow_size': 8,              # pixel size for arrows
    'phaseplot.separatrices_radius': 0.1    # initial radius for separatrices
}

def phase_plane_plot(
        sys, pointdata=None, timedata=None, gridtype=None, gridspec=None,
        plot_streamlines=True, plot_vectorfield=False, plot_equilpoints=True,
        plot_separatrices=True, ax=None, suppress_warnings=False, title=None,
        **kwargs
):
    """Plot phase plane diagram.

    This function plots phase plane data, including vector fields, stream
    lines, equilibrium points, and contour curves.

    Parameters
    ----------
    sys : NonlinearIOSystem or callable(t, x, ...)
        I/O system or function used to generate phase plane data. If a
        function is given, the remaining arguments are drawn from the
        `params` keyword.
    pointdata : list or 2D array
        List of the form [xmin, xmax, ymin, ymax] describing the
        boundaries of the phase plot or an array of shape (N, 2)
        giving points of at which to plot the vector field.
    timedata : int or list of int
        Time to simulate each streamline.  If a list is given, a different
        time can be used for each initial condition in `pointdata`.
    gridtype : str, optional
        The type of grid to use for generating initial conditions:
        'meshgrid' (default) generates a mesh of initial conditions within
        the specified boundaries, 'boxgrid' generates initial conditions
        along the edges of the boundary, 'circlegrid' generates a circle of
        initial conditions around each point in point data.
    gridspec : list, optional
        If the gridtype is 'meshgrid' and 'boxgrid', `gridspec` gives the
        size of the grid in the x and y axes on which to generate points.
        If gridtype is 'circlegrid', then `gridspec` is a 2-tuple
        specifying the radius and number of points around each point in the
        `pointdata` array.
    params : dict, optional
        Parameters to pass to system. For an I/O system, `params` should be
        a dict of parameters and values. For a callable, `params` should be
        dict with key 'args' and value given by a tuple (passed to callable).
    color : matplotlib color spec, optional
        Plot all elements in the given color (use `plot_<fcn>={'color': c}`
        to set the color in one element of the phase plot.
    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.

    Returns
    -------
    cplt : :class:`ControlPlot` object
        Object containing the data that were plotted:

          * cplt.lines: array of list of :class:`matplotlib.artist.Artist`
            objects:

              - lines[0] = list of Line2D objects (streamlines, separatrices).
              - lines[1] = Quiver object (vector field arrows).
              - lines[2] = list of Line2D objects (equilibrium points).

          * cplt.axes: 2D array of :class:`matplotlib.axes.Axes` for the plot.

          * cplt.figure: :class:`matplotlib.figure.Figure` containing the plot.

        See :class:`ControlPlot` for more detailed information.


    Other parameters
    ----------------
    plot_streamlines : bool or dict, optional
        If `True` (default) then plot streamlines based on the pointdata
        and gridtype.  If set to a dict, pass on the key-value pairs in
        the dict as keywords to :func:`~control.phaseplot.streamlines`.
    plot_vectorfield : bool or dict, optional
        If `True` (default) then plot the vector field based on the pointdata
        and gridtype.  If set to a dict, pass on the key-value pairs in
        the dict as keywords to :func:`~control.phaseplot.vectorfield`.
    plot_equilpoints : bool or dict, optional
        If `True` (default) then plot equilibrium points based in the phase
        plot boundary. If set to a dict, pass on the key-value pairs in the
        dict as keywords to :func:`~control.phaseplot.equilpoints`.
    plot_separatrices : bool or dict, optional
        If `True` (default) then plot separatrices starting from each
        equilibrium point.  If set to a dict, pass on the key-value pairs
        in the dict as keywords to :func:`~control.phaseplot.separatrices`.
    suppress_warnings : bool, optional
        If set to `True`, suppress warning messages in generating trajectories.
    title : str, optional
        Set the title of the plot.  Defaults to plot type and system name(s).

    """
    # Process arguments
    params = kwargs.get('params', None)
    sys = _create_system(sys, params)
    pointdata = [-1, 1, -1, 1] if pointdata is None else pointdata
    rcParams = config._get_param('ctrlplot', 'rcParams', kwargs, pop=True)

    # Create axis if needed
    user_ax = ax
    fig, ax = _process_ax_keyword(user_ax, squeeze=True, rcParams=rcParams)

    # Create copy of kwargs for later checking to find unused arguments
    initial_kwargs = dict(kwargs)
    passed_kwargs = False

    # Utility function to create keyword arguments
    def _create_kwargs(global_kwargs, local_kwargs, **other_kwargs):
        new_kwargs = dict(global_kwargs)
        new_kwargs.update(other_kwargs)
        if isinstance(local_kwargs, dict):
            new_kwargs.update(local_kwargs)
        return new_kwargs

    # Create list for storing outputs
    out = np.array([[], None, None], dtype=object)

    # Plot out the main elements
    if plot_streamlines:
        kwargs_local = _create_kwargs(
            kwargs, plot_streamlines, gridspec=gridspec, gridtype=gridtype,
            ax=ax)
        out[0] += streamlines(
            sys, pointdata, timedata, check_kwargs=False,
            suppress_warnings=suppress_warnings, **kwargs_local)

        # Get rid of keyword arguments handled by streamlines
        for kw in ['arrows', 'arrow_size', 'arrow_style', 'color',
                   'dir', 'params']:
            initial_kwargs.pop(kw, None)

    # Reset the gridspec for the remaining commands, if needed
    if gridtype not in [None, 'boxgrid', 'meshgrid']:
        gridspec = None

    if plot_separatrices:
        kwargs_local = _create_kwargs(
            kwargs, plot_separatrices, gridspec=gridspec, ax=ax)
        out[0] += separatrices(
            sys, pointdata, check_kwargs=False, **kwargs_local)

        # Get rid of keyword arguments handled by separatrices
        for kw in ['arrows', 'arrow_size', 'arrow_style', 'params']:
            initial_kwargs.pop(kw, None)

    if plot_vectorfield:
        kwargs_local = _create_kwargs(
            kwargs, plot_vectorfield, gridspec=gridspec, ax=ax)
        out[1] = vectorfield(
            sys, pointdata, check_kwargs=False, **kwargs_local)

        # Get rid of keyword arguments handled by vectorfield
        for kw in ['color', 'params']:
            initial_kwargs.pop(kw, None)

    if plot_equilpoints:
        kwargs_local = _create_kwargs(
            kwargs, plot_equilpoints, gridspec=gridspec, ax=ax)
        out[2] = equilpoints(
            sys, pointdata, check_kwargs=False, **kwargs_local)

        # Get rid of keyword arguments handled by equilpoints
        for kw in ['params']:
            initial_kwargs.pop(kw, None)

    # Make sure all keyword arguments were used
    if initial_kwargs:
        raise TypeError("unrecognized keywords: ", str(initial_kwargs))

    if user_ax is None:
        if title is None:
            title = f"Phase portrait for {sys.name}"
        _update_plot_title(title, use_existing=False, rcParams=rcParams)
        ax.set_xlabel(sys.state_labels[0])
        ax.set_ylabel(sys.state_labels[1])
        plt.tight_layout()

    return ControlPlot(out, ax, fig)


def vectorfield(
        sys, pointdata, gridspec=None, ax=None, suppress_warnings=False,
        check_kwargs=True, **kwargs):
    """Plot a vector field in the phase plane.

    This function plots a vector field for a two-dimensional state
    space system.

    Parameters
    ----------
    sys : NonlinearIOSystem or callable(t, x, ...)
        I/O system or function used to generate phase plane data.  If a
        function is given, the remaining arguments are drawn from the
        `params` keyword.
    pointdata : list or 2D array
        List of the form [xmin, xmax, ymin, ymax] describing the
        boundaries of the phase plot or an array of shape (N, 2)
        giving points of at which to plot the vector field.
    gridtype : str, optional
        The type of grid to use for generating initial conditions:
        'meshgrid' (default) generates a mesh of initial conditions within
        the specified boundaries, 'boxgrid' generates initial conditions
        along the edges of the boundary, 'circlegrid' generates a circle of
        initial conditions around each point in point data.
    gridspec : list, optional
        If the gridtype is 'meshgrid' and 'boxgrid', `gridspec` gives the
        size of the grid in the x and y axes on which to generate points.
        If gridtype is 'circlegrid', then `gridspec` is a 2-tuple
        specifying the radius and number of points around each point in the
        `pointdata` array.
    params : dict or list, optional
        Parameters to pass to system. For an I/O system, `params` should be
        a dict of parameters and values. For a callable, `params` should be
        dict with key 'args' and value given by a tuple (passed to callable).
    color : matplotlib color spec, optional
        Plot the vector field in the given color.
    ax : matplotlib.axes.Axes
        Use the given axes for the plot, otherwise use the current axes.

    Returns
    -------
    out : Quiver

    Other parameters
    ----------------
    suppress_warnings : bool, optional
        If set to `True`, suppress warning messages in generating trajectories.

    """
    # Process keywords
    rcParams = config._get_param('ctrlplot', 'rcParams', kwargs, pop=True)

    # Get system parameters
    params = kwargs.pop('params', None)

    # Create system from callable, if needed
    sys = _create_system(sys, params)

    # Determine the points on which to generate the vector field
    points, _ = _make_points(pointdata, gridspec, 'meshgrid')

    # Create axis if needed
    if ax is None:
        ax = plt.gca()

    # Set the plotting limits
    xlim, ylim, maxlim = _set_axis_limits(ax, pointdata)

    # Figure out the color to use
    color = _get_color(kwargs, ax=ax)

    # Make sure all keyword arguments were processed
    if check_kwargs and kwargs:
        raise TypeError("unrecognized keywords: ", str(kwargs))

    # Generate phase plane (quiver) data
    vfdata = np.zeros((points.shape[0], 4))
    sys._update_params(params)
    for i, x in enumerate(points):
        vfdata[i, :2] = x
        vfdata[i, 2:] = sys._rhs(0, x, 0)

    with plt.rc_context(rcParams):
        out = ax.quiver(
            vfdata[:, 0], vfdata[:, 1], vfdata[:, 2], vfdata[:, 3],
            angles='xy', color=color)

    return out


def streamlines(
        sys, pointdata, timedata=1, gridspec=None, gridtype=None, dir=None,
        ax=None, check_kwargs=True, suppress_warnings=False, **kwargs):
    """Plot stream lines in the phase plane.

    This function plots stream lines for a two-dimensional state space
    system.

    Parameters
    ----------
    sys : NonlinearIOSystem or callable(t, x, ...)
        I/O system or function used to generate phase plane data.  If a
        function is given, the remaining arguments are drawn from the
        `params` keyword.
    pointdata : list or 2D array
        List of the form [xmin, xmax, ymin, ymax] describing the
        boundaries of the phase plot or an array of shape (N, 2)
        giving points of at which to plot the vector field.
    timedata : int or list of int
        Time to simulate each streamline.  If a list is given, a different
        time can be used for each initial condition in `pointdata`.
    gridtype : str, optional
        The type of grid to use for generating initial conditions:
        'meshgrid' (default) generates a mesh of initial conditions within
        the specified boundaries, 'boxgrid' generates initial conditions
        along the edges of the boundary, 'circlegrid' generates a circle of
        initial conditions around each point in point data.
    gridspec : list, optional
        If the gridtype is 'meshgrid' and 'boxgrid', `gridspec` gives the
        size of the grid in the x and y axes on which to generate points.
        If gridtype is 'circlegrid', then `gridspec` is a 2-tuple
        specifying the radius and number of points around each point in the
        `pointdata` array.
    params : dict or list, optional
        Parameters to pass to system. For an I/O system, `params` should be
        a dict of parameters and values. For a callable, `params` should be
        dict with key 'args' and value given by a tuple (passed to callable).
    color : str
        Plot the streamlines in the given color.
    ax : matplotlib.axes.Axes
        Use the given axes for the plot, otherwise use the current axes.

    Returns
    -------
    out : list of Line2D objects

    Other parameters
    ----------------
    suppress_warnings : bool, optional
        If set to `True`, suppress warning messages in generating trajectories.

    """
    # Process keywords
    rcParams = config._get_param('ctrlplot', 'rcParams', kwargs, pop=True)

    # Get system parameters
    params = kwargs.pop('params', None)

    # Create system from callable, if needed
    sys = _create_system(sys, params)

    # Parse the arrows keyword
    arrow_pos, arrow_style = _parse_arrow_keywords(kwargs)

    # Determine the points on which to generate the streamlines
    points, gridspec = _make_points(pointdata, gridspec, gridtype=gridtype)
    if dir is None:
        dir = 'both' if gridtype == 'meshgrid' else 'forward'

    # Create axis if needed
    if ax is None:
        ax = plt.gca()

    # Set the axis limits
    xlim, ylim, maxlim = _set_axis_limits(ax, pointdata)

    # Figure out the color to use
    color = _get_color(kwargs, ax=ax)

    # Make sure all keyword arguments were processed
    if check_kwargs and kwargs:
        raise TypeError("unrecognized keywords: ", str(kwargs))

    # Create reverse time system, if needed
    if dir != 'forward':
        revsys = NonlinearIOSystem(
            lambda t, x, u, params: -np.asarray(sys.updfcn(t, x, u, params)),
            sys.outfcn, states=sys.nstates, inputs=sys.ninputs,
            outputs=sys.noutputs, params=sys.params)
    else:
        revsys = None

    # Generate phase plane (streamline) data
    out = []
    for i, X0 in enumerate(points):
        # Create the trajectory for this point
        timepts = _make_timepts(timedata, i)
        traj = _create_trajectory(
            sys, revsys, timepts, X0, params, dir,
            gridtype=gridtype, gridspec=gridspec, xlim=xlim, ylim=ylim,
            suppress_warnings=suppress_warnings)

        # Plot the trajectory (if there is one)
        if traj.shape[1] > 1:
            with plt.rc_context(rcParams):
                out += ax.plot(traj[0], traj[1], color=color)

                # Add arrows to the lines at specified intervals
                _add_arrows_to_line2D(
                    ax, out[-1], arrow_pos, arrowstyle=arrow_style, dir=1)
    return out


def equilpoints(
        sys, pointdata, gridspec=None, color='k', ax=None, check_kwargs=True,
        **kwargs):
    """Plot equilibrium points in the phase plane.

    This function plots the equilibrium points for a planar dynamical system.

    Parameters
    ----------
    sys : NonlinearIOSystem or callable(t, x, ...)
        I/O system or function used to generate phase plane data. If a
        function is given, the remaining arguments are drawn from the
        `params` keyword.
    pointdata : list or 2D array
        List of the form [xmin, xmax, ymin, ymax] describing the
        boundaries of the phase plot or an array of shape (N, 2)
        giving points of at which to plot the vector field.
    gridtype : str, optional
        The type of grid to use for generating initial conditions:
        'meshgrid' (default) generates a mesh of initial conditions within
        the specified boundaries, 'boxgrid' generates initial conditions
        along the edges of the boundary, 'circlegrid' generates a circle of
        initial conditions around each point in point data.
    gridspec : list, optional
        If the gridtype is 'meshgrid' and 'boxgrid', `gridspec` gives the
        size of the grid in the x and y axes on which to generate points.
        If gridtype is 'circlegrid', then `gridspec` is a 2-tuple
        specifying the radius and number of points around each point in the
        `pointdata` array.
    params : dict or list, optional
        Parameters to pass to system. For an I/O system, `params` should be
        a dict of parameters and values. For a callable, `params` should be
        dict with key 'args' and value given by a tuple (passed to callable).
    color : str
        Plot the equilibrium points in the given color.
    ax : matplotlib.axes.Axes
        Use the given axes for the plot, otherwise use the current axes.

    Returns
    -------
    out : list of Line2D objects

    """
    # Process keywords
    rcParams = config._get_param('ctrlplot', 'rcParams', kwargs, pop=True)

    # Get system parameters
    params = kwargs.pop('params', None)

    # Create system from callable, if needed
    sys = _create_system(sys, params)

    # Create axis if needed
    if ax is None:
        ax = plt.gca()

    # Set the axis limits
    xlim, ylim, maxlim = _set_axis_limits(ax, pointdata)

    # Determine the points on which to generate the vector field
    gridspec = [5, 5] if gridspec is None else gridspec
    points, _ = _make_points(pointdata, gridspec, 'meshgrid')

    # Make sure all keyword arguments were processed
    if check_kwargs and kwargs:
        raise TypeError("unrecognized keywords: ", str(kwargs))

    # Search for equilibrium points
    equilpts = _find_equilpts(sys, points, params=params)

    # Plot the equilibrium points
    out = []
    for xeq in equilpts:
        with plt.rc_context(rcParams):
            out.append(
                ax.plot(xeq[0], xeq[1], marker='o', color=color))
    return out


def separatrices(
        sys, pointdata, timedata=None, gridspec=None, ax=None,
        check_kwargs=True, suppress_warnings=False, **kwargs):
    """Plot separatrices in the phase plane.

    This function plots separatrices for a two-dimensional state space
    system.

    Parameters
    ----------
    sys : NonlinearIOSystem or callable(t, x, ...)
        I/O system or function used to generate phase plane data. If a
        function is given, the remaining arguments are drawn from the
        `params` keyword.
    pointdata : list or 2D array
        List of the form [xmin, xmax, ymin, ymax] describing the
        boundaries of the phase plot or an array of shape (N, 2)
        giving points of at which to plot the vector field.
    timedata : int or list of int
        Time to simulate each streamline.  If a list is given, a different
        time can be used for each initial condition in `pointdata`.
    gridtype : str, optional
        The type of grid to use for generating initial conditions:
        'meshgrid' (default) generates a mesh of initial conditions within
        the specified boundaries, 'boxgrid' generates initial conditions
        along the edges of the boundary, 'circlegrid' generates a circle of
        initial conditions around each point in point data.
    gridspec : list, optional
        If the gridtype is 'meshgrid' and 'boxgrid', `gridspec` gives the
        size of the grid in the x and y axes on which to generate points.
        If gridtype is 'circlegrid', then `gridspec` is a 2-tuple
        specifying the radius and number of points around each point in the
        `pointdata` array.
    params : dict or list, optional
        Parameters to pass to system. For an I/O system, `params` should be
        a dict of parameters and values. For a callable, `params` should be
        dict with key 'args' and value given by a tuple (passed to callable).
    color : matplotlib color spec, optional
        Plot the separatrics in the given color.  If a single color
        specification is given, this is used for both stable and unstable
        separatrices.  If a tuple is given, the first element is used as
        the color specification for stable separatrices and the second
        elmeent for unstable separatrices.
    ax : matplotlib.axes.Axes
        Use the given axes for the plot, otherwise use the current axes.

    Returns
    -------
    out : list of Line2D objects

    Other parameters
    ----------------
    suppress_warnings : bool, optional
        If set to `True`, suppress warning messages in generating trajectories.

    """
    # Process keywords
    rcParams = config._get_param('ctrlplot', 'rcParams', kwargs, pop=True)

    # Get system parameters
    params = kwargs.pop('params', None)

    # Create system from callable, if needed
    sys = _create_system(sys, params)

    # Parse the arrows keyword
    arrow_pos, arrow_style = _parse_arrow_keywords(kwargs)

    # Determine the initial states to use in searching for equilibrium points
    gridspec = [5, 5] if gridspec is None else gridspec
    points, _ = _make_points(pointdata, gridspec, 'meshgrid')

    # Find the equilibrium points
    equilpts = _find_equilpts(sys, points, params=params)
    radius = config._get_param('phaseplot', 'separatrices_radius')

    # Create axis if needed
    if ax is None:
        ax = plt.gca()

    # Set the axis limits
    xlim, ylim, maxlim = _set_axis_limits(ax, pointdata)

    # Figure out the color to use for stable, unstable subspaces
    color = _get_color(kwargs)
    match color:
        case None:
            stable_color = 'r'
            unstable_color = 'b'
        case (stable_color, unstable_color) | [stable_color, unstable_color]:
            pass
        case single_color:
            stable_color = unstable_color = color

    # Make sure all keyword arguments were processed
    if check_kwargs and kwargs:
        raise TypeError("unrecognized keywords: ", str(kwargs))

    # Create a "reverse time" system to use for simulation
    revsys = NonlinearIOSystem(
        lambda t, x, u, params: -np.array(sys.updfcn(t, x, u, params)),
        sys.outfcn, states=sys.nstates, inputs=sys.ninputs,
        outputs=sys.noutputs, params=sys.params)

    # Plot separatrices by flowing backwards in time along eigenspaces
    out = []
    for i, xeq in enumerate(equilpts):
        # Plot the equilibrium points
        with plt.rc_context(rcParams):
            out.append(
                ax.plot(xeq[0], xeq[1], marker='o', color='k'))

        # Figure out the linearization and eigenvectors
        evals, evecs = np.linalg.eig(sys.linearize(xeq, 0, params=params).A)

        # See if we have real eigenvalues (=> evecs are meaningful)
        if evals[0].imag > 0:
            continue

        # Create default list of time points
        if timedata is not None:
            timepts = _make_timepts(timedata, i)

        # Generate the traces
        for j, dir in enumerate(evecs.T):
            # Figure out time vector if not yet computed
            if timedata is None:
                timescale = math.log(maxlim / radius) / abs(evals[j].real)
                timepts = np.linspace(0, timescale)

            # Run the trajectory starting in eigenvector directions
            for eps in [-radius, radius]:
                x0 = xeq + dir * eps
                if evals[j].real < 0:
                    traj = _create_trajectory(
                        sys, revsys, timepts, x0, params, 'reverse',
                        gridtype='boxgrid', xlim=xlim, ylim=ylim,
                        suppress_warnings=suppress_warnings)
                    color = stable_color
                    linestyle = '--'
                elif evals[j].real > 0:
                    traj = _create_trajectory(
                        sys, revsys, timepts, x0, params, 'forward',
                        gridtype='boxgrid', xlim=xlim, ylim=ylim,
                        suppress_warnings=suppress_warnings)
                    color = unstable_color
                    linestyle = '-'

                # Plot the trajectory (if there is one)
                if traj.shape[1] > 1:
                    with plt.rc_context(rcParams):
                        out.append(ax.plot(
                            traj[0], traj[1], color=color, linestyle=linestyle))

                    # Add arrows to the lines at specified intervals
                    with plt.rc_context(rcParams):
                        _add_arrows_to_line2D(
                            ax, out[-1][0], arrow_pos, arrowstyle=arrow_style,
                            dir=1)
    return out


#
# User accessible utility functions
#

# Utility function to generate boxgrid (in the form needed here)
def boxgrid(xvals, yvals):
    """Generate list of points along the edge of box.

    points = boxgrid(xvals, yvals) generates a list of points that
    corresponds to a grid given by the cross product of the x and y values.

    Parameters
    ----------
    xvals, yvals: 1D array-like
        Array of points defining the points on the lower and left edges of
        the box.

    Returns
    -------
    grid: 2D array
        Array with shape (p, 2) defining the points along the edges of the
        box, where p is the number of points around the edge.

    """
    return np.array(
        [(x, yvals[0]) for x in xvals[:-1]] +           # lower edge
        [(xvals[-1], y) for y in yvals[:-1]] +          # right edge
        [(x, yvals[-1]) for x in xvals[:0:-1]] +        # upper edge
        [(xvals[0], y) for y in yvals[:0:-1]]           # left edge
    )


# Utility function to generate meshgrid (in the form needed here)
# TODO: add examples of using grid functions directly
def meshgrid(xvals, yvals):
    """Generate list of points forming a mesh.

    points = meshgrid(xvals, yvals) generates a list of points that
    corresponds to a grid given by the cross product of the x and y values.

    Parameters
    ----------
    xvals, yvals: 1D array-like
        Array of points defining the points on the lower and left edges of
        the box.

    Returns
    -------
    grid: 2D array
        Array of points with shape (n * m, 2) defining the mesh

    """
    xvals, yvals = np.meshgrid(xvals, yvals)
    grid = np.zeros((xvals.shape[0] * xvals.shape[1], 2))
    grid[:, 0] = xvals.reshape(-1)
    grid[:, 1] = yvals.reshape(-1)

    return grid


# Utility function to generate circular grid
def circlegrid(centers, radius, num):
    """Generate list of points around a circle.

    points = circlegrid(centers, radius, num) generates a list of points
    that form a circle around a list of centers.

    Parameters
    ----------
    centers : 2D array-like
        Array of points with shape (p, 2) defining centers of the circles.
    radius : float
        Radius of the points to be generated around each center.
    num : int
        Number of points to generate around the circle.

    Returns
    -------
    grid: 2D array
        Array of points with shape (p * num, 2) defining the circles.

    """
    centers = np.atleast_2d(np.array(centers))
    grid = np.zeros((centers.shape[0] * num, 2))
    for i, center in enumerate(centers):
        grid[i * num: (i + 1) * num, :] = center + np.array([
            [radius * math.cos(theta), radius * math.sin(theta)] for
            theta in np.linspace(0, 2 * math.pi, num, endpoint=False)])
    return grid

#
# Internal utility functions
#

# Create a system from a callable
def _create_system(sys, params):
    if isinstance(sys, NonlinearIOSystem):
        if sys.nstates != 2:
            raise ValueError("system must be planar")
        return sys

    # Make sure that if params is present, it has 'args' key
    if params and not params.get('args', None):
        raise ValueError("params must be dict with key 'args'")

    _update = lambda t, x, u, params: sys(t, x, *params.get('args', ()))
    _output = lambda t, x, u, params: np.array([])
    return NonlinearIOSystem(
        _update, _output, states=2, inputs=0, outputs=0, name="_callable")

# Set axis limits for the plot
def _set_axis_limits(ax, pointdata):
    # Get the current axis limits
    if ax.lines:
        xlim, ylim = ax.get_xlim(), ax.get_ylim()
    else:
        # Nothing on the plot => always use new limits
        xlim, ylim = [np.inf, -np.inf], [np.inf, -np.inf]

    # Short utility function for updating axis limits
    def _update_limits(cur, new):
        return [min(cur[0], np.min(new)), max(cur[1], np.max(new))]

    # If we were passed a box, use that to update the limits
    if isinstance(pointdata, list) and len(pointdata) == 4:
        xlim = _update_limits(xlim, [pointdata[0], pointdata[1]])
        ylim = _update_limits(ylim, [pointdata[2], pointdata[3]])

    elif isinstance(pointdata, np.ndarray):
        pointdata = np.atleast_2d(pointdata)
        xlim = _update_limits(
            xlim, [np.min(pointdata[:, 0]), np.max(pointdata[:, 0])])
        ylim = _update_limits(
            ylim, [np.min(pointdata[:, 1]), np.max(pointdata[:, 1])])

    # Keep track of the largest dimension on the plot
    maxlim = max(xlim[1] - xlim[0], ylim[1] - ylim[0])

    # Set the new limits
    ax.autoscale(enable=True, axis='x', tight=True)
    ax.autoscale(enable=True, axis='y', tight=True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    return xlim, ylim, maxlim


# Find equilibrium points
def _find_equilpts(sys, points, params=None):
    equilpts = []
    for i, x0 in enumerate(points):
        # Look for an equilibrium point near this point
        xeq, ueq = find_eqpt(sys, x0, 0, params=params)

        if xeq is None:
            continue            # didn't find anything

        # See if we have already found this point
        seen = False
        for x in equilpts:
            if np.allclose(np.array(x), xeq):
                seen = True
        if seen:
            continue

        # Save a new point
        equilpts += [xeq.tolist()]

    return equilpts


def _make_points(pointdata, gridspec, gridtype):
    # Check to see what type of data we got
    if isinstance(pointdata, np.ndarray) and gridtype is None:
        pointdata = np.atleast_2d(pointdata)
        if pointdata.shape[1] == 2:
            # Given a list of points => no action required
            return pointdata, None

    # Utility function to parse (and check) input arguments
    def _parse_args(defsize):
        if gridspec is None:
            return defsize

        elif not isinstance(gridspec, (list, tuple)) or \
             len(gridspec) != len(defsize):
            raise ValueError("invalid grid specification")

        return gridspec

    # Generate points based on grid type
    match gridtype:
        case 'boxgrid' | None:
            gridspec = _parse_args([6, 4])
            points = boxgrid(
                np.linspace(pointdata[0], pointdata[1], gridspec[0]),
                np.linspace(pointdata[2], pointdata[3], gridspec[1]))

        case 'meshgrid':
            gridspec = _parse_args([9, 6])
            points = meshgrid(
                np.linspace(pointdata[0], pointdata[1], gridspec[0]),
                np.linspace(pointdata[2], pointdata[3], gridspec[1]))

        case 'circlegrid':
            gridspec = _parse_args((0.5, 10))
            if isinstance(pointdata, np.ndarray):
                # Create circles around each point
                points = circlegrid(pointdata, gridspec[0], gridspec[1])
            else:
                # Create circle around center of the plot
                points = circlegrid(
                    np.array(
                        [(pointdata[0] + pointdata[1]) / 2,
                         (pointdata[0] + pointdata[1]) / 2]),
                    gridspec[0], gridspec[1])

        case _:
            raise ValueError(f"unknown grid type '{gridtype}'")

    return points, gridspec


def _parse_arrow_keywords(kwargs):
    # Get values for params (and pop from list to allow keyword use in plot)
    # TODO: turn this into a utility function (shared with nyquist_plot?)
    arrows = config._get_param(
        'phaseplot', 'arrows', kwargs, None, pop=True)
    arrow_size = config._get_param(
        'phaseplot', 'arrow_size', kwargs, None, pop=True)
    arrow_style = config._get_param('phaseplot', 'arrow_style', kwargs, None)

    # Parse the arrows keyword
    if not arrows:
        arrow_pos = []
    elif isinstance(arrows, int):
        N = arrows
        # Space arrows out, starting midway along each "region"
        arrow_pos = np.linspace(0.5/N, 1 + 0.5/N, N, endpoint=False)
    elif isinstance(arrows, (list, np.ndarray)):
        arrow_pos = np.sort(np.atleast_1d(arrows))
    else:
        raise ValueError("unknown or unsupported arrow location")

    # Set the arrow style
    if arrow_style is None:
        arrow_style = mpl.patches.ArrowStyle(
            'simple', head_width=int(2 * arrow_size / 3),
            head_length=arrow_size)

    return arrow_pos, arrow_style


# TODO: move to ctrlplot?
def _create_trajectory(
        sys, revsys, timepts, X0, params, dir, suppress_warnings=False,
        gridtype=None, gridspec=None, xlim=None, ylim=None):
    # Comput ethe forward trajectory
    if dir == 'forward' or dir == 'both':
        fwdresp = input_output_response(
            sys, timepts, X0=X0, params=params, ignore_errors=True)
        if not fwdresp.success and not suppress_warnings:
            warnings.warn(f"{X0=}, {fwdresp.message}")

    # Compute the reverse trajectory
    if dir == 'reverse' or dir == 'both':
        revresp = input_output_response(
            revsys, timepts, X0=X0, params=params, ignore_errors=True)
        if not revresp.success and not suppress_warnings:
            warnings.warn(f"{X0=}, {revresp.message}")

    # Create the trace to plot
    if dir == 'forward':
        traj = fwdresp.states
    elif dir == 'reverse':
        traj = revresp.states[:, ::-1]
    elif dir == 'both':
        traj = np.hstack([revresp.states[:, :1:-1], fwdresp.states])

    # Remove points outside the window (keep first point beyond boundary)
    inrange = np.asarray(
        (traj[0] >= xlim[0]) & (traj[0] <= xlim[1]) &
        (traj[1] >= ylim[0]) & (traj[1] <= ylim[1]))
    inrange[:-1] = inrange[:-1] | inrange[1:]   # keep if next point in range
    inrange[1:] = inrange[1:] | inrange[:-1]    # keep if prev point in range

    return traj[:, inrange]


def _make_timepts(timepts, i):
    if timepts is None:
        return np.linspace(0, 1)
    elif isinstance(timepts, (int, float)):
        return np.linspace(0, timepts)
    elif timepts.ndim == 2:
        return timepts[i]
    return timepts


#
# Legacy phase plot function
#
# Author: Richard Murray
# Date: 24 July 2011, converted from MATLAB version (2002); based on
# a version by Kristi Morgansen
#
def phase_plot(odefun, X=None, Y=None, scale=1, X0=None, T=None,
               lingrid=None, lintime=None, logtime=None, timepts=None,
               parms=None, params=(), tfirst=False, verbose=True):

    """(legacy) Phase plot for 2D dynamical systems.

    Produces a vector field or stream line plot for a planar system.  This
    function has been replaced by the :func:`~control.phase_plane_map` and
    :func:`~control.phase_plane_plot` functions.

    Call signatures:
      phase_plot(func, X, Y, ...) - display vector field on meshgrid
      phase_plot(func, X, Y, scale, ...) - scale arrows
      phase_plot(func. X0=(...), T=Tmax, ...) - display stream lines
      phase_plot(func, X, Y, X0=[...], T=Tmax, ...) - plot both
      phase_plot(func, X0=[...], T=Tmax, lingrid=N, ...) - plot both
      phase_plot(func, X0=[...], lintime=N, ...) - stream lines with arrows

    Parameters
    ----------
    func : callable(x, t, ...)
        Computes the time derivative of y (compatible with odeint).  The
        function should be the same for as used for :mod:`scipy.integrate`.
        Namely, it should be a function of the form dxdt = F(t, x) that
        accepts a state x of dimension 2 and returns a derivative dx/dt of
        dimension 2.
    X, Y: 3-element sequences, optional, as [start, stop, npts]
        Two 3-element sequences specifying x and y coordinates of a
        grid.  These arguments are passed to linspace and meshgrid to
        generate the points at which the vector field is plotted.  If
        absent (or None), the vector field is not plotted.
    scale: float, optional
        Scale size of arrows; default = 1
    X0: ndarray of initial conditions, optional
        List of initial conditions from which streamlines are plotted.
        Each initial condition should be a pair of numbers.
    T: array-like or number, optional
        Length of time to run simulations that generate streamlines.
        If a single number, the same simulation time is used for all
        initial conditions.  Otherwise, should be a list of length
        len(X0) that gives the simulation time for each initial
        condition.  Default value = 50.
    lingrid : integer or 2-tuple of integers, optional
        Argument is either N or (N, M).  If X0 is given and X, Y are missing,
        a grid of arrows is produced using the limits of the initial
        conditions, with N grid points in each dimension or N grid points in x
        and M grid points in y.
    lintime : integer or tuple (integer, float), optional
        If a single integer N is given, draw N arrows using equally space time
        points.  If a tuple (N, lambda) is given, draw N arrows using
        exponential time constant lambda
    timepts : array-like, optional
        Draw arrows at the given list times [t1, t2, ...]
    tfirst : bool, optional
        If True, call `func` with signature `func(t, x, ...)`.
    params: tuple, optional
        List of parameters to pass to vector field: `func(x, t, *params)`

    See also
    --------
    box_grid : construct box-shaped grid of initial conditions

    """
    # Generate a deprecation warning
    warnings.warn(
        "phase_plot is deprecated; use phase_plot_plot instead",
        FutureWarning)

    #
    # Figure out ranges for phase plot (argument processing)
    #
    #! TODO: need to add error checking to arguments
    #! TODO: think through proper action if multiple options are given
    #
    autoFlag = False
    logtimeFlag = False
    timeptsFlag = False
    Narrows = 0

    # Get parameters to pass to function
    if parms:
        warnings.warn(
            f"keyword 'parms' is deprecated; use 'params'", FutureWarning)
        if params:
            raise ControlArgument(f"duplicate keywords 'parms' and 'params'")
        else:
            params = parms

    if lingrid is not None:
        autoFlag = True
        Narrows = lingrid
        if (verbose):
            print('Using auto arrows\n')

    elif logtime is not None:
        logtimeFlag = True
        Narrows = logtime[0]
        timefactor = logtime[1]
        if (verbose):
            print('Using logtime arrows\n')

    elif timepts is not None:
        timeptsFlag = True
        Narrows = len(timepts)

    # Figure out the set of points for the quiver plot
    #! TODO: Add sanity checks
    elif X is not None and Y is not None:
        x1, x2 = np.meshgrid(
            np.linspace(X[0], X[1], X[2]),
            np.linspace(Y[0], Y[1], Y[2]))
        Narrows = len(x1)

    else:
        # If we weren't given any grid points, don't plot arrows
        Narrows = 0

    if not autoFlag and not logtimeFlag and not timeptsFlag and Narrows > 0:
        # Now calculate the vector field at those points
        (nr,nc) = x1.shape
        dx = np.empty((nr, nc, 2))
        for i in range(nr):
            for j in range(nc):
                if tfirst:
                    dx[i, j, :] = np.squeeze(
                        odefun(0, [x1[i,j], x2[i,j]], *params))
                else:
                    dx[i, j, :] = np.squeeze(
                        odefun([x1[i,j], x2[i,j]], 0, *params))

        # Plot the quiver plot
        #! TODO: figure out arguments to make arrows show up correctly
        if scale is None:
            plt.quiver(x1, x2, dx[:,:,1], dx[:,:,2], angles='xy')
        elif (scale != 0):
            #! TODO: optimize parameters for arrows
            #! TODO: figure out arguments to make arrows show up correctly
            xy = plt.quiver(x1, x2, dx[:,:,0]*np.abs(scale),
                            dx[:,:,1]*np.abs(scale), angles='xy')
            # set(xy, 'LineWidth', PP_arrow_linewidth, 'Color', 'b')

        #! TODO: Tweak the shape of the plot
        # a=gca; set(a,'DataAspectRatio',[1,1,1])
        # set(a,'XLim',X(1:2)); set(a,'YLim',Y(1:2))
        plt.xlabel('x1'); plt.ylabel('x2')

    # See if we should also generate the streamlines
    if X0 is None or len(X0) == 0:
        return

    # Convert initial conditions to a numpy array
    X0 = np.array(X0)
    (nr, nc) = np.shape(X0)

    # Generate some empty matrices to keep arrow information
    x1 = np.empty((nr, Narrows))
    x2 = np.empty((nr, Narrows))
    dx = np.empty((nr, Narrows, 2))

    # See if we were passed a simulation time
    if T is None:
        T = 50

    # Parse the time we were passed
    TSPAN = T
    if isinstance(T, (int, float)):
        TSPAN = np.linspace(0, T, 100)

    # Figure out the limits for the plot
    if scale is None:
        # Assume that the current axis are set as we want them
        alim = plt.axis()
        xmin = alim[0]; xmax = alim[1]
        ymin = alim[2]; ymax = alim[3]
    else:
        # Use the maximum extent of all trajectories
        xmin = np.min(X0[:,0]); xmax = np.max(X0[:,0])
        ymin = np.min(X0[:,1]); ymax = np.max(X0[:,1])

    # Generate the streamlines for each initial condition
    for i in range(nr):
        state = odeint(odefun, X0[i], TSPAN, args=params, tfirst=tfirst)
        time = TSPAN

        plt.plot(state[:,0], state[:,1])
        #! TODO: add back in colors for stream lines
        # PP_stream_color(np.mod(i-1, len(PP_stream_color))+1))
        # set(h[i], 'LineWidth', PP_stream_linewidth)

        # Plot arrows if quiver parameters were 'auto'
        if autoFlag or logtimeFlag or timeptsFlag:
            # Compute the locations of the arrows
            #! TODO: check this logic to make sure it works in python
            for j in range(Narrows):

                # Figure out starting index; headless arrows start at 0
                k = -1 if scale is None else 0

                # Figure out what time index to use for the next point
                if autoFlag:
                    # Use a linear scaling based on ODE time vector
                    tind = np.floor((len(time)/Narrows) * (j-k)) + k
                elif logtimeFlag:
                    # Use an exponential time vector
                    # MATLAB: tind = find(time < (j-k) / lambda, 1, 'last')
                    tarr = _find(time < (j-k) / timefactor)
                    tind = tarr[-1] if len(tarr) else 0
                elif timeptsFlag:
                    # Use specified time points
                    # MATLAB: tind = find(time < Y[j], 1, 'last')
                    tarr = _find(time < timepts[j])
                    tind = tarr[-1] if len(tarr) else 0

                # For tailless arrows, skip the first point
                if tind == 0 and scale is None:
                    continue

                # Figure out the arrow at this point on the curve
                x1[i,j] = state[tind, 0]
                x2[i,j] = state[tind, 1]

                # Skip arrows outside of initial condition box
                if (scale is not None or
                     (x1[i,j] <= xmax and x1[i,j] >= xmin and
                      x2[i,j] <= ymax and x2[i,j] >= ymin)):
                    if tfirst:
                        pass
                        v = odefun(0, [x1[i,j], x2[i,j]], *params)
                    else:
                        v = odefun([x1[i,j], x2[i,j]], 0, *params)
                    dx[i, j, 0] = v[0]; dx[i, j, 1] = v[1]
                else:
                    dx[i, j, 0] = 0; dx[i, j, 1] = 0

    # Set the plot shape before plotting arrows to avoid warping
    # a=gca
    # if (scale != None):
    #     set(a,'DataAspectRatio', [1,1,1])
    # if (xmin != xmax and ymin != ymax):
    #     plt.axis([xmin, xmax, ymin, ymax])
    # set(a, 'Box', 'on')

    # Plot arrows on the streamlines
    if scale is None and Narrows > 0:
        # Use a tailless arrow
        #! TODO: figure out arguments to make arrows show up correctly
        plt.quiver(x1, x2, dx[:,:,0], dx[:,:,1], angles='xy')
    elif scale != 0 and Narrows > 0:
        #! TODO: figure out arguments to make arrows show up correctly
        xy = plt.quiver(x1, x2, dx[:,:,0]*abs(scale), dx[:,:,1]*abs(scale),
                        angles='xy')
        # set(xy, 'LineWidth', PP_arrow_linewidth)
        # set(xy, 'AutoScale', 'off')
        # set(xy, 'AutoScaleFactor', 0)

    if scale < 0:
        bp = plt.plot(x1, x2, 'b.');        # add dots at base
        # set(bp, 'MarkerSize', PP_arrow_markersize)


# Utility function for generating initial conditions around a box
def box_grid(xlimp, ylimp):
    """box_grid   generate list of points on edge of box

    list = box_grid([xmin xmax xnum], [ymin ymax ynum]) generates a
    list of points that correspond to a uniform grid at the end of the
    box defined by the corners [xmin ymin] and [xmax ymax].
    """

    # Generate a deprecation warning
    warnings.warn(
        "box_grid is deprecated; use phaseplot.boxgrid instead",
        FutureWarning)

    return boxgrid(
        np.linspace(xlimp[0], xlimp[1], xlimp[2]),
        np.linspace(ylimp[0], ylimp[1], ylimp[2]))


# TODO: rename to something more useful (or remove??)
def _find(condition):
    """Returns indices where ravel(a) is true.
    Private implementation of deprecated matplotlib.mlab.find
    """
    return np.nonzero(np.ravel(condition))[0]

python-control / python-control / 10312443515

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous