unify color processesing

murrayrm · murrayrm · commit a9ffba5be7f3 · 2024-07-25T23:00:01.000-07:00
diff --git a/control/ctrlplot.py b/control/ctrlplot.py
@@ -5,7 +5,7 @@
 
 # Code pattern for control system plotting functions:
 #
-# def name_plot(sysdata, plot=None, **kwargs):
+# def name_plot(sysdata, *fmt, plot=None, **kwargs):
 #     # Process keywords and set defaults
 #     ax = kwargs.pop('ax', None)
 #     color = kwargs.pop('color', None)
@@ -37,12 +37,13 @@
 #     # Plot the data
 #     lines = np.full(ax_array.shape, [])
 #     line_labels = _process_line_labels(label, ntraces, nrows, ncols)
+#     color_offset, color_cycle = _get_color_offset(ax)
 #     for i, j in itertools.product(range(nrows), range(ncols)):
 #         ax = ax_array[i, j]
-#         color_cycle, color_offset = _process_color_keyword(ax)
 #         for k in range(ntraces):
 #             if color is None:
-#                 color = color_cycle[(k + color_offset) % len(color_cycle)]
+#                 color = _get_color(
+#                     color, fmt=fmt, offset=k, color_cycle=color_cycle)
 #             label = line_labels[k, i, j]
 #             lines[i, j] += ax.plot(data.x, data.y, color=color, label=label)
 #
