new indent code that adds a consistent number of points at every indent if omega_range_given is false

sawyerbfuller · sawyerbfuller · commit 23b3e8c74cdb · 2021-11-22T12:58:15.000-08:00
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -522,7 +522,7 @@ def gen_zero_centered_series(val_min, val_max, period):
     'nyquist.mirror_style': '--',
     'nyquist.arrows': 2,
     'nyquist.arrow_size': 8,
-    'nyquist.indent_radius': 1e-2,
+    'nyquist.indent_radius': 1e-6,
     'nyquist.maximum_magnitude': 5,
     'nyquist.indent_direction': 'right',
 }
@@ -552,7 +552,9 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
         If True, plot magnitude
 
     omega : array_like
-        Set of frequencies to be evaluated, in rad/sec.
+        Set of frequencies to be evaluated, in rad/sec. Remark: if omega is
+        specified, you may need to specify specify a larger `indent_radius` 
+        than the default to get reasonable contours. 
 
     omega_limits : array_like of two values
         Limits to the range of frequencies. Ignored if omega is provided, and
@@ -725,7 +727,7 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
         # do indentations in s-plane where it is more convenient
         splane_contour = 1j * omega_sys
 
-        # Bend the contour around any poles on/near the imaginary axis
+        # indent the contour around any poles on/near the imaginary axis
         if isinstance(sys, (StateSpace, TransferFunction)) \
                 and indent_direction != 'none':
             if sys.isctime():
@@ -738,30 +740,49 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
                 zplane_poles = zplane_poles[~np.isclose(abs(zplane_poles), 0.)]
                 splane_poles = np.log(zplane_poles)/sys.dt
 
-            if splane_contour[1].imag > indent_radius \
-                    and np.any(np.isclose(abs(splane_poles), 0)) \
-                    and not omega_range_given:
-                # add some points for quarter circle around poles at origin
-                splane_contour = np.concatenate(
-                    (1j * np.linspace(0., indent_radius, 50),
-                    splane_contour[1:]))
-            for i, s in enumerate(splane_contour):
-                # Find the nearest pole
-                p = splane_poles[(np.abs(splane_poles - s)).argmin()]
-                # See if we need to indent around it
-                if abs(s - p) < indent_radius:
-                    if p.real < 0 or (np.isclose(p.real, 0) \
-                            and indent_direction == 'right'):
-                        # Indent to the right
-                        splane_contour[i] += \
-                            np.sqrt(indent_radius ** 2 - (s-p).imag ** 2)
-                    elif p.real > 0 or (np.isclose(p.real, 0) \
-                            and indent_direction == 'left'):
-                        # Indent to the left
-                        splane_contour[i] -= \
-                            np.sqrt(indent_radius ** 2 - (s-p).imag ** 2)
-                    else:
-                        ValueError("unknown value for indent_direction")
+            if omega_range_given:
+                # indent given contour
+                for i, s in enumerate(splane_contour):
+                    # Find the nearest pole
+                    p = splane_poles[(np.abs(splane_poles - s)).argmin()]
+                    # See if we need to indent around it
+                    if abs(s - p) < indent_radius:
+                        if p.real < 0 or (np.isclose(p.real, 0) \
+                                and indent_direction == 'right'):
+                            # Indent to the right
+                            splane_contour[i] += \
+                                np.sqrt(indent_radius ** 2 - (s-p).imag ** 2)
+                        elif p.real > 0 or (np.isclose(p.real, 0) \
+                                and indent_direction == 'left'):
+                            # Indent to the left
+                            splane_contour[i] -= \
+                                np.sqrt(indent_radius ** 2 - (s-p).imag ** 2)
+                        else:
+                            ValueError("unknown value for indent_direction")
+            else:
+                # find poles that are near imag axis and replace contour 
+                # near there with indented contour 
+                if indent_direction == 'right': 
+                    indent = indent_radius * \
+                        np.exp(1j * np.linspace(-np.pi/2, np.pi/2, 51))                
+                elif indent_direction == 'left':
+                    indent = -indent_radius * \
+                        np.exp(1j * np.linspace(np.pi/2, -np.pi/2, 51))                
+                else: 
+                    ValueError("unknown value for indent_direction")
+                for p in splane_poles[np.isclose(splane_poles.real, 0) \
+                        & (splane_poles.imag >= 0)]:
+                    start = np.searchsorted(
+                        splane_contour.imag, p.imag - indent_radius)
+                    end = np.searchsorted(
+                        splane_contour.imag, p.imag + indent_radius)
+                    # only keep indent parts that are > 0 and within contour
+                    indent_mask = ((indent + p).imag >= 0) \
+                        & ((indent + p).imag <= splane_contour[-1].imag)
+                    splane_contour = np.concatenate((
+                        splane_contour[:start], 
+                        indent[indent_mask] + p, 
+                        splane_contour[end:]))
 
         # map contour to z-plane if necessary
         if sys.isctime():
