Update bode_plot and gangof4_plot axes label handling

murrayrm · murrayrm · commit eabe562b8e0b · 2018-01-20T10:11:09.000-08:00
Changed the labels for subplots to be of the form `control-plotname-subfigname` for consistency. Moved around some of the code from @roryyorke to put it in a single place. Added documentation.
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -135,20 +135,6 @@ def bode_plot(syslist, omega=None, dB=None, Hz=None, deg=None,
             else:
                 omega = sp.logspace(np.log10(omega_limits[0]), np.log10(omega_limits[1]), endpoint=True)
 
-    if Plot:
-        fig = plt.gcf()
-        ax_mag = None
-        ax_phase = None
-        for ax in fig.axes:
-            if ax.get_label() == 'pycontrol-bode-mag':
-                ax_mag = ax
-            elif ax.get_label() == 'pycontrol-bode-phs':
-                ax_phase = ax
-        if ax_mag is None or ax_phase is None:
-            plt.clf()
-            ax_mag = plt.subplot(211, label = 'pycontrol-bode-mag')
-            ax_phase = plt.subplot(212, label = 'pycontrol-bode-phs', sharex=ax_mag)
-
     mags, phases, omegas, nyquistfrqs = [], [], [], []
     for sys in syslist:
         if (sys.inputs > 1 or sys.outputs > 1):
@@ -185,10 +171,30 @@ def bode_plot(syslist, omega=None, dB=None, Hz=None, deg=None,
             #! TODO: Not current implemented; just use subplot for now
 
             if (Plot):
-                # Create a unique label to fix bug in matplotlib<=2.1
-                # See https://github.com/matplotlib/matplotlib/issues/9024
-                import random
-                figlabel = str(random.randint(1, 1e6))
+                # Set up the axes with labels so that multiple calls to
+                # bode_plot will superimpose the data.  This was implicit
+                # before matplotlib 2.1, but changed after that (See
+                # https://github.com/matplotlib/matplotlib/issues/9024). 
+                # The code below should work on all cases.
+
+                # Get the current figure 
+                fig = plt.gcf()
+                ax_mag = None
+                ax_phase = None
+
+                # Get the current axes if they already exist
+                for ax in fig.axes:
+                    if ax.get_label() == 'control-bode-magnitude':
+                        ax_mag = ax
+                    elif ax.get_label() == 'control-bode-phase':
+                        ax_phase = ax
+
+                # If no axes present, create them from scratch
+                if ax_mag is None or ax_phase is None:
+                    plt.clf()
+                    ax_mag = plt.subplot(211, label = 'control-bode-magnitude')
+                    ax_phase = plt.subplot(212, label = 'control-bode-phase',
+                                           sharex=ax_mag)
 
                 # Magnitude plot
                 if dB:
@@ -372,25 +378,29 @@ def gangof4_plot(P, C, omega=None):
         L = P * C;
         S = feedback(1, L);
         T = L * S;
-
+        
+        # Set up the axes with labels so that multiple calls to
+        # gangof4_plot will superimpose the data.  See details in bode_plot.
         plot_axes = {'t' : None, 's' : None, 'ps' : None, 'cs' : None}
         for ax in plt.gcf().axes:
             label = ax.get_label()
-            if label.startswith('pycontrol-gof-'):
-                key = label[len('pycontrol-gof-'):]
+            if label.startswith('control-gangof4-'):
+                key = label[len('control-gangof4-'):]
                 if key not in plot_axes:
-                    raise RuntimeError("unknown gof axis type '{}'".format(label))
+                    raise RuntimeError("unknown gangof4 axis type '{}'".format(label))
                 plot_axes[key] = ax
 
-        # if any are missing, start from scratch
+        # if any of the axes are missing, start from scratch
         if any((ax is None for ax in plot_axes.values())):
             plt.clf()
-            plot_axes = {'t' : plt.subplot(221,label='pycontrol-gof-t'),
-                         'ps' : plt.subplot(222,label='pycontrol-gof-ps'),
-                         'cs' : plt.subplot(223,label='pycontrol-gof-cs'),
-                         's' : plt.subplot(224,label='pycontrol-gof-s')}
+            plot_axes = {'t' : plt.subplot(221,label='control-gangof4-t'),
+                         'ps' : plt.subplot(222,label='control-gangof4-ps'),
+                         'cs' : plt.subplot(223,label='control-gangof4-cs'),
+                         's' : plt.subplot(224,label='control-gangof4-s')}
 
+        #
         # Plot the four sensitivity functions
+        #
 
         #! TODO: Need to add in the mag = 1 lines
         mag_tmp, phase_tmp, omega = T.freqresp(omega);
diff --git a/doc/control.rst b/doc/control.rst
@@ -44,6 +44,15 @@ Frequency domain plotting
     gangof4_plot
     nichols_plot
 
+Note: For plotting commands that create multiple axes on the same plot, the
+individual axes can be retrieved using the axes label (retrieved using the
+`get_label` method for the matplotliib axes object).  The following labels
+are currently defined:
+
+* Bode plots: `control-bode-magnitude`, `control-bode-phase`
+* Gang of 4 plots: `control-gangof4-s`, `control-gangof4-cs`,
+  `control-gangof4-ps`, `control-gangof4-t`
+
 Time domain simulation
 ======================
 
