Add pyplot.fig_subplot, for easier creation of figures with multiple subplots.

fperez · fperez · commit 9c60f86b4d7f · 2010-02-24T18:25:49.000Z
svn path=/trunk/matplotlib/; revision=8151
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,7 @@
+2010-02-24 Added pyplot.fig_subplot(), to create a figure and a group of
+	   subplots in a single call.  This offers an easier pattern than
+	   manually making figures and calling add_subplot() multiple times. FP
+	   
 2010-02-17 Added Gokhan's and Mattias' customizable keybindings patch
            for the toolbar.  You can now set the keymap.* properties
            in the matplotlibrc file.  Newbindings were added for
diff --git a/examples/pylab_examples/fig_subplot_demo.py b/examples/pylab_examples/fig_subplot_demo.py
@@ -0,0 +1,42 @@
+"""
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+
+x = np.linspace(0, 2*np.pi, 400)
+y = np.sin(x**2)
+
+plt.close('all')
+
+# Just a figure and one subplot
+f, ax = plt.fig_subplot()
+ax.plot(x, y)
+ax.set_title('Simple plot')
+
+# Two subplots, grab the whole fig_axes list
+fax = plt.fig_subplot(2, sharex=True)
+fax[1].plot(x, y)
+fax[1].set_title('Sharing X axis')
+fax[2].scatter(x, y)
+
+# Two subplots, unpack the output immediately
+f, ax1, ax2 = plt.fig_subplot(1, 2, sharey=True)
+ax1.plot(x, y)
+ax1.set_title('Sharing Y axis')
+ax2.scatter(x, y)
+
+# Three subplots sharing both x/y axes
+f, ax1, ax2, ax3 = plt.fig_subplot(3, sharex=True, sharey=True)
+ax1.plot(x, y)
+ax1.set_title('Sharing both axes')
+ax2.scatter(x, y)
+ax3.scatter(x, 2*y**2-1,color='r')
+# Fine-tune figure; make subplots close to each other and hide x ticks for
+# all but bottom plot.
+f.subplots_adjust(hspace=0)
+plt.setp([a.get_xticklabels() for a in f.axes[:-1]], visible=False)
+
+# Four polar axes
+plt.fig_subplot(2, 2, subplot_kw=dict(polar=True))
+
+plt.show()
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -649,6 +649,80 @@ def subplot(*args, **kwargs):
     return a
 
 
+def fig_subplot(nrows=1, ncols=1, sharex=False, sharey=False,
+                subplot_kw=None, **fig_kw):
+    """Create a figure with a set of subplots already made.
+
+    This utility wrapper makes it convenient to create common layouts of
+    subplots, including the enclosing figure object, in a single call.
+
+    Keyword arguments:
+    
+    nrows : int
+      Number of rows of the subplot grid.  Defaults to 1.
+
+    nrows : int
+      Number of columns of the subplot grid.  Defaults to 1.
+
+    sharex : bool
+      If True, the X axis will be shared amongst all subplots.
+
+    sharex : bool
+      If True, the Y axis will be shared amongst all subplots.
+
+    subplot_kw : dict
+      Dict with keywords passed to the add_subplot() call used to create each
+      subplots.
+
+    fig_kw : dict
+      Dict with keywords passed to the figure() call.  Note that all keywords
+      not recognized above will be automatically included here.
+
+    Returns:
+
+    fig_axes : list    
+      A list containing [fig, ax1, ax2, ...], where fig is the Matplotlib
+      Figure object and the rest are the axes.
+
+    **Examples:**
+
+    x = np.linspace(0, 2*np.pi, 400)
+    y = np.sin(x**2)
+
+    # Just a figure and one subplot
+    f, ax = plt.fig_subplot()
+    ax.plot(x, y)
+    ax.set_title('Simple plot')
+    
+    # Two subplots, unpack the output immediately
+    f, ax1, ax2 = plt.fig_subplot(1, 2, sharey=True)
+    ax1.plot(x, y)
+    ax1.set_title('Sharing Y axis')
+    ax2.scatter(x, y)
+
+    # Four polar axes
+    plt.fig_subplot(2, 2, subplot_kw=dict(polar=True))
+    """
+
+    if subplot_kw is None:
+        subplot_kw = {}
+        
+    fig = figure(**fig_kw)
+
+    # Create first subplot separately, so we can share it if requested
+    ax1 = fig.add_subplot(nrows, ncols, 1, **subplot_kw)
+    if sharex:
+        subplot_kw['sharex'] = ax1
+    if sharey:
+        subplot_kw['sharey'] = ax1
+
+    # Valid indices for axes start at 1, since fig is at 0: 
+    axes = [ fig.add_subplot(nrows, ncols, i, **subplot_kw)
+             for i in range(2, nrows*ncols+1)]
+
+    return [fig, ax1] + axes
+
+
 def twinx(ax=None):
     """
     Make a second axes overlay *ax* (or the current axes if *ax* is
