Cleanup and simplify two_scales example

dstansby · dstansby · commit 213403fd49b1 · 2017-01-14T21:45:12.000Z
diff --git a/examples/api/two_scales.py b/examples/api/two_scales.py
@@ -14,63 +14,67 @@
 `Axes.twiny` is available to generate axes that share a *y* axis but
 have different top and bottom scales.
 """
-
 import numpy as np
 import matplotlib.pyplot as plt
 
-def two_scales(ax1, ax2, time, data1, data2, param1_dic , param2_dic):
+
+def two_scales(ax1, time, data1, data2, c1, c2):
     """
 
     Parameters
     ----------
-    ax : (type of axis)
-        A description of axis
+    ax : axis
+        Axis to put two scales on
 
-    data1: (first dataset)
-        A description of data1
+    time : array-like
+        x-axis values for both datasets
 
-    data2 : (first dataset)
-        A description of data2
+    data1: array-like
+        Data for left hand scale
 
-    param_dic : This is a dictionary of the parameters of the style and color e.g. {line style: '-', text color = 'r'}
-    Returns
-    -------
-    Overlays
-    data1 : (Plot first data set)
-    data2 : (Plot second data set)
+    data2 : array-like
+        Data for right hand scale
+
+    c1 : color
+        Color for line 1
 
-     """
-    def color_y_axes(ax, color):
-        """Color your axes."""
-        for t in ax.get_yticklabels():
-            t.set_color(color)
-        return None
+    c2 : color
+        Color for line 2
 
+    Returns
+    -------
+    ax : axis
+        Original axis
+    ax2 : axis
+        New twin axis
+    """
+    ax2 = ax1.twinx()
 
-    ax1.plot(time, data1, param1_dic['color'] + param1_dic['style'])
+    ax1.plot(time, data1, color=c1)
     ax1.set_xlabel('time (s)')
-    # Make the y-axis label and tick labels match the line color.
-    ax1.set_ylabel('exp', color=param1_dic['color'])
-    color_y_axes(ax1, param1_dic['color'])
+    ax1.set_ylabel('exp')
 
+    ax2.plot(time, data2, color=c2)
+    ax2.set_ylabel('sin')
+    return ax1, ax2
 
-    ax2.plot(time, data2, param2_dic['color'] + param2_dic['style'])
-    ax2.set_ylabel('sin', color=param2_dic['color'])
-    color_y_axes(ax2, param2_dic['color'])
-    return plt.show()
 
-#Create some mock data
+# Create some mock data
 t = np.arange(0.01, 10.0, 0.01)
 s1 = np.exp(t)
-s2 = np.sin(2*np.pi*t)
+s2 = np.sin(2 * np.pi * t)
 
-#Specify your parameter dictionary
-d1 = {'style': '-', 'color' : 'r'}
-d2 = {'style': '.', 'color' :'b'}
-
-#create your axes
+# Create axes
 fig, ax = plt.subplots()
-ax2 = ax.twinx()
-
-#Call the function
-two_scales(ax, ax2, t, s1, s2, d1, d2)
+ax1, ax2 = two_scales(ax, t, s1, s2, 'r', 'b')
+
+
+# Change color of each axis
+def color_y_axis(ax, color):
+    """Color your axes."""
+    for t in ax.get_yticklabels():
+        t.set_color(color)
+    return None
+color_y_axis(ax1, 'r')
+color_y_axis(ax2, 'b')
+plt.show()
