Merge pull request #723 from roryyorke/rory/nichols-improvements

bnavigator · web-flow · commit ae8d5867d3b5 · 2022-04-30T11:56:30.000+02:00
Improvements to Nichols chart plotting
diff --git a/control/nichols.py b/control/nichols.py
@@ -51,6 +51,8 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
+import matplotlib.transforms
+
 from .ctrlutil import unwrap
 from .freqplot import _default_frequency_range
 from . import config
@@ -119,7 +121,18 @@ def nichols_plot(sys_list, omega=None, grid=None):
         nichols_grid()
 
 
-def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted'):
+def _inner_extents(ax):
+    # intersection of data and view extents
+    # if intersection empty, return view extents
+    _inner = matplotlib.transforms.Bbox.intersection(ax.viewLim, ax.dataLim)
+    if _inner is None:
+        return ax.ViewLim.extents
+    else:
+        return _inner.extents
+
+
+def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted', ax=None,
+                 label_cl_phases=True):
     """Nichols chart grid
 
     Plots a Nichols chart grid on the current axis, or creates a new chart
@@ -136,17 +149,36 @@ def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted'):
     line_style : string, optional
         :doc:`Matplotlib linestyle \
             <matplotlib:gallery/lines_bars_and_markers/linestyles>`
+    ax : matplotlib.axes.Axes, optional
+        Axes to add grid to.  If ``None``, use ``plt.gca()``.
+    label_cl_phases: bool, optional
+        If True, closed-loop phase lines will be labelled.
 
+    Returns
+    -------
+    cl_mag_lines: list of `matplotlib.line.Line2D`
+      The constant closed-loop gain contours
+    cl_phase_lines: list of `matplotlib.line.Line2D`
+      The constant closed-loop phase contours
+    cl_mag_labels: list of `matplotlib.text.Text`
+      mcontour labels; each entry corresponds to the respective entry
+      in ``cl_mag_lines``
+    cl_phase_labels: list of `matplotlib.text.Text`
+      ncontour labels; each entry corresponds to the respective entry
+      in ``cl_phase_lines``
     """
+    if ax is None:
+        ax = plt.gca()
+
     # Default chart size
     ol_phase_min = -359.99
     ol_phase_max = 0.0
     ol_mag_min = -40.0
     ol_mag_max = default_ol_mag_max = 50.0
 
-    # Find bounds of the current dataset, if there is one.
-    if plt.gcf().gca().has_data():
-        ol_phase_min, ol_phase_max, ol_mag_min, ol_mag_max = plt.axis()
+    if ax.has_data():
+        # Find extent of intersection the current dataset or view
+        ol_phase_min, ol_mag_min, ol_phase_max, ol_mag_max = _inner_extents(ax)
 
     # M-circle magnitudes.
     if cl_mags is None:
@@ -165,19 +197,22 @@ def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted'):
                                      ol_mag_min + cl_mag_step, cl_mag_step)
         cl_mags = np.concatenate((extended_cl_mags, key_cl_mags))
 
+    # a minimum 360deg extent containing the phases
+    phase_round_max = 360.0*np.ceil(ol_phase_max/360.0)
+    phase_round_min = min(phase_round_max-360,
+                          360.0*np.floor(ol_phase_min/360.0))
+
     # N-circle phases (should be in the range -360 to 0)
     if cl_phases is None:
