ENH/WIP: implement the simplest version of subplots from layout

tacaswell · tacaswell · commit 63e5ba3942f7 · 2020-02-28T03:05:58.000-05:00
This implements an idea @anntzer and @story645 had in the simplest possible way: trade a 2D array of strings for a dictionary of new Axes.
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -18,6 +18,7 @@
 from matplotlib import rcParams
 from matplotlib import docstring, projections
 from matplotlib import __version__ as _mpl_version
+import matplotlib.gridspec as gridspec
 
 import matplotlib.artist as martist
 from matplotlib.artist import Artist, allow_rasterization
@@ -1576,6 +1577,66 @@ def subplots(self, nrows=1, ncols=1, sharex=False, sharey=False,
             # Returned axis array will be always 2-d, even if nrows=ncols=1.
             return axarr
 
+    def build_grid(self, layout, subplot_kw=None, gridspec_kw=None):
+        """
+        Builds a layout based on an grid of strings
+
+        This is a helper function to build complex GridSpec layouts visually
+
+        Parameters
+        ----------
+        layout : list of list of str
+            A visual layout of how you want your Axes to be arranged labeled as strings.
+            For example ::
+
+               x = [['A', 'A', 'B'],
+                    ['C', 'D', 'B']]
+
+            Produces 4 axes:
+
+               - 'A' which is 1 row high and spans the first two columns
+               - 'B' which is 2 rows high and is on the right edge
+               - 'C' which in 1 row and 1 column wide in the bottom left
+               - 'D' which in 1 row and 1 column wide in the bottom center
+
+        subplot_kw : dict, optional
+            Dict with keywords passed to the
+            :meth:`~matplotlib.figure.Figure.add_subplot` call used to create
+            each subplot.
+
+        gridspec_kw : dict, optional
+            Dict with keywords passed to the
+            `~matplotlib.gridspec.GridSpec` constructor used to create
+            the grid the subplots are placed on.
+
+        Returns
+        -------
+        dict[str, Axes]
+           A dictionary mapping the string labels to the new Axes objects.
+
+        """
+        subplot_kw = subplot_kw or {}
+        gridspec_kw = gridspec_kw or {}
+        layout = np.asarray(layout)
+        rows, cols = layout.shape
+        unique_ids = set(layout.ravel())
+
+        gs = gridspec.GridSpec(rows, cols, figure=self, **gridspec_kw)
+
+        covered = np.zeros((rows, cols), dtype=bool)
+        output = dict()
+        for name in unique_ids:
+            indx = np.stack(np.where(layout == name)).T
+            start_row, start_col = np.min(indx, axis=0)
+            end_row, end_col = np.max(indx, axis=0) + 1
+            slc = (slice(start_row, end_row), slice(start_col, end_col))
+            if not np.all(covered[slc] == False):
+                raise ValueError
+            covered[slc] = True
+            output[name] = self.add_subplot(gs[slc], **subplot_kw)
+
+        return output
+
     def delaxes(self, ax):
         """
         Remove the `~.axes.Axes` *ax* from the figure; update the current axes.
diff --git a/lib/matplotlib/tests/test_build_grid.py b/lib/matplotlib/tests/test_build_grid.py
@@ -0,0 +1,24 @@
+from matplotlib.testing.decorators import check_figures_equal
+import matplotlib.gridspec as gridspec
+
+
+@check_figures_equal()
+def test_label_loc_vertical(fig_test, fig_ref):
+    x = [["A", "A", "B"], ["C", "D", "B"]]
+    grid_axes = fig_test.build_grid(x)
+
+    for k, ax in grid_axes.items():
+        ax.set_title(k)
+
+    gs = gridspec.GridSpec(2, 3, figure=fig_ref)
+    axA = fig_ref.add_subplot(gs[:1, :2])
+    axA.set_title("A")
+
+    axB = fig_ref.add_subplot(gs[:, 2])
+    axB.set_title("B")
+
+    axC = fig_ref.add_subplot(gs[1, 0])
+    axC.set_title("C")
+
+    axD = fig_ref.add_subplot(gs[1, 1])
+    axD.set_title("D")
