TST: copy matplotlib subplots function for compat with mpl < 1.0

Author: wesm

pandas/core/frame.py

@@ -3571,11 +3571,12 @@ def plot(self, subplots=False, sharex=True, sharey=False, use_index=True,
         and will error.
         """
         import matplotlib.pyplot as plt
+        import pandas.tools.plotting as gfx
 
         if subplots:
-            fig, axes = plt.subplots(nrows=len(self.columns),
-                                   sharex=sharex, sharey=sharey,
-                                   figsize=figsize)
+            fig, axes = gfx.subplots(nrows=len(self.columns),
+                                     sharex=sharex, sharey=sharey,
+                                     figsize=figsize)
         else:
             if ax is None:
                 fig = plt.figure(figsize=figsize)
@@ -3676,13 +3677,14 @@ def hist(self, grid=True, **kwds):
         kwds : other plotting keyword arguments
             To be passed to hist function
         """
+        import pandas.tools.plotting as gfx
         import matplotlib.pyplot as plt
 
         n = len(self.columns)
         k = 1
         while k ** 2 < n:
             k += 1
-        _, axes = plt.subplots(nrows=k, ncols=k)
+        _, axes = gfx.subplots(nrows=k, ncols=k)
 
         for i, col in enumerate(_try_sort(self.columns)):
             ax = axes[i / k][i % k]

pandas/tools/plotting.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 def scatter_matrix(data):
     pass
 
@@ -123,14 +125,12 @@ def plot_group(group, ax):
 
 def _grouped_plot(plotf, data, by=None, numeric_only=True, figsize=(10, 5),
                   sharex=True, sharey=True):
-    import matplotlib.pyplot as plt
-
     grouped = data.groupby(by)
     ngroups = len(grouped)
 
     nrows, ncols = _get_layout(ngroups)
-    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize,
-                             sharex=sharex, sharey=sharey)
+    fig, axes = subplots(nrows=nrows, ncols=ncols, figsize=figsize,
+                         sharex=sharex, sharey=sharey)
 
     ravel_axes = []
     for row in axes:
@@ -155,8 +155,8 @@ def _grouped_plot_by_column(plotf, data, columns=None, by=None,
     ngroups = len(columns)
 
     nrows, ncols = _get_layout(ngroups)
-    fig, axes = plt.subplots(nrows=nrows, ncols=ncols,
-                             sharex=True, sharey=True)
+    fig, axes = subplots(nrows=nrows, ncols=ncols,
+                         sharex=True, sharey=True)
 
     if isinstance(axes, plt.Axes):
         ravel_axes = [axes]
@@ -198,6 +198,113 @@ def _get_layout(nplots):
     else:
         return k, k
 
+# copied from matplotlib/pyplot.py for compatibility with matplotlib < 1.0
+
+def subplots(nrows=1, ncols=1, sharex=False, sharey=False, squeeze=True,
+              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.
+
+    ncols : 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.
+
+    squeeze : bool
+
+      If True, extra dimensions are squeezed out from the returned axis object:
+        - if only one subplot is constructed (nrows=ncols=1), the resulting
+        single Axis object is returned as a scalar.
+        - for Nx1 or 1xN subplots, the returned object is a 1-d numpy object
+        array of Axis objects are returned as numpy 1-d arrays.
+        - for NxM subplots with N>1 and M>1 are returned as a 2d array.
+
+      If False, no squeezing at all is done: the returned axis object is always
+      a 2-d array contaning Axis instances, even if it ends up being 1x1.
+
+    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, ax : tuple
+      - fig is the Matplotlib Figure object
+      - ax can be either a single axis object or an array of axis objects if
+      more than one supblot was created.  The dimensions of the resulting array
+      can be controlled with the squeeze keyword, see above.
+
+    **Examples:**
+
+    x = np.linspace(0, 2*np.pi, 400)
+    y = np.sin(x**2)
+
+    # Just a figure and one subplot
+    f, ax = plt.subplots()
+    ax.plot(x, y)
+    ax.set_title('Simple plot')
+
+    # Two subplots, unpack the output array immediately
+    f, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
+    ax1.plot(x, y)
+    ax1.set_title('Sharing Y axis')
+    ax2.scatter(x, y)
+
+    # Four polar axes
+    plt.subplots(2, 2, subplot_kw=dict(polar=True))
+    """
+    import matplotlib.pyplot as plt
+
+    if subplot_kw is None:
+        subplot_kw = {}
+
+    fig = plt.figure(**fig_kw)
+
+    # Create empty object array to hold all axes.  It's easiest to make it 1-d
+    # so we can just append subplots upon creation, and then
+    nplots = nrows*ncols
+    axarr = np.empty(nplots, dtype=object)
+
+    # Create first subplot separately, so we can share it if requested
+    ax0 = fig.add_subplot(nrows, ncols, 1, **subplot_kw)
+    if sharex:
+        subplot_kw['sharex'] = ax0
+    if sharey:
+        subplot_kw['sharey'] = ax0
+    axarr[0] = ax0
+
+    # Note off-by-one counting because add_subplot uses the MATLAB 1-based
+    # convention.
+    for i in range(1, nplots):
+        axarr[i] = fig.add_subplot(nrows, ncols, i+1, **subplot_kw)
+
+    if squeeze:
+        # Reshape the array to have the final desired dimension (nrow,ncol),
+        # though discarding unneeded dimensions that equal 1.  If we only have
+        # one subplot, just return it instead of a 1-element array.
+        if nplots==1:
+            return fig, axarr[0]
+        else:
+            return fig, axarr.reshape(nrows, ncols).squeeze()
+    else:
+        # returned axis array will be always 2-d, even if nrows=ncols=1
+        return fig, axarr.reshape(nrows, ncols)
+
 if __name__ == '__main__':
     import pandas.rpy.common as com
     sales = com.load_data('sanfrancisco.home.sales', package='nutshell')