Merge pull request matplotlib#914 from braingram/master

more flexible axis sharing for pyplot.subplots

Author: mdboom

examples/pylab_examples/subplots_demo.py

@@ -10,8 +10,8 @@
 import numpy as np
 
 # Simple data to display in various forms
-x = np.linspace(0, 2*np.pi, 400)
-y = np.sin(x**2)
+x = np.linspace(0, 2 * np.pi, 400)
+y = np.sin(x ** 2)
 
 plt.close('all')
 
@@ -37,25 +37,33 @@
 ax1.plot(x, y)
 ax1.set_title('Sharing both axes')
 ax2.scatter(x, y)
-ax3.scatter(x, 2*y**2-1,color='r')
+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)
 
+# row and column sharing
+f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex='col', sharey='row')
+ax1.plot(x, y)
+ax1.set_title('Sharing x per column, y per row')
+ax2.scatter(x, y)
+ax3.scatter(x, 2 * y ** 2 - 1, color='r')
+ax4.plot(x, 2 * y ** 2 - 1, color='r')
+
 # Four axes, returned as a 2-d array
 f, axarr = plt.subplots(2, 2)
-axarr[0,0].plot(x, y)
-axarr[0,0].set_title('Axis [0,0]')
-axarr[0,1].scatter(x, y)
-axarr[0,1].set_title('Axis [0,1]')
-axarr[1,0].plot(x, y**2)
-axarr[1,0].set_title('Axis [1,0]')
-axarr[1,1].scatter(x, y**2)
-axarr[1,1].set_title('Axis [1,1]')
+axarr[0, 0].plot(x, y)
+axarr[0, 0].set_title('Axis [0,0]')
+axarr[0, 1].scatter(x, y)
+axarr[0, 1].set_title('Axis [0,1]')
+axarr[1, 0].plot(x, y ** 2)
+axarr[1, 0].set_title('Axis [1,0]')
+axarr[1, 1].scatter(x, y ** 2)
+axarr[1, 1].set_title('Axis [1,1]')
 # Fine-tune figure; hide x ticks for top plots and y ticks for right plots
-plt.setp([a.get_xticklabels() for a in axarr[0,:]], visible=False)
-plt.setp([a.get_yticklabels() for a in axarr[:,1]], visible=False)
+plt.setp([a.get_xticklabels() for a in axarr[0, :]], visible=False)
+plt.setp([a.get_yticklabels() for a in axarr[:, 1]], visible=False)
 
 # Four polar axes
 plt.subplots(2, 2, subplot_kw=dict(polar=True))

lib/matplotlib/pyplot.py