@@ -816,7 +837,7 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
             _add_arrows_to_line2D(
                 ax, p[0], arrow_pos, arrowstyle=arrow_style, dir=1)
             # plot large magnitude contour
-            p = plt.plot(x_lg, y_lg, ':', color='red', *args, **kwargs)
+            p = plt.plot(x_lg, y_lg, color='red', *args, **kwargs)
             _add_arrows_to_line2D(
                 ax, p[0], arrow_pos, arrowstyle=arrow_style, dir=1)
 
@@ -825,7 +846,8 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
                 p = plt.plot(x, -y, mirror_style, color=c, *args, **kwargs)
                 _add_arrows_to_line2D(
                     ax, p[0], arrow_pos, arrowstyle=arrow_style, dir=-1)
-                p = plt.plot(x_lg, -y_lg, ':', color='red', *args, **kwargs)
+                p = plt.plot(x_lg, -y_lg, mirror_style, 
+                        color='red', *args, **kwargs)
                 _add_arrows_to_line2D(
                     ax, p[0], arrow_pos, arrowstyle=arrow_style, dir=-1)
 
diff --git a/control/tests/nyquist_test.py b/control/tests/nyquist_test.py
@@ -69,7 +69,8 @@ def test_nyquist_basic():
 
     # Make sure that we can turn off frequency modification
     count, contour_indented = ct.nyquist_plot(
-        sys, np.linspace(1e-4, 1e2, 100), return_contour=True)
+        sys, np.linspace(1e-4, 1e2, 100), indent_radius=0.01, 
+        return_contour=True)
     assert not all(contour_indented.real == 0)
     count, contour = ct.nyquist_plot(
         sys, np.linspace(1e-4, 1e2, 100), return_contour=True,
@@ -87,14 +88,17 @@ def test_nyquist_basic():
     assert _Z(sys) == count + _P(sys)
 
     # Nyquist plot with poles on imaginary axis, omega specified
+    # when specifying omega, usually need to specify larger indent radius
     sys = ct.tf([1], [1, 3, 2]) * ct.tf([1], [1, 0, 1])
-    count = ct.nyquist_plot(sys, np.linspace(1e-3, 1e1, 1000))
+    count = ct.nyquist_plot(sys, np.linspace(1e-3, 1e1, 1000), 
+        indent_radius=0.1)
     assert _Z(sys) == count + _P(sys)
 
     # Nyquist plot with poles on imaginary axis, omega specified, with contour
     sys = ct.tf([1], [1, 3, 2]) * ct.tf([1], [1, 0, 1])
     count, contour = ct.nyquist_plot(
-        sys, np.linspace(1e-3, 1e1, 1000), return_contour=True)
+        sys, np.linspace(1e-3, 1e1, 1000), indent_radius=0.1,
+        return_contour=True)
     assert _Z(sys) == count + _P(sys)
 
     # Nyquist plot with poles on imaginary axis, return contour
@@ -193,13 +197,10 @@ def test_nyquist_indent():
     plt.title("Pole at origin; indent_radius=default")
     assert _Z(sys) == count + _P(sys)
 
-    # first value of default omega vector was 0.1, replaced by 0. for contour
-    # indent_radius is larger than 0.1 -> no extra quater circle around origin
+    # confirm that first element of indent is properly indented
     count, contour = ct.nyquist_plot(sys, plot=False, indent_radius=.1007,
                                      return_contour=True)
     np.testing.assert_allclose(contour[0], .1007+0.j)
-    # second value of omega_vector is larger than indent_radius: not indented
-    assert np.all(contour.real[2:] == 0.)
 
     plt.figure();
     count, contour = ct.nyquist_plot(sys, indent_radius=0.01,
@@ -208,7 +209,7 @@ def test_nyquist_indent():
     assert _Z(sys) == count + _P(sys)
     # indent radius is smaller than the start of the default omega vector
     # check that a quarter circle around the pole at origin has been added.
-    np.testing.assert_allclose(contour[:50].real**2 + contour[:50].imag**2,
+    np.testing.assert_allclose(contour[:25].real**2 + contour[:25].imag**2,
                                0.01**2)
 
     plt.figure();
