add unit tests for common plotting functionality

murrayrm · murrayrm · commit fb5c194fb249 · 2024-08-08T14:54:49.000-07:00
diff --git a/control/phaseplot.py b/control/phaseplot.py
@@ -36,7 +36,7 @@
 from scipy.integrate import odeint
 
 from . import config
-from .ctrlplot import ControlPlot, _add_arrows_to_line2D
+from .ctrlplot import ControlPlot, _add_arrows_to_line2D, _process_ax_keyword
 from .exception import ControlNotImplemented
 from .nlsys import NonlinearIOSystem, find_eqpt, input_output_response
 
@@ -141,6 +141,9 @@ def phase_plane_plot(
     pointdata = [-1, 1, -1, 1] if pointdata is None else pointdata
 
     # Create axis if needed
+    user_ax = ax
+    # TODO: make use of _process_ax_keyword
+    # fig, ax = _process_ax_keyword(user_ax, squeeze=True)
     if ax is None:
         fig, ax = plt.gcf(), plt.gca()
     else:
@@ -212,7 +215,8 @@ def _create_kwargs(global_kwargs, local_kwargs, **other_kwargs):
     if initial_kwargs:
         raise TypeError("unrecognized keywords: ", str(initial_kwargs))
 
-    if fig is not None:
+    # TODO: update to common code pattern
+    if user_ax is None:
         ax.set_title(f"Phase portrait for {sys.name}")
         ax.set_xlabel(sys.state_labels[0])
         ax.set_ylabel(sys.state_labels[1])
diff --git a/control/pzmap.py b/control/pzmap.py
@@ -315,9 +315,15 @@ def pole_zero_plot(
 
     # Initialize the figure
     # TODO: turn into standard utility function (from plotutil.py?)
+    # fig, ax = _process_ax_keyword(
+    #     user_ax, rcParams=freqplot_rcParams, squeeze=True, create_axes=False)
+    # axs = [ax] if ax is not None else []
     if user_ax is None:
         fig = plt.gcf()
         axs = fig.get_axes()
+    elif isinstance(user_ax, np.ndarray):
+        axs = user_ax.reshape(-1)
+        fig = axs[0].figure
     else:
         fig = ax.figure
         axs = [ax]
diff --git a/control/tests/ctrlplot_test.py b/control/tests/ctrlplot_test.py
@@ -1,11 +1,126 @@
 # ctrlplot_test.py - test out control plotting utilities
 # RMM, 27 Jun 2024
 
+import inspect
+import warnings
+
 import matplotlib.pyplot as plt
+import numpy as np
 import pytest
 
 import control as ct
 
+# List of all plotting functions
+resp_plot_fcns = [
+    # response function                 plotting function
+    (ct.frequency_response,             ct.bode_plot),
+    (ct.frequency_response,             ct.nichols_plot),
+    (ct.singular_values_response,       ct.singular_values_plot),
+    (ct.gangof4_response,               ct.gangof4_plot),
+    (ct.describing_function_response,   ct.describing_function_plot),
+    (None,                              ct.phase_plane_plot),
+    (ct.pole_zero_map,                  ct.pole_zero_plot),
+    (ct.nyquist_response,               ct.nyquist_plot),
+    (ct.root_locus_map,                 ct.root_locus_plot),
+    (ct.initial_response,               ct.time_response_plot),
+    (ct.step_response,                  ct.time_response_plot),
+    (ct.impulse_response,               ct.time_response_plot),
+    (ct.forced_response,                ct.time_response_plot),
+    (ct.input_output_response,          ct.time_response_plot),
+]
+
+deprecated_fcns = [ct.phase_plot]
+
+# Make sure we didn't miss any plotting functions
+def test_find_respplot_functions():
+    # Get the list of plotting functions
+    plot_fcns = {respplot[1] for respplot in resp_plot_fcns}
+
+    # Look through every object in the package
+    found = 0
+    for name, obj in inspect.getmembers(ct):
+        # Skip anything that is outside of this module
+        if inspect.getmodule(obj) is not None and \
+           not inspect.getmodule(obj).__name__.startswith('control'):
+            # Skip anything that isn't part of the control package
+            continue
+
+        # Only look for non-deprecated functions ending in 'plot'
+        if not inspect.isfunction(obj) or name[-4:] != 'plot' or \
+           obj in deprecated_fcns:
+            continue
+
+        # Make sure that we have this on our list of functions
+        assert obj in plot_fcns
+        found += 1
+
+    assert found == len(plot_fcns)
+
+
+@pytest.mark.parametrize("resp_fcn, plot_fcn", resp_plot_fcns)
+@pytest.mark.usefixtures('mplcleanup')
+def test_plot_functions(resp_fcn, plot_fcn):
+    # Create some systems to use
+    sys1 = ct.rss(2, 1, 1, strictly_proper=True)
+    sys2 = ct.rss(4, 1, 1, strictly_proper=True)
+
+    # Set up arguments
+    kwargs = meth_kwargs = plot_fcn_kwargs = {}
+    match resp_fcn, plot_fcn:
+        case ct.describing_function_response, _:
+            F = ct.descfcn.saturation_nonlinearity(1)
+            amp = np.linspace(1, 4, 10)
+            args = (sys1, F, amp)
+
+        case ct.gangof4_response, _:
+            args = (sys1, sys2)
+
+        case ct.frequency_response, ct.nichols_plot:
+            args = (sys1, )
+            meth_kwargs = {'plot_type': 'nichols'}
+
+        case ct.root_locus_map, ct.root_locus_plot:
+            args = (sys1, )
+            meth_kwargs = plot_fcn_kwargs = {'interactive': False}
+
+        case (ct.forced_response | ct.input_output_response, _):
+            timepts = np.linspace(1, 10)
+            U = np.sin(timepts)
+            args = (sys1, timepts, U)
+
+        case _, _:
+            args = (sys1, )
+
+    # Call the plot through the response function
+    if resp_fcn is not None:
+        resp = resp_fcn(*args, **kwargs)
+        cplt1 = resp.plot(**kwargs, **meth_kwargs)
+        assert isinstance(cplt1, ct.ControlPlot)
+
+    # Call the plot directly, plotting on top of previous plot
+    if plot_fcn not in [None, ct.time_response_plot]:
+        cplt2 = plot_fcn(*args, **kwargs, **plot_fcn_kwargs)
+        assert isinstance(cplt2, ct.ControlPlot)
+
+        # Plot should have landed on top of previous plot
+        if resp_fcn is not None:
+            assert cplt2.figure == cplt1.figure
+            if plot_fcn != ct.root_locus_plot:
+                assert np.all(cplt2.axes == cplt1.axes)
+            else:
+                warnings.warn("test skipped for root locus plot")
+            assert len(cplt2.lines[0]) == len(cplt1.lines[0])
+
+    # Pass axes explicitly
+    if resp_fcn is not None:
+        cplt3 = resp.plot(**kwargs, **meth_kwargs, ax=cplt1.axes)
+        assert cplt3.figure == cplt1.figure
+        if plot_fcn != ct.root_locus_plot:
+            assert np.all(cplt3.axes == cplt1.axes)
+        else:
+            warnings.warn("test skipped for root locus plot")
+        assert len(cplt3.lines[0]) == len(cplt1.lines[0])
+
 
 @pytest.mark.usefixtures('mplcleanup')
 def test_rcParams():