@@ -793,15 +793,27 @@ def subplots(nrows=1, ncols=1, sharex=False, sharey=False, squeeze=True,
       *ncols* : int
         Number of columns of the subplot grid.  Defaults to 1.
 
-      *sharex* : bool
+      *sharex* : string or bool
         If *True*, the X axis will be shared amongst all subplots.  If
         *True* and you have multiple rows, the x tick labels on all but
         the last row of plots will have visible set to *False*
-
-      *sharey* : bool
+        If a string must be one of "row", "col", "all", or "none".
+        "all" has the same effect as *True*, "none" has the same effect
+        as *False*.
+        If "row", each subplot row will share a X axis.
+        If "col", each subplot column will share a X axis and the x tick
+        labels on all but the last row will have visible set to *False*.
+
+      *sharey* : string or bool
         If *True*, the Y axis will be shared amongst all subplots. If
         *True* and you have multiple columns, the y tick labels on all but
         the first column of plots will have visible set to *False*
+        If a string must be one of "row", "col", "all", or "none".
+        "all" has the same effect as *True*, "none" has the same effect
+        as *False*.
+        If "row", each subplot row will share a Y axis.
+        If "col", each subplot column will share a Y axis and the y tick
+        labels on all but the last row will have visible set to *False*.
 
       *squeeze* : bool
         If *True*, extra dimensions are squeezed out from the
@@ -859,7 +871,36 @@ def subplots(nrows=1, ncols=1, sharex=False, sharey=False, squeeze=True,
 
         # Four polar axes
         plt.subplots(2, 2, subplot_kw=dict(polar=True))
+
+        # Share a X axis with each column of subplots
+        plt.subplots(2, 2, sharex='col')
+
+        # Share a Y axis with each row of subplots
+        plt.subplots(2, 2, sharey='row')
+
+        # Share a X and Y axis with all subplots
+        plt.subplots(2, 2, sharex='all', sharey='all')
+        # same as
+        plt.subplots(2, 2, sharex=True, sharey=True)
     """
+    # for backwards compatability
+    if isinstance(sharex, bool):
+        if sharex:
+            sharex = "all"
+        else:
+            sharex = "none"
+    if isinstance(sharey, bool):
+        if sharey:
+            sharey = "all"
+        else:
+            sharey = "none"
+    share_values = ["all", "row", "col", "none"]
+    if sharex not in share_values:
+        raise ValueError("sharex [%s] must be one of %s" % \
+                (sharex, share_values))
+    if sharey not in share_values:
+        raise ValueError("sharey [%s] must be one of %s" % \
+                (sharey, share_values))
 
     if subplot_kw is None:
         subplot_kw = {}
@@ -873,34 +914,52 @@ def subplots(nrows=1, ncols=1, sharex=False, sharey=False, squeeze=True,
 
     # 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
+    #if sharex:
+    #    subplot_kw['sharex'] = ax0
+    #if sharey:
+    #    subplot_kw['sharey'] = ax0
     axarr[0] = ax0
 
+    r, c = np.mgrid[:nrows, :ncols]
+    r = r.flatten() * ncols
+    c = c.flatten()
+    lookup = {
+            "none": np.arange(nplots),
+            "all": np.zeros(nplots, dtype=int),
+            "row": r,
+            "col": c,
+            }
+    sxs = lookup[sharex]
+    sys = lookup[sharey]
+
     # 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 sxs[i] == i:
+            subplot_kw['sharex'] = None
+        else:
+            subplot_kw['sharex'] = axarr[sxs[i]]
+        if sys[i] == i:
+            subplot_kw['sharey'] = None
+        else:
+            subplot_kw['sharey'] = axarr[sys[i]]
+        axarr[i] = fig.add_subplot(nrows, ncols, i + 1, **subplot_kw)
 
     # returned axis array will be always 2-d, even if nrows=ncols=1
     axarr = axarr.reshape(nrows, ncols)
 
-
     # turn off redundant tick labeling
-    if sharex and nrows>1:
+    if sharex in ["col", "all"] and nrows > 1:
+    #if sharex and nrows>1:
         # turn off all but the bottom row
-        for ax in axarr[:-1,:].flat:
+        for ax in axarr[:-1, :].flat:
             for label in ax.get_xticklabels():
                 label.set_visible(False)
 
-
-    if sharey and ncols>1:
+    if sharey in ["row", "all"] and ncols > 1:
+    #if sharey and ncols>1:
         # turn off all but the first column
-        for ax in axarr[:,1:].flat:
+        for ax in axarr[:, 1:].flat:
             for label in ax.get_yticklabels():
                 label.set_visible(False)
 

lib/matplotlib/tests/test_subplots.py

@@ -0,0 +1,110 @@
+import numpy
+import matplotlib.pyplot as plt
+
+from nose.tools import assert_raises
+
+
+def check_shared(results, f, axs):
+    """
+    results is a 4 x 4 x 2 matrix of boolean values where
+        if [i, j, 0] == True, X axis for subplots i and j should be shared
+        if [i, j, 1] == False, Y axis for subplots i and j should not be shared
+    """
+    shared_str = ['x', 'y']
+    shared = [axs[0]._shared_x_axes, axs[0]._shared_y_axes]
+    #shared = {
+    #        'x': a1._shared_x_axes,
+    #        'y': a1._shared_y_axes,
+    #        }
+    tostr = lambda r: "not " if r else ""
+    for i1 in xrange(len(axs)):
+        for i2 in xrange(i1 + 1, len(axs)):
+            for i3 in xrange(len(shared)):
+                assert shared[i3].joined(axs[i1], axs[i2]) == \
+                        results[i1, i2, i3], \
+                        "axes %i and %i incorrectly %ssharing %s axis" % \
+                        (i1, i2, tostr(results[i1, i2, i3]), shared_str[i3])
+
+
+def check_visible(result, f, axs):
+    tostr = lambda v: "invisible" if v else "visible"
+    for (ax, vx, vy) in zip(axs, result['x'], result['y']):
+        for l in ax.get_xticklabels():
+            assert l.get_visible() == vx, \
+                    "X axis was incorrectly %s" % (tostr(vx))
+        for l in ax.get_yticklabels():
+            assert l.get_visible() == vy, \
+                    "Y axis was incorrectly %s" % (tostr(vy))
+
+
+def test_shared():
+    rdim = (4, 4, 2)
+    share = {
+            'all': numpy.ones(rdim[:2], dtype=bool),
+            'none': numpy.zeros(rdim[:2], dtype=bool),
+            'row': numpy.array([
+                [False, True, False, False],
+                [True, False, False, False],
+                [False, False, False, True],
+                [False, False, True, False]]),
+            'col': numpy.array([
+                [False, False, True, False],
+                [False, False, False, True],
+                [True, False, False, False],
+                [False, True, False, False]]),
+            }
+    visible = {
+            'x': {
+                'all': [False, False, True, True],
+                'col': [False, False, True, True],
+                'row': [True] * 4,
+                'none': [True] * 4,
+                False: [True] * 4,
+                True: [False, False, True, True],
+                },
+            'y': {
+                'all': [True, False, True, False],
+                'col': [True] * 4,
+                'row': [True, False, True, False],
+                'none': [True] * 4,
+                False: [True] * 4,
+                True: [True, False, True, False],
+                },
+            }
+    share[False] = share['none']
+    share[True] = share['all']
+
+    # test default
+    f, ((a1, a2), (a3, a4)) = plt.subplots(2, 2)
+    axs = [a1, a2, a3, a4]
+    check_shared(numpy.dstack((share['none'], share['none'])), \
+            f, axs)
+    plt.close(f)
+
+    # test all option combinations
+    ops = [False, True, 'all', 'none', 'row', 'col']
+    for xo in ops:
+        for yo in ops:
+            f, ((a1, a2), (a3, a4)) = plt.subplots(2, 2, sharex=xo, sharey=yo)
+            axs = [a1, a2, a3, a4]
+            check_shared(numpy.dstack((share[xo], share[yo])), \
+                    f, axs)
+            check_visible(dict(x=visible['x'][xo], y=visible['y'][yo]), \
+                    f, axs)
+            plt.close(f)
+
+
+def test_exceptions():
+    # TODO should this test more options?
+    with assert_raises(ValueError):
+        plt.subplots(2, 2, sharex='blah')
+        plt.subplots(2, 2, sharey='blah')
+
+
+def test_subplots():
+    # things to test
+    # - are axes actually shared?
+    # - are tickmarks correctly hidden?
+    test_shared()
+    # - are exceptions thrown correctly
+    test_exceptions()