refactoring to put common functions in ctrlplot.py

murrayrm · murrayrm · commit 44123c9d969c · 2024-06-29T12:41:18.000-07:00
diff --git a/control/ctrlplot.py b/control/ctrlplot.py
@@ -3,12 +3,14 @@
 #
 # Collection of functions that are used by various plotting functions.
 
+from os.path import commonprefix
+
 import matplotlib.pyplot as plt
 import numpy as np
 
 from . import config
 
-__all__ = ['suptitle']
+__all__ = ['suptitle', 'get_plot_axes']
 
 
 def suptitle(
@@ -56,6 +58,94 @@ def suptitle(
         raise ValueError(f"unknown frame '{frame}'")
 
 
+# Create vectorized function to find axes from lines
+def get_plot_axes(line_array):
+    """Get a list of axes from an array of lines.
+
+    This function can be used to return the set of axes corresponding to
+    the line array that is returned by `time_response_plot`.  This is useful for
+    generating an axes array that can be passed to subsequent plotting
+    calls.
+
+    Parameters
+    ----------
+    line_array : array of list of Line2D
+        A 2D array with elements corresponding to a list of lines appearing
+        in an axes, matching the return type of a time response data plot.
+
+    Returns
+    -------
+    axes_array : array of list of Axes
+        A 2D array with elements corresponding to the Axes assocated with
+        the lines in `line_array`.
+
+    Notes
+    -----
+    Only the first element of each array entry is used to determine the axes.
+
+    """
+    _get_axes = np.vectorize(lambda lines: lines[0].axes)
+    return _get_axes(line_array)
+
+#
+# Utility functions
+#
+
+
+# Utility function to make legend labels
+def _make_legend_labels(labels, ignore_common=False):
+
+    # Look for a common prefix (up to a space)
+    common_prefix = commonprefix(labels)
+    last_space = common_prefix.rfind(', ')
+    if last_space < 0 or ignore_common:
+        common_prefix = ''
+    elif last_space > 0:
+        common_prefix = common_prefix[:last_space]
+    prefix_len = len(common_prefix)
+
+    # Look for a common suffice (up to a space)
+    common_suffix = commonprefix(
+        [label[::-1] for label in labels])[::-1]
+    suffix_len = len(common_suffix)
+    # Only chop things off after a comma or space
+    while suffix_len > 0 and common_suffix[-suffix_len] != ',':
+        suffix_len -= 1
+
+    # Strip the labels of common information
+    if suffix_len > 0 and not ignore_common:
+        labels = [label[prefix_len:-suffix_len] for label in labels]
+    else:
+        labels = [label[prefix_len:] for label in labels]
+
+    return labels
+
+
+def _update_suptitle(fig, title, rcParams=None, frame='axes'):
+    if fig is not None and isinstance(title, str):
+        # Get the current title, if it exists
+        old_title = None if fig._suptitle is None else fig._suptitle._text
+        new_title = title
+
+        if old_title is not None:
+            # Find the common part of the titles
+            common_prefix = commonprefix([old_title, new_title])
+
+            # Back up to the last space
+            last_space = common_prefix.rfind(' ')
+            if last_space > 0:
+                common_prefix = common_prefix[:last_space]
+            common_len = len(common_prefix)
+
+            # Add the new part of the title (usually the system name)
+            if old_title[common_len:] != new_title[common_len:]:
+                separator = ',' if len(common_prefix) > 0 else ';'
+                new_title = old_title + separator + new_title[common_len:]
+
+        # Add the title
+        suptitle(title, fig=fig, rcParams=rcParams, frame=frame)
+
+
 def _find_axes_center(fig, axs):
     """Find the midpoint between axes in display coordinates.
 
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -19,14 +19,14 @@
 
 from . import config
 from .bdalg import feedback
-from .ctrlplot import suptitle, _find_axes_center
+from .ctrlplot import suptitle, _find_axes_center, _make_legend_labels, \
+    _update_suptitle
 from .ctrlutil import unwrap
 from .exception import ControlMIMONotImplemented
 from .frdata import FrequencyResponseData
 from .lti import LTI, _process_frequency_response, frequency_response
 from .margins import stability_margins
 from .statesp import StateSpace
-from .timeplot import _make_legend_labels
 from .xferfcn import TransferFunction
 
 __all__ = ['bode_plot', 'NyquistResponseData', 'nyquist_response',
@@ -954,28 +954,7 @@ def gen_zero_centered_series(val_min, val_max, period):
         else:
             title = data[0].title
 
-    if fig is not None and isinstance(title, str):
-        # Get the current title, if it exists
-        old_title = None if fig._suptitle is None else fig._suptitle._text
-        new_title = title
-
-        if old_title is not None:
-            # Find the common part of the titles
-            common_prefix = commonprefix([old_title, new_title])
-
-            # Back up to the last space
-            last_space = common_prefix.rfind(' ')
-            if last_space > 0:
-                common_prefix = common_prefix[:last_space]
-            common_len = len(common_prefix)
-
-            # Add the new part of the title (usually the system name)
-            if old_title[common_len:] != new_title[common_len:]:
-                separator = ',' if len(common_prefix) > 0 else ';'
-                new_title = old_title + separator + new_title[common_len:]
-
-        # Add the title
-        suptitle(title, fig=fig, rcParams=rcParams, frame=suptitle_frame)
+    _update_suptitle(fig, title, rcParams=rcParams, frame=suptitle_frame)
 
     #
     # Create legends
diff --git a/control/tests/timeplot_test.py b/control/tests/timeplot_test.py
@@ -191,7 +191,7 @@ def test_response_plots(
 
 
 def test_axes_setup():
-    get_plot_axes = ct.timeplot.get_plot_axes
+    get_plot_axes = ct.get_plot_axes
 
     sys_2x3 = ct.rss(4, 2, 3)
     sys_2x3b = ct.rss(4, 2, 3)
@@ -377,7 +377,7 @@ def test_rcParams():
     assert ax.title.get_fontsize() == 10
     assert ax.xaxis._get_tick_label_size('x') == 10
     assert ax.yaxis._get_tick_label_size('y') == 10
-    assert fig._suptitle.get_fontsize() == 12
+    assert fig._suptitle.get_fontsize() == 10
 
 def test_relabel():
     sys1 = ct.rss(2, inputs='u', outputs='y')
diff --git a/control/timeplot.py b/control/timeplot.py
@@ -8,15 +8,16 @@
 # Note: It might eventually make sense to put the functions here
 # directly into timeresp.py.
 
-import numpy as np
+from warnings import warn
+
 import matplotlib as mpl
 import matplotlib.pyplot as plt
-from os.path import commonprefix
-from warnings import warn
+import numpy as np
 
 from . import config
+from .ctrlplot import _make_legend_labels, _update_suptitle
 
-__all__ = ['time_response_plot', 'combine_time_responses', 'get_plot_axes']
+__all__ = ['time_response_plot', 'combine_time_responses']
 
 # Default font dictionary
 _timeplot_rcParams = mpl.rcParams.copy()
@@ -157,7 +158,6 @@ def time_response_plot(
     #
     # Process keywords and set defaults
     #
-
     # Set up defaults
     time_label = config._get_param(
         'timeplot', 'time_label', kwargs, _timeplot_defaults, pop=True)
@@ -597,29 +597,7 @@ def _make_line_label(signal_index, signal_labels, trace_index):
     # list of systems (e.g., "Step response for sys[1], sys[2]").
     #
 
-    if fig is not None and title is not None:
-        # Get the current title, if it exists
-        old_title = None if fig._suptitle is None else fig._suptitle._text
-        new_title = title
-
-        if old_title is not None:
-            # Find the common part of the titles
-            common_prefix = commonprefix([old_title, new_title])
-
-            # Back up to the last space
-            last_space = common_prefix.rfind(' ')
-            if last_space > 0:
-                common_prefix = common_prefix[:last_space]
-            common_len = len(common_prefix)
-
-            # Add the new part of the title (usually the system name)
-            if old_title[common_len:] != new_title[common_len:]:
-                separator = ',' if len(common_prefix) > 0 else ';'
-                new_title = old_title + separator + new_title[common_len:]
-
-        # Add the title
-        with plt.rc_context(rcParams):
-            fig.suptitle(new_title)
+    _update_suptitle(fig, title, rcParams=rcParams)
 
     return out
 
@@ -730,62 +708,3 @@ def combine_time_responses(response_list, trace_labels=None, title=None):
         return_x=base.return_x, squeeze=base.squeeze, sysname=base.sysname,
         trace_labels=trace_labels, trace_types=trace_types,
         plot_inputs=base.plot_inputs)
-
-
-# Create vectorized function to find axes from lines
-def get_plot_axes(line_array):
-    """Get a list of axes from an array of lines.
-
-    This function can be used to return the set of axes corresponding to
-    the line array that is returned by `time_response_plot`.  This is useful for
-    generating an axes array that can be passed to subsequent plotting
-    calls.
-
-    Parameters
-    ----------
-    line_array : array of list of Line2D
-        A 2D array with elements corresponding to a list of lines appearing
-        in an axes, matching the return type of a time response data plot.
-
-    Returns
-    -------
-    axes_array : array of list of Axes
-        A 2D array with elements corresponding to the Axes assocated with
-        the lines in `line_array`.
-
-    Notes
-    -----
-    Only the first element of each array entry is used to determine the axes.
-
-    """
-    _get_axes = np.vectorize(lambda lines: lines[0].axes)
-    return _get_axes(line_array)
-
-
-# Utility function to make legend labels
-def _make_legend_labels(labels, ignore_common=False):
-
-    # Look for a common prefix (up to a space)
-    common_prefix = commonprefix(labels)
-    last_space = common_prefix.rfind(', ')
-    if last_space < 0 or ignore_common:
-        common_prefix = ''
-    elif last_space > 0:
-        common_prefix = common_prefix[:last_space]
-    prefix_len = len(common_prefix)
-
-    # Look for a common suffice (up to a space)
-    common_suffix = commonprefix(
-        [label[::-1] for label in labels])[::-1]
-    suffix_len = len(common_suffix)
-    # Only chop things off after a comma or space
-    while suffix_len > 0 and common_suffix[-suffix_len] != ',':
-        suffix_len -= 1
-
-    # Strip the labels of common information
-    if suffix_len > 0 and not ignore_common:
-        labels = [label[prefix_len:-suffix_len] for label in labels]
-    else:
-        labels = [label[prefix_len:] for label in labels]
-
-    return labels
