Added docstring and unittest

murrayrm · Apr 16, 2019 · May 23, 2018 · May 24, 2018 · May 25, 2018 · May 25, 2018
commit 1e3a925ccc804348c52a1f50fda460df30d2240d
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -213,8 +213,10 @@ def bode_plot(syslist, omega=None, dB=None, Hz=None, deg=None,
                     pltline = ax_mag.semilogx(omega_plot, 20 * np.log10(mag),
                                               *args, **kwargs)
                 else:
+
                     pltline = ax_mag.loglog(omega_plot, mag, *args, **kwargs)
 
+
                 if nyquistfrq_plot:
                     ax_mag.axvline(nyquistfrq_plot,
                                    color=pltline[0].get_color())

diff --git a/control/rlocus.py b/control/rlocus.py
@@ -380,6 +380,7 @@ def _RLFeedbackClicksSisotool(event,sys,fig,bode_plot_params,tvect):
 def _RLFeedbackClicksPoint(event,sys,fig,ax_rlocus=None,sisotool=False):
     """Display root-locus gain feedback point for clicks on the root-locus plot
     """
+    print(event)
     if sisotool == False:
         ax_rlocus = fig.axes[0]
 
@@ -413,8 +414,6 @@ def _RLFeedbackClicksPoint(event,sys,fig,ax_rlocus=None,sisotool=False):
 
         return True
 
-
-
 def _sgrid_func(fig=None, zeta=None, wn=None):
     if fig is None:
         fig = pylab.gcf()

diff --git a/control/sisotool.py b/control/sisotool.py
@@ -5,9 +5,56 @@
 from .lti import issiso
 import matplotlib
 import matplotlib.pyplot as plt
-
-def sisotool(sys, kvect = None, xlim = None, ylim = None, plotstr_rlocus = 'b' if int(matplotlib.__version__[0]) == 1 else 'C0',rlocus_grid = False, omega = None, dB = None, Hz = None, deg = None, omega_limits = None, omega_num = None,margins = True, tvect=None):
-
+import warnings
+
+def sisotool(sys, kvect = None, xlim_rlocus = None, ylim_rlocus = None, plotstr_rlocus = 'b' if int(matplotlib.__version__[0]) == 1 else 'C0',rlocus_grid = False, omega = None, dB = None, Hz = None, deg = None, omega_limits = None, omega_num = None,margins_bode = True, tvect=None):
+
+    """Sisotool
+
+    Sisotool style collection of plots inspired by the matlab sisotool.
+    The left two plots contain the bode magnitude and phase diagrams.
+    The top right plot is a clickable root locus plot, clicking on the
+    root locus will change the gain of the system. The bottom left plot
+    shows a closed loop time response.
+
+    Parameters
+    ----------
+    sys : LTI object
+        Linear input/output systems (SISO only)
+    kvect : list or ndarray, optional
+        List of gains to use for plotting root locus
+    xlim_rlocus : tuple or list, optional
+        control of x-axis range, normally with tuple (see matplotlib.axes)
+    ylim_rlocus : tuple or list, optional
+        control of y-axis range
+    plotstr_rlocus : Additional options to matplotlib
+        plotting style for the root locus plot(color, linestyle, etc)
+    rlocus_grid: boolean (default = False)
+        If True plot s-plane grid.
+    omega : freq_range
+        Range of frequencies in rad/sec for the bode plot
+    dB : boolean
+        If True, plot result in dB for the bode plot
+    Hz : boolean
+        If True, plot frequency in Hz for the bode plot (omega must be provided in rad/sec)
+    deg : boolean
+        If True, plot phase in degrees for the bode plot (else radians)
+    omega_limits: tuple, list, ... of two values
+        Limits of the to generate frequency vector.
+        If Hz=True the limits are in Hz otherwise in rad/s.
+    omega_num: int
+        number of samples
+    margins_bode : boolean
+        If True, plot gain and phase margin in the bode plot
+    tvect : list or ndarray, optional
+        List of timesteps to use for closed loop step response
+
+    Examples
+    --------
+    >>> sys = tf([1000], [1,25,100,0])
+    >>> sisotool(sys)
+
+    """
     from .rlocus import root_locus
 
     # Check if it is a single SISO system
@@ -30,14 +77,14 @@ def sisotool(sys, kvect = None, xlim = None, ylim = None, plotstr_rlocus = 'b' i
         'omega_num' : omega_num,
         'sisotool': True,
         'fig': fig,
-        'margins': margins
+        'margins': margins_bode
     }
 
     # First time call to setup the bode and step response plots
     _SisotoolUpdate(sys, fig,1 if kvect is None else kvect[0],bode_plot_params)
 
     # Setup the root-locus plot window
-    root_locus(sys,kvect=kvect,xlim=xlim,ylim = ylim,plotstr=plotstr_rlocus,grid = rlocus_grid,fig=fig,bode_plot_params=bode_plot_params,tvect=tvect,sisotool=True)
+    root_locus(sys,kvect=kvect,xlim=xlim_rlocus,ylim = ylim_rlocus,plotstr=plotstr_rlocus,grid = rlocus_grid,fig=fig,bode_plot_params=bode_plot_params,tvect=tvect,sisotool=True)
 
 def _SisotoolUpdate(sys,fig,K,bode_plot_params,tvect=None):
 
@@ -50,7 +97,11 @@ def _SisotoolUpdate(sys,fig,K,bode_plot_params,tvect=None):
 
     # Get the subaxes and clear them
     ax_mag,ax_rlocus,ax_phase,ax_step = fig.axes[0],fig.axes[1],fig.axes[2],fig.axes[3]
-    ax_mag.cla(),ax_phase.cla(),ax_step.cla()
+
+    # Catch matplotlib 2.1.x and higher userwarnings when clearing a log axis
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        ax_step.clear(), ax_mag.clear(), ax_phase.clear()
 
     # Update the bodeplot
     bode_plot_params['syslist'] = sys*K.real

diff --git a/control/tests/sisotool_test.py b/control/tests/sisotool_test.py
@@ -0,0 +1,78 @@
+import unittest
+import numpy as np
+from control.sisotool import sisotool
+from control.tests.margin_test import assert_array_almost_equal
+from control.rlocus import _RLFeedbackClicksSisotool
+from control.xferfcn import TransferFunction
+import matplotlib.pyplot as plt
+
+class TestSisotool(unittest.TestCase):
+    """These are tests for the sisotool in sisotool.py."""
+
+    def setUp(self):
+        """This contains some random LTI systems and scalars for testing."""
+
+        # Two random SISO systems.
+        sys1 = TransferFunction([1000],[1,25,100,0])
+        sys2 = TransferFunction([1,1,1],[1,1])
+        self.systems = (sys1, sys2)
+
+    def test_sisotool(self):
+        sisotool(self.systems[0])
+        fig = plt.gcf()
+        ax_mag,ax_rlocus,ax_phase,ax_step = fig.axes[0],fig.axes[1],fig.axes[2],fig.axes[3]
+
+        # Check the initial root locus plot points
+        initial_point_0 = (np.array([-22.53155977]),np.array([0.]))
+        initial_point_1 = (np.array([-1.23422011]), np.array([-6.54667031]))
+        initial_point_2 = (np.array([-1.23422011]), np.array([06.54667031]))
+        assert_array_almost_equal(ax_rlocus.lines[0].get_data(),initial_point_0)
+        assert_array_almost_equal(ax_rlocus.lines[1].get_data(),initial_point_1)
+        assert_array_almost_equal(ax_rlocus.lines[2].get_data(),initial_point_2)
+
+        # Check the bode plot magnitude line
+        bode_mag_original = np.array([ 15.12670407,  12.48358602,  10.28367521,   8.44961576,   6.91740918, 5.63436301,   4.55750665,   3.65237575,   2.89200358,   2.25587634])
+        assert_array_almost_equal(ax_mag.lines[0].get_data()[1][10:20],bode_mag_original)
+
+        # Check the step response before moving the point
+        print(ax_step.lines[0].get_data()[1][:10])
+        step_response_original = np.array([ 0.,          0.02233651,  0.13118374,  0.33078542,  0.5907113,   0.87041549, 1.13038536,  1.33851053,  1.47374666,  1.52757114])
+        assert_array_almost_equal(ax_step.lines[0].get_data()[1][:10],step_response_original)
+
+        bode_plot_params = {
+            'omega': None,
+            'dB': False,
+            'Hz': False,
+            'deg': True,
+            'omega_limits': None,
+            'omega_num': None,
+            'sisotool': True,
+            'fig': fig,
+            'margins': True
+        }
+
+        # Move the rootlocus to another point
+        event = type('test', (object,), {'xdata': 2.31206868287,'ydata':15.5983051046, 'inaxes':ax_rlocus.axes})()
+        _RLFeedbackClicksSisotool(event=event, sys=self.systems[0], fig=fig, bode_plot_params=bode_plot_params, tvect=None)
+
+        # Check the moved root locus plot points
+        moved_point_0 = (np.array([-29.91742755]), np.array([0.]))
+        moved_point_1 = (np.array([2.45871378]), np.array([-15.52647768]))
+        moved_point_2 = (np.array([2.45871378]), np.array([15.52647768]))
+        assert_array_almost_equal(ax_rlocus.lines[-3].get_data(),moved_point_0)
+        assert_array_almost_equal(ax_rlocus.lines[-2].get_data(),moved_point_1)
+        assert_array_almost_equal(ax_rlocus.lines[-1].get_data(),moved_point_2)
+
+        # Check if the bode_mag line has moved
+        bode_mag_moved = np.array([ 111.83321224,   92.29238035,   76.02822315,   62.46884113,   51.14108703, 41.6554004,    33.69409534,   27.00237344,   21.38086717,   16.67791585])
+        assert_array_almost_equal(ax_mag.lines[0].get_data()[1][10:20],bode_mag_moved)
+
+        # Check if the step response has changed
+        step_response_moved = np.array([[ 0.,          0.02458187,  0.16529784 , 0.46602716 , 0.91012035 , 1.43364313, 1.93996334 , 2.3190105  , 2.47041552 , 2.32724853] ])
+        assert_array_almost_equal(ax_step.lines[0].get_data()[1][:10],step_response_moved)
+
+def test_suite():
+    return unittest.TestLoader().loadTestsFromTestCase(TestSisotool)
+
+if __name__ == "__main__":
+    unittest.main()