fix rlocus timeout due to inefficient _default_wn calculation

murrayrm · bnavigator · commit ff7d82728b81 · 2021-02-05T10:59:37.000+01:00
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -646,8 +646,9 @@ def _sgrid_func(fig=None, zeta=None, wn=None):
     else:
         ax = fig.axes[1]
 
-    # Get locator function for x-axis tick marks
+    # Get locator function for x-axis, y-axis tick marks
     xlocator = ax.get_xaxis().get_major_locator()
+    ylocator = ax.get_yaxis().get_major_locator()
 
     # Decide on the location for the labels (?)
     ylim = ax.get_ylim()
@@ -690,7 +691,7 @@ def _sgrid_func(fig=None, zeta=None, wn=None):
     # omega-constant lines
     angles = np.linspace(-90, 90, 20) * np.pi/180
     if wn is None:
-        wn = _default_wn(xlocator(), ylim)
+        wn = _default_wn(xlocator(), ylocator())
 
     for om in wn:
         if om < 0:
@@ -746,7 +747,7 @@ def _default_zetas(xlim, ylim):
     return zeta.tolist()
 
 
-def _default_wn(xloc, ylim):
+def _default_wn(xloc, yloc, max_lines=7):
     """Return default wn for root locus plot
 
     This function computes a list of natural frequencies based on the grid
@@ -758,23 +759,30 @@ def _default_wn(xloc, ylim):
         List of x-axis tick values
     ylim : array_like
         List of y-axis limits [min, max]
+    max_lines : int, optional
+        Maximum number of frequencies to generate (default = 7)
 
     Returns
     -------
     wn : list
         List of default natural frequencies for the plot
 
     """
+    sep = xloc[1]-xloc[0]       # separation between x-ticks
+    
+    # Decide whether to use the x or y axis for determining wn
+    if yloc[-1] / sep > max_lines*10:
+        # y-axis scale >> x-axis scale
+        wn = yloc                   # one frequency per y-axis tick mark
+    else:
+        wn = xloc                   # one frequency per x-axis tick mark
 
-    wn = xloc                   # one frequency per x-axis tick mark
-    sep = xloc[1]-xloc[0]       # separation between ticks
-
-    # Insert additional frequencies to span the y-axis
-    while np.abs(wn[0]) < ylim[1]:
-        wn = np.insert(wn, 0, wn[0]-sep)
+        # Insert additional frequencies to span the y-axis
+        while np.abs(wn[0]) < yloc[-1]:
+            wn = np.insert(wn, 0, wn[0]-sep)
 
     # If there are too many values, cut them in half
-    while len(wn) > 7:
+    while len(wn) > max_lines:
         wn = wn[0:-1:2]
 
     return wn
diff --git a/control/tests/rlocus_test.py b/control/tests/rlocus_test.py
@@ -8,6 +8,7 @@
 from numpy.testing import assert_array_almost_equal
 import pytest
 
+import control as ct
 from control.rlocus import root_locus, _RLClickDispatcher
 from control.xferfcn import TransferFunction
 from control.statesp import StateSpace
@@ -74,3 +75,31 @@ def test_root_locus_zoom(self):
 
         assert_array_almost_equal(zoom_x, zoom_x_valid)
         assert_array_almost_equal(zoom_y, zoom_y_valid)
+
+    def test_rlocus_default_wn(self):
+        """Check that default wn calculation works properly"""
+        #
+        # System that triggers use of y-axis as basis for wn (for coverage)
+        #
+        # This system generates a root locus plot that used to cause the
+        # creation (and subsequent deletion) of a large number of natural
+        # frequency contours within the `_default_wn` function in `rlocus.py`.
+        # This unit test makes sure that is fixed by generating a test case
+        # that will take a long time to do the calculation (minutes).
+        #
+        import scipy as sp
+        import signal
+
+        # Define a system that exhibits this behavior
+        sys = ct.tf(*sp.signal.zpk2tf(
+            [-1e-2, 1-1e7j, 1+1e7j], [0, -1e7j, 1e7j], 1))
+
+        # Set up a timer to catch execution time
+        def signal_handler(signum, frame):
+            raise Exception("rlocus took too long to complete")
+        signal.signal(signal.SIGALRM, signal_handler)
+
+        # Run the command and reset the alarm
+        signal.alarm(2)         # 2 second timeout
+        ct.root_locus(sys)
+        signal.alarm(0)         # reset the alarm