@@ -656,11 +657,78 @@ def _add_arrows_to_line2D(
     return arrows
 
 
-def _get_color(colorspec, ax=None, lines=None, color_cycle=None):
+def _get_color_offset(ax, color_cycle=None):
+    """Get color offset based on current lines.
+
+    This function determines that the current offset is for the next color
+    to use based on current colors in a plot.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes
+        Axes containing already plotted lines.
+    color_cycle : list of matplotlib color specs, optional
+        Colors to use in plotting lines.  Defaults to matplotlib rcParams
+        color cycle.
+
+    Returns
+    -------
+    color_offset : matplotlib color spec
+        Starting color for next line to be drawn.
+    color_cycle : list of matplotlib color specs
+        Color cycle used to determine colors.
+
+    """
+    if color_cycle is None:
+        color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']
+
+    color_offset = 0
+    if len(ax.lines) > 0:
+        last_color = ax.lines[-1].get_color()
+        if last_color in color_cycle:
+            color_offset = color_cycle.index(last_color) + 1
+
+    return color_offset % len(color_cycle), color_cycle
+
+
+def _get_color(
+        colorspec, offset=None, fmt=None, ax=None, lines=None,
+        color_cycle=None):
+    """Get color to use for plotting line.
+
+    This function returns the color to be used for the line to be drawn (or
+    None if the detault color cycle for the axes should be used).
+
+    Parameters
+    ----------
+    colorspec : matplotlib color specification
+        User-specified color (or None).
+    offset : int, optional
+        Offset into the color cycle (for multi-trace plots).
+    fmt : str, optional
+        Format string passed to plotting command.
+    ax : matplotlib.axes.Axes, optional
+        Axes containing already plotted lines.
+    lines : list of matplotlib.lines.Line2D, optional
+        List of plotted lines.  If not given, use ax.get_lines().
+    color_cycle : list of matplotlib color specs, optional
+        Colors to use in plotting lines.  Defaults to matplotlib rcParams
+        color cycle.
+
+    Returns
+    -------
+    color : matplotlib color spec
+        Color to use for this line (or None for matplotlib default).
+
+    """
     # See if the color was explicitly specified by the user
     if isinstance(colorspec, dict):
         if 'color' in colorspec:
             return colorspec.pop('color')
+    elif fmt is not None and \
+         [isinstance(arg, str) and
+          any([c in arg for c in "bgrcmykw#"]) for arg in fmt]:
+        return None             # *fmt will set the color
     elif colorspec != None:
         return colorspec
 
@@ -673,7 +741,9 @@ def _get_color(colorspec, ax=None, lines=None, color_cycle=None):
         lines = ax.lines
 
     # If we were passed a set of lines, try to increment color from previous
-    if lines is not None:
+    if offset is not None:
+        return color_cycle[offset]
+    elif lines is not None:
         color_offset = 0
         if len(ax.lines) > 0:
             last_color = ax.lines[-1].get_color()
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -20,8 +20,9 @@
 from . import config
 from .bdalg import feedback
 from .ctrlplot import ControlPlot, _add_arrows_to_line2D, _find_axes_center, \
-    _get_line_labels, _make_legend_labels, _process_ax_keyword, \
-    _process_legend_keywords, _process_line_labels, _update_plot_title
+    _get_color, _get_color_offset, _get_line_labels, _make_legend_labels, \
+    _process_ax_keyword, _process_legend_keywords, _process_line_labels, \
+    _update_plot_title
 from .ctrlutil import unwrap
 from .exception import ControlMIMONotImplemented
 from .frdata import FrequencyResponseData
@@ -2306,6 +2307,7 @@ def singular_values_plot(
 
     """
     # Keyword processing
+    color = kwargs.pop('color', None)
     dB = config._get_param(
         'freqplot', 'dB', kwargs, _freqplot_defaults, pop=True)
     Hz = config._get_param(
@@ -2370,14 +2372,8 @@ def singular_values_plot(
     legend_loc, _, show_legend = _process_legend_keywords(
         kwargs, None, 'center right')
 
-    # Handle color cycle manually as all singular values
-    # of the same systems are expected to be of the same color
-    color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']
-    color_offset = 0
-    if len(ax_sigma.lines) > 0:
-        last_color = ax_sigma.lines[-1].get_color()
-        if last_color in color_cycle:
-            color_offset = color_cycle.index(last_color) + 1
+    # Get color offset for first (new) line to be drawn
+    color_offset, color_cycle = _get_color_offset(ax_sigma)
 
     # Create a list of lines for the output
     out = np.empty(len(data), dtype=object)
@@ -2392,14 +2388,13 @@ def singular_values_plot(
         else:
             nyq_freq = None
 
-        # See if the color was specified, otherwise rotate
-        if kwargs.get('color', None) or any(
-                [isinstance(arg, str) and
-                 any([c in arg for c in "bgrcmykw#"]) for arg in fmt]):
-            color_arg = {}                      # color set by *fmt, **kwargs
-        else:
-            color_arg = {'color': color_cycle[
-                (idx_sys + color_offset) % len(color_cycle)]}
+        # Determine the color to use for this response
+        color = _get_color(
+            color, fmt=fmt, offset=color_offset + idx_sys,
+            color_cycle=color_cycle)
+
+        # To avoid conflict with *fmt, only pass color kw if non-None
+        color_arg = {} if color is None else {'color': color}
 
         # Decide on the system name
         sysname = response.sysname if response.sysname is not None \
diff --git a/control/phaseplot.py b/control/phaseplot.py
@@ -548,8 +548,12 @@ def separatrices(
         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.
+    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.
 
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -19,8 +19,9 @@
 
 from . import config
 from .config import _process_legacy_keyword
-from .ctrlplot import ControlPlot, _get_line_labels, _process_ax_keyword, \
-    _process_legend_keywords, _process_line_labels, _update_plot_title
+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 .freqplot import _freqplot_defaults
 from .grid import nogrid, sgrid, zgrid
 from .iosys import isctime, isdtime
@@ -367,15 +368,8 @@ def pole_zero_plot(
         if grid is not None:
             warnings.warn("axis already exists; grid keyword ignored")
 
-    # Handle color cycle manually as all root locus segments
-    # of the same system are expected to be of the same color
-    # TODO: replace with common function?
-    color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']
-    color_offset = 0
-    if len(ax.lines) > 0:
-        last_color = ax.lines[-1].get_color()
-        if last_color in color_cycle:
-            color_offset = color_cycle.index(last_color) + 1
+    # 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(
@@ -388,11 +382,8 @@ def pole_zero_plot(
         poles = response.poles
         zeros = response.zeros
 
-        # Get the color to use for this system
-        if user_color is None:
-            color = color_cycle[(color_offset + idx) % len(color_cycle)]
-        else:
-            color = user_color
+        # 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:
diff --git a/control/tests/ctrlplot_test.py b/control/tests/ctrlplot_test.py
@@ -331,7 +331,7 @@ def test_plot_linestyle_processing(resp_fcn, plot_fcn):
     sys2 = ct.rss(4, 1, 1, strictly_proper=True, name="sys[2]")
 
     # Set up arguments
-    args, _, _, kwargs, meth_kwargs, plot_kwargs, resp_kwargs = \
+    args1, args2, _, kwargs, meth_kwargs, plot_kwargs, resp_kwargs = \
         setup_plot_arguments(resp_fcn, plot_fcn)
     default_labels = ["sys[1]", "sys[2]"]
     expected_labels = ["sys1_", "sys2_"]
@@ -340,8 +340,12 @@ def test_plot_linestyle_processing(resp_fcn, plot_fcn):
             default_labels = ["P=sys[1]", "P=sys[2]"]
 
     # Set line color
-    cplt = plot_fcn(*args, **kwargs, **plot_kwargs, color='r')
-    assert cplt.lines.reshape(-1)[0][0].get_color() == 'r'
+    cplt1 = plot_fcn(*args1, **kwargs, **plot_kwargs, color='r')
+    assert cplt1.lines.reshape(-1)[0][0].get_color() == 'r'
+
+    # Second plot, new line color
+    cplt2 = plot_fcn(*args2, **kwargs, **plot_kwargs, color='g')
+    assert cplt2.lines.reshape(-1)[0][0].get_color() == 'g'
 
     # Make sure that docstring documents line properties
     if plot_fcn not in legacy_plot_fcns:
@@ -350,12 +354,12 @@ def test_plot_linestyle_processing(resp_fcn, plot_fcn):
 
     # Set other characteristics if documentation says we can
     if "line properties" in plot_fcn.__doc__:
-        cplt = plot_fcn(*args, **kwargs, **plot_kwargs, linewidth=5)
+        cplt = plot_fcn(*args1, **kwargs, **plot_kwargs, linewidth=5)
         assert cplt.lines.reshape(-1)[0][0].get_linewidth() == 5
 
     # If fmt string is allowed, use it to set line color and style
     if "*fmt" in plot_fcn.__doc__:
-        cplt = plot_fcn(*args, 'r--', **kwargs, **plot_kwargs)
+        cplt = plot_fcn(*args1, 'r--', **kwargs, **plot_kwargs)
         assert cplt.lines.reshape(-1)[0][0].get_color() == 'r'
         assert cplt.lines.reshape(-1)[0][0].get_linestyle() == '--'
 
diff --git a/doc/plotting.rst b/doc/plotting.rst
@@ -449,15 +449,15 @@ various ways.  The following general rules apply:
   2x2 array of subplots should be given a 2x2 array of axes for the ``ax``
   keyword).
 
-* The ```color``, ``linestyle``, ``linewidth``, and other matplotlib line
+* The ``color``, ``linestyle``, ``linewidth``, and other matplotlib line
   property arguments can be used to override the default line properties.
   If these arguments are absent, the default matplotlib line properties are
   used and the color cycles through the default matplotlib color cycle.
 
-  The :func:`~control.bode_plot`, :func:`~control.time_response_plot`, and
-  selected other commands can also accept a matplotlib format string (e.g.,
-  'r--').  The format string that must appear as a positional argument
-  right after the required data argumnt.
+  The :func:`~control.bode_plot`, :func:`~control.time_response_plot`,
+  and selected other commands can also accept a matplotlib format
+  string (e.g., 'r--').  The format string must appear as a positional
+  argument right after the required data argumnt.
 
   Note that line property arguments are the same for all lines generated as
   part of a single plotting command call, including when multiple responses