-        # Choose a reasonable set of default phases (denser if the open-loop
-        # data is restricted to a relatively small range of phases).
-        key_cl_phases = np.array([-0.25, -45.0, -90.0, -180.0, -270.0,
-                                  -325.0, -359.75])
-        if np.abs(ol_phase_max - ol_phase_min) < 90.0:
-            other_cl_phases = np.arange(-10.0, -360.0, -10.0)
-        else:
-            other_cl_phases = np.arange(-10.0, -360.0, -20.0)
-        cl_phases = np.concatenate((key_cl_phases, other_cl_phases))
-    else:
-        assert ((-360.0 < np.min(cl_phases)) and (np.max(cl_phases) < 0.0))
+        # aim for 9 lines, but always show (-360+eps, -180, -eps)
+        # smallest spacing is 45, biggest is 180
+        phase_span = phase_round_max - phase_round_min
+        spacing = np.clip(round(phase_span / 8 / 45) * 45, 45, 180)
+        key_cl_phases = np.array([-0.25, -359.75])
+        other_cl_phases = np.arange(-spacing, -360.0, -spacing)
+        cl_phases = np.unique(np.concatenate((key_cl_phases, other_cl_phases)))
+    elif not ((-360 < np.min(cl_phases)) and (np.max(cl_phases) < 0.0)):
+        raise ValueError('cl_phases must between -360 and 0, exclusive')
 
     # Find the M-contours
     m = m_circles(cl_mags, phase_min=np.min(cl_phases),
@@ -196,27 +231,57 @@ def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted'):
     # over the range -360 < phase < 0. Given the range
     # the base chart is computed over, the phase offset should be 0
     # for -360 < ol_phase_min < 0.
-    phase_offset_min = 360.0*np.ceil(ol_phase_min/360.0)
-    phase_offset_max = 360.0*np.ceil(ol_phase_max/360.0) + 360.0
-    phase_offsets = np.arange(phase_offset_min, phase_offset_max, 360.0)
+    phase_offsets = 360 + np.arange(phase_round_min, phase_round_max, 360.0)
+
+    cl_mag_lines = []
+    cl_phase_lines = []
+    cl_mag_labels = []
+    cl_phase_labels = []
 
     for phase_offset in phase_offsets:
         # Draw M and N contours
-        plt.plot(m_phase + phase_offset, m_mag, color='lightgray',
-                 linestyle=line_style, zorder=0)
-        plt.plot(n_phase + phase_offset, n_mag, color='lightgray',
-                 linestyle=line_style, zorder=0)
+        cl_mag_lines.extend(
+            ax.plot(m_phase + phase_offset, m_mag, color='lightgray',
+                    linestyle=line_style, zorder=0))
+        cl_phase_lines.extend(
+            ax.plot(n_phase + phase_offset, n_mag, color='lightgray',
+                    linestyle=line_style, zorder=0))
 
         # Add magnitude labels
         for x, y, m in zip(m_phase[:][-1] + phase_offset, m_mag[:][-1],
                            cl_mags):
             align = 'right' if m < 0.0 else 'left'
-            plt.text(x, y, str(m) + ' dB', size='small', ha=align,
-                     color='gray')
+            cl_mag_labels.append(
+                ax.text(x, y, str(m) + ' dB', size='small', ha=align,
+                        color='gray', clip_on=True))
+
+        # phase labels
+        if label_cl_phases:
+            for x, y, p in zip(n_phase[:][0] + phase_offset,
+                               n_mag[:][0],
+                               cl_phases):
+                if p > -175:
+                    align = 'right'
+                elif p > -185:
+                    align = 'center'
+                else:
+                    align = 'left'
+                cl_phase_labels.append(
+                    ax.text(x, y, f'{round(p)}\N{DEGREE SIGN}',
+                            size='small',
+                            ha=align,
+                            va='bottom',
+                            color='gray',
+                            clip_on=True))
+
 
     # Fit axes to generated chart
-    plt.axis([phase_offset_min - 360.0, phase_offset_max - 360.0,
-              np.min(cl_mags), np.max([ol_mag_max, default_ol_mag_max])])
+    ax.axis([phase_round_min,
+             phase_round_max,
+             np.min(np.concatenate([cl_mags,[ol_mag_min]])),
+             np.max([ol_mag_max, default_ol_mag_max])])
+
+    return cl_mag_lines, cl_phase_lines, cl_mag_labels, cl_phase_labels
 
 #
 # Utility functions
diff --git a/control/tests/nichols_test.py b/control/tests/nichols_test.py
@@ -3,9 +3,11 @@
 RMM, 31 Mar 2011
 """
 
+import matplotlib.pyplot as plt
+
 import pytest
 
-from control import StateSpace, nichols_plot, nichols
+from control import StateSpace, nichols_plot, nichols, nichols_grid, pade, tf
 
 
 @pytest.fixture()
@@ -26,3 +28,68 @@ def test_nichols(tsys, mplcleanup):
 def test_nichols_alias(tsys, mplcleanup):
     """Test the control.nichols alias and the grid=False parameter"""
     nichols(tsys, grid=False)
+
+
+@pytest.mark.usefixtures("mplcleanup")
+class TestNicholsGrid:
+    def test_ax(self):
+        # check grid is plotted into gca, or specified axis
+        fig, axs = plt.subplots(2,2)
+        plt.sca(axs[0,1])
+
+        cl_mag_lines = nichols_grid()[1]
+        assert cl_mag_lines[0].axes is axs[0, 1]
+
+        cl_mag_lines = nichols_grid(ax=axs[1,1])[1]
+        assert cl_mag_lines[0].axes is axs[1, 1]
+        # nichols_grid didn't change what the "current axes" are
+        assert plt.gca() is axs[0, 1]
+
+
+    def test_cl_phase_label_control(self):
+        # test label_cl_phases argument
+        cl_mag_lines, cl_phase_lines, cl_mag_labels, cl_phase_labels \
+            = nichols_grid()
+        assert len(cl_phase_labels) > 0
+
+        cl_mag_lines, cl_phase_lines, cl_mag_labels, cl_phase_labels \
+            = nichols_grid(label_cl_phases=False)
+        assert len(cl_phase_labels) == 0
+
+
+    def test_labels_clipped(self):
+        # regression test: check that contour labels are clipped
+        mcontours, ncontours, mlabels, nlabels = nichols_grid()
+        assert all(ml.get_clip_on() for ml in mlabels)
+        assert all(nl.get_clip_on() for nl in nlabels)
+
+
+    def test_minimal_phase(self):
+        # regression test: phase extent is minimal
+        g = tf([1],[1,1]) * tf([1],[1/1, 2*0.1/1, 1])
+        nichols(g)
+        ax = plt.gca()
+        assert ax.get_xlim()[1] <= 0
+
+
+    def test_fixed_view(self):
+        # respect xlim, ylim set by user
+        g = (tf([1],[1/1, 2*0.01/1, 1])
+             * tf([1],[1/100**2, 2*0.001/100, 1])
+             * tf(*pade(0.01, 5)))
+
+        # normally a broad axis
+        nichols(g)
+
+        assert(plt.xlim()[0] == -1440)
+        assert(plt.ylim()[0] <= -240)
+
+        nichols(g, grid=False)
+
+        # zoom in
+        plt.axis([-360,0,-40,50])
+
+        # nichols_grid doesn't expand limits
+        nichols_grid()
+        assert(plt.xlim()[0] == -360)
+        assert(plt.ylim()[1] >= -40)
