unify color processesing

murrayrm · murrayrm · commit 9a407ac04f07 · 2024-07-25T22:56:53.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
@@ -454,10 +454,10 @@ various ways.  The following general rules apply:
   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
