Added unit test for rootlocus zoom functionality

icam0 · icam0 · commit c26187f9d46f · 2018-06-03T15:15:12.000+02:00
diff --git a/control/rlocus.py b/control/rlocus.py
@@ -174,7 +174,6 @@ def _default_gains(num, den, xlim, ylim,zoom_xlim=None,zoom_ylim=None):
 
     k_break, real_break = _break_points(num, den)
     kmax = _k_max(num, den, real_break, k_break)
-    print(kmax)
     kvect = np.hstack((np.linspace(0, kmax, 50), np.real(k_break)))
     kvect.sort()
 
diff --git a/control/tests/rlocus_test.py b/control/tests/rlocus_test.py
@@ -5,10 +5,13 @@
 
 import unittest
 import numpy as np
-from control.rlocus import root_locus
+from control.rlocus import root_locus, _RLClickDispatcher
 from control.xferfcn import TransferFunction
 from control.statesp import StateSpace
 from control.bdalg import feedback
+import matplotlib.pyplot as plt
+from control.tests.margin_test import assert_array_almost_equal
+
 
 class TestRootLocus(unittest.TestCase):
     """These are tests for the feedback function in rlocus.py."""
@@ -42,6 +45,31 @@ def test_without_gains(self):
             roots, kvect = root_locus(sys, Plot=False)
             self.check_cl_poles(sys, roots, kvect)
 
+    def testRootLocusZoom(self):
+        """Check the zooming functionality of the Root locus plot"""
+        system = TransferFunction([1000], [1, 25, 100, 0])
+        root_locus(system)
+        fig = plt.gcf()
+        ax_rlocus = fig.axes[0]
+
+        event = type('test', (object,), {'xdata': 0.9957380594313321, 'ydata': 1.7825491928580846, 'inaxes': ax_rlocus.axes})()
+        ax_rlocus.set_xlim((-4.420022219849855, 0.9957380594313321))
+        ax_rlocus.set_ylim((1.7825491928580846, 10.695295157148486))
+        plt.get_current_fig_manager().toolbar.mode = 'zoom rect'
+        _RLClickDispatcher(event, system, fig, ax_rlocus, '-')
+
+        zoom_x = ax_rlocus.lines[-2].get_data()[0][65:75]
+        zoom_y = ax_rlocus.lines[-2].get_data()[1][65:75]
+        zoom_y = [abs(y) for y in zoom_y]
+
+        zoom_x_valid = [-2.23659192, -2.23659121, -2.23659103, -2.23659086, -2.23659068, -2.23659063,
+         -2.23659059, -2.23659055, -2.2365905,  -2.23658766]
+        zoom_y_valid = [1.78253589, 1.78254318, 1.782545,   1.78254682, 1.78254864, 1.7825491,
+         1.78254955, 1.78255001, 1.78255046, 1.7825796]
+
+        assert_array_almost_equal(zoom_x,zoom_x_valid)
+        assert_array_almost_equal(zoom_y,zoom_y_valid)
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromTestCase(TestRootLocus)
 
