nichols_grid returns artists of grid elements; further tighten phase extent

roryyorke · roryyorke · commit 3281a8d6eafe · 2022-04-27T10:05:23.000+02:00
diff --git a/control/nichols.py b/control/nichols.py
@@ -153,6 +153,19 @@ def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted', ax=None,
         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()
@@ -163,8 +176,8 @@ def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted', ax=None,
     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 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.
@@ -184,20 +197,22 @@ def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted', ax=None,
                                      ol_mag_min + cl_mag_step, cl_mag_step)
         cl_mags = np.concatenate((extended_cl_mags, key_cl_mags))
 
-    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
+    # 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:
         # aim for 9 lines, but always show (-360+eps, -180, -eps)
         # smallest spacing is 45, biggest is 180
-        phase_span = phase_offset_max - phase_offset_min
+        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)))
-    else:
-        assert ((-360.0 < np.min(cl_phases)) and (np.max(cl_phases) < 0.0))
+    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),
@@ -216,21 +231,29 @@ def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted', ax=None,
     # 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_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
-        ax.plot(m_phase + phase_offset, m_mag, color='lightgray',
-                 linestyle=line_style, zorder=0)
-        ax.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'
-            ax.text(x, y, str(m) + ' dB', size='small', ha=align,
-                     color='gray', clip_on=True)
+            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:
@@ -243,20 +266,23 @@ def nichols_grid(cl_mags=None, cl_phases=None, line_style='dotted', ax=None,
                     align = 'center'
                 else:
                     align = 'left'
-                ax.text(x, y, f'{round(p)}\N{DEGREE SIGN}',
-                        size='small',
-                        ha=align,
-                        va='bottom',
-                        color='gray',
-                        clip_on=True)
+                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
-    ax.axis([phase_offset_min - 360.0,
-             phase_offset_max - 360.0,
+    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
 #
