Function wrapper for examples/api/two_scales.py

examples/api/two_scales.py

@@ -13,28 +13,68 @@
 Such axes are generated by calling the `Axes.twinx` method.  Likewise,
 `Axes.twiny` is available to generate axes that share a *y* axis but
 have different top and bottom scales.
-
-The twinx and twiny methods are also exposed as pyplot functions.
-
 """
-
 import numpy as np
 import matplotlib.pyplot as plt
 
-fig, ax1 = plt.subplots()
+
+def two_scales(ax1, time, data1, data2, c1, c2):
+    """
+
+    Parameters
+    ----------
+    ax : axis
+        Axis to put two scales on
+
+    time : array-like
+        x-axis values for both datasets
+
+    data1: array-like
+        Data for left hand scale
+
+    data2 : array-like
+        Data for right hand scale
+
+    c1 : color
+        Color for line 1
+
+    c2 : color
+        Color for line 2
+
+    Returns
+    -------
+    ax : axis
+        Original axis
+    ax2 : axis
+        New twin axis
+    """
+    ax2 = ax1.twinx()
+
+    ax1.plot(time, data1, color=c1)
+    ax1.set_xlabel('time (s)')
+    ax1.set_ylabel('exp')
+
+    ax2.plot(time, data2, color=c2)
+    ax2.set_ylabel('sin')
+    return ax1, ax2
+
+
+# Create some mock data
 t = np.arange(0.01, 10.0, 0.01)
 s1 = np.exp(t)
-ax1.plot(t, s1, 'b-')
-ax1.set_xlabel('time (s)')
-# Make the y-axis label, ticks and tick labels match the line color.
-ax1.set_ylabel('exp', color='b')
-ax1.tick_params('y', colors='b')
-
-ax2 = ax1.twinx()
 s2 = np.sin(2 * np.pi * t)
-ax2.plot(t, s2, 'r.')
-ax2.set_ylabel('sin', color='r')
-ax2.tick_params('y', colors='r')
 
-fig.tight_layout()
+# Create axes
+fig, ax = plt.subplots()
+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()