TST: add unit test for multiple bode plots on same axes

murrayrm · murrayrm · commit 244429666b55 · 2018-01-20T12:41:22.000-08:00
diff --git a/control/tests/freqresp_test.py b/control/tests/freqresp_test.py
@@ -8,6 +8,7 @@
 
 import unittest
 import numpy as np
+import control as ctrl
 from control.statesp import StateSpace
 from control.matlab import ss, tf, bode
 from control.exception import slycot_check
@@ -34,6 +35,50 @@ def test_siso(self):
       systf = tf(sys)
       bode(systf)
 
+   def test_superimpose(self):
+      # Test to make sure that multiple calls to plots superimpose their
+      # data on the same axes unless told to do otherwise
+
+      # Generate two plots in a row; should be on the same axes
+      plt.figure(1); plt.clf()
+      ctrl.bode_plot(ctrl.tf([1], [1,2,1]))
+      ctrl.bode_plot(ctrl.tf([5], [1, 1]))
+
+      # Check to make sure there are two axes and that each axes has two lines
+      assert len(plt.gcf().axes) == 2
+      for ax in plt.gcf().axes:
+         # Make sure there are 2 lines in each subplot
+         assert len(ax.get_lines()) == 2
+      
+      # Generate two plots as a list; should be on the same axes
+      plt.figure(2); plt.clf();
+      ctrl.bode_plot([ctrl.tf([1], [1,2,1]), ctrl.tf([5], [1, 1])])
+
+      # Check to make sure there are two axes and that each axes has two lines
+      assert len(plt.gcf().axes) == 2
+      for ax in plt.gcf().axes:
+         # Make sure there are 2 lines in each subplot
+         assert len(ax.get_lines()) == 2
+
+      # Generate two separate plots; only the second should appear
+      plt.figure(3); plt.clf();
+      ctrl.bode_plot(ctrl.tf([1], [1,2,1]))
+      plt.clf()
+      ctrl.bode_plot(ctrl.tf([5], [1, 1]))
+
+      # Check to make sure there are two axes and that each axes has one line
+      assert len(plt.gcf().axes) == 2
+      for ax in plt.gcf().axes:
+         # Make sure there is only 1 line in the subplot
+         assert len(ax.get_lines()) == 1
+
+      # Now add a line to the magnitude plot and make sure if is there
+      for ax in plt.gcf().axes:
+         if ax.get_label() == 'control-bode-magnitude':
+            break
+      ax.semilogx([1e-2, 1e1], 20 * np.log10([1, 1]), 'k-')
+      assert len(ax.get_lines()) == 2
+
    def test_doubleint(self):
       # 30 May 2016, RMM: added to replicate typecast bug in freqresp.py
       A = np.matrix('0, 1; 0, 0');
