add plot_streamplot option to phase_plane_plot to use matplotlibs streamplot

lkies · lkies · commit c0e4cb42576d · 2025-02-01T20:58:16.000+01:00
diff --git a/control/phaseplot.py b/control/phaseplot.py
@@ -24,6 +24,7 @@
 * separatrices: Plot separatrices in the phase plane
 * streamlines: Plot stream lines in the phase plane
 * vectorfield: Plot a vector field in the phase plane
+* streampot: Plot streamlines using matplotlib's streamplot function
 
 """
 
@@ -55,7 +56,7 @@ 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_streamplot=False, **kwargs
 ):
     """Plot phase plane diagram.
 
@@ -133,6 +134,10 @@ def phase_plane_plot(
         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_streamplot : bool or dict, optional
+        If `True` then use matplotlib's streamplot function to plot the
+        streamlines.  If set to a dict, pass on the key-value pairs in the
+        dict as keywords to :func:`~
     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
@@ -213,6 +218,16 @@ def _create_kwargs(global_kwargs, local_kwargs, **other_kwargs):
         for kw in ['color', 'params']:
             initial_kwargs.pop(kw, None)
 
+    if plot_streamplot:
+        kwargs_local = _create_kwargs(
+            kwargs, plot_streamplot, gridspec=gridspec, ax=ax)
+        streamplot(
+            sys, pointdata, _check_kwargs=False, **kwargs_local)
+
+        # Get rid of keyword arguments handled by streamplot
+        for kw in ['color', 'params']:
+            initial_kwargs.pop(kw, None)
+
     if plot_equilpoints:
         kwargs_local = _create_kwargs(
             kwargs, plot_equilpoints, gridspec=gridspec, ax=ax)
@@ -331,6 +346,116 @@ def vectorfield(
     return out
 
 
+def streamplot(
+        sys, pointdata, gridspec=None, ax=None, vary_color=False,
+        vary_linewidth=False, cmap=None, norm=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 from which to make the streamplot. In the latter case,
+        the points lie on a grid like that generated by `meshgrid`.
+    gridspec : list, optional
+        Specifies the size of the grid in the x and y axes on which to
+        generate points.
+    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.
+    vary_color : bool, optional
+        If set to `True`, vary the color of the streamlines based on the magnitude
+    vary_linewidth : bool, optional
+        If set to `True`, vary the linewidth of the streamlines based on the magnitude
+    cmap : str or Colormap, optional
+        Colormap to use for varying the color of the streamlines
+    norm : Normalize, optional
+        An instance of Normalize to use for scaling the colormap and linewidths
+    ax : matplotlib.axes.Axes
+        Use the given axes for the plot, otherwise use the current axes.
+
+    Returns
+    -------
+    out : Quiver
+
+    Other parameters
+    ----------------
+    rcParams : dict
+        Override the default parameters used for generating plots.
+        Default is set by config.default['ctrlplot.rcParams'].
+    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 streamplot field
+    points, gridspec = _make_points(pointdata, gridspec, 'meshgrid')
+
+    # attempt to recover the grid by counting the jumps in xvals
+    if gridspec is None:
+        nrows = np.sum(np.diff(points[:, 0]) < 0) + 1
+        ncols = points.shape[0] // nrows
+        if nrows * ncols != points.shape[0]:
+            raise ValueError("Could not recover grid from points.")
+        gridspec = [nrows, ncols]
+
+    grid_arr_shape = gridspec[::-1]
+    xs, ys = points[:, 0].reshape(grid_arr_shape), points[:, 1].reshape(grid_arr_shape)
+
+    # 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
+    sys._update_params(params)
+    us_flat, vs_flat = np.transpose([sys._rhs(0, x, np.zeros(sys.ninputs)) for x in points])
+    us, vs = us_flat.reshape(grid_arr_shape), vs_flat.reshape(grid_arr_shape)
+
+    magnitudes = np.linalg.norm([us, vs], axis=0)
+    norm = norm or mpl.colors.Normalize()
+    normalized = norm(magnitudes)
+    cmap =  plt.get_cmap(cmap)
+
+    with plt.rc_context(rcParams):
+        default_lw = plt.rcParams['lines.linewidth']
+        min_lw, max_lw = 0.25*default_lw, 2*default_lw
+        linewidths = normalized * (max_lw - min_lw) + min_lw if vary_linewidth else None
+        color = magnitudes if vary_color else color
+
+        out = ax.streamplot(xs, ys, us, vs, color=color, linewidth=linewidths, cmap=cmap, norm=norm)
+
+    return out
+
 def streamlines(
         sys, pointdata, timedata=1, gridspec=None, gridtype=None, dir=None,
         ax=None, _check_kwargs=True, suppress_warnings=False, **kwargs):
diff --git a/control/tests/kwargs_test.py b/control/tests/kwargs_test.py
@@ -171,6 +171,7 @@ def test_unrecognized_kwargs(function, nsssys, ntfsys, moreargs, kwargs,
      (control.phase_plane_plot, 1, ([-1, 1, -1, 1], 1), {}),
      (control.phaseplot.streamlines, 1, ([-1, 1, -1, 1], 1), {}),
      (control.phaseplot.vectorfield, 1, ([-1, 1, -1, 1], ), {}),
+     (control.phaseplot.streamplot, 1, ([-1, 1, -1, 1], ), {}),
      (control.phaseplot.equilpoints, 1, ([-1, 1, -1, 1], ), {}),
      (control.phaseplot.separatrices, 1, ([-1, 1, -1, 1], ), {}),
      (control.singular_values_plot, 1, (), {})]
@@ -347,6 +348,7 @@ def test_response_plot_kwargs(data_fcn, plot_fcn, mimo):
         optimal_test.test_oep_argument_errors,
     'phaseplot.streamlines': test_matplotlib_kwargs,
     'phaseplot.vectorfield': test_matplotlib_kwargs,
+    'phaseplot.streamplot': test_matplotlib_kwargs,
     'phaseplot.equilpoints': test_matplotlib_kwargs,
     'phaseplot.separatrices': test_matplotlib_kwargs,
 }
diff --git a/control/tests/phaseplot_test.py b/control/tests/phaseplot_test.py
@@ -227,6 +227,18 @@ def oscillator_update(t, x, u, params):
     if savefigs:
         plt.savefig('phaseplot-oscillator-helpers.png')
 
+    plt.figure()
+    ct.phase_plane_plot(
+        invpend, [-2*pi, 2*pi, -2, 2], plot_streamlines=False,
+        plot_streamplot=dict(vary_color=True, vary_density=True),
+        gridspec=[60, 20], params={'m': 1, 'l': 1, 'b': 0.2, 'g': 1}
+    )
+    plt.xlabel(r"$\theta$ [rad]")
+    plt.ylabel(r"$\dot\theta$ [rad/sec]")
+
+    if savefigs:
+        plt.savefig('phaseplot-invpend-streamplot.png')
+
 
 if __name__ == "__main__":
     #
