ENH: add ability to pass per-Axes subplot_kw through subplot_mosaic

tacaswell · tacaswell · commit f74fe35df8c9 · 2022-12-21T15:26:35.000-05:00
diff --git a/doc/users/next_whats_new/per_subplot_mosaic.rst b/doc/users/next_whats_new/per_subplot_mosaic.rst
@@ -0,0 +1,22 @@
+Per-subplot keyword arguments  in ``subplot_mosaic``
+----------------------------------------------------
+
+It is now possible to pass keyword arguments through to the creation to each
+specific call to ``add_subplot`` in `.Figure.subplot_mosaic` and
+`.pyplot.subplot_mosaic` :
+
+.. plot::
+   :include-source: true
+
+   fig, axd = plt.subplot_mosaic(
+       "AB;CD",
+       per_subplot_kw={
+           "A": {"projection": "polar"},
+           ("C", "D"): {"xscale": "log"},
+           "B": {"projection": "3d"},
+       },
+   )
+
+
+This is particularly useful for creating mosaics with mixed projections, but
+any keyword arguments can be passed through.
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -1759,6 +1759,25 @@ def get_tightbbox(self, renderer=None, bbox_extra_artists=None):
 
         return _bbox
 
+    @staticmethod
+    def _norm_per_subplot_kw(per_subplot_kw):
+        expanded = {}
+        for k, v in per_subplot_kw.items():
+            if isinstance(k, tuple):
+                for sub_key in k:
+                    if sub_key in expanded:
+                        raise ValueError(
+                            f'The key {sub_key!r} appears multiple times.'
+                            )
+                    expanded[sub_key] = v
+            else:
+                if k in expanded:
+                    raise ValueError(
+                        f'The key {k!r} appears multiple times.'
+                    )
+                expanded[k] = v
+        return expanded
+
     @staticmethod
     def _normalize_grid_string(layout):
         if '\n' not in layout:
@@ -1771,7 +1790,8 @@ def _normalize_grid_string(layout):
 
     def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False,
                        width_ratios=None, height_ratios=None,
-                       empty_sentinel='.', subplot_kw=None, gridspec_kw=None):
+                       empty_sentinel='.', subplot_kw=None, gridspec_kw=None,
+                       per_subplot_kw=None):
         """
         Build a layout of Axes based on ASCII art or nested lists.
 
@@ -1821,6 +1841,9 @@ def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False,
             The string notation allows only single character Axes labels and
             does not support nesting but is very terse.
 
+            Tuples may be used instead of lists and tuples may not be used
+            as Axes identifiers.
+
         sharex, sharey : bool, default: False
             If True, the x-axis (*sharex*) or y-axis (*sharey*) will be shared
             among all subplots.  In that case, tick label visibility and axis
@@ -1843,7 +1866,21 @@ def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False,
 
         subplot_kw : dict, optional
             Dictionary with keywords passed to the `.Figure.add_subplot` call
-            used to create each subplot.
+            used to create each subplot.  These values may be overridden by
+            values in *per_subplot_kw*.
+
+        per_subplot_kw : dict, optional
+            A dictionary mapping the Axes idenitfies or tuples of identifies to
+            a dictionary of keyword arguments to be passed to the
+            `.Figure.add_subplot` call used to create each subplot.  The values
+            in these dictionaries have precedence over the values in
+            *subplot_kw*.
+
+            In the special case *mosaic* being a string, multi-character keys
+            in *per_subplot_kw* will be applied to all of the Axes named
+            treating the string as a sequence.
+
+            .. versionadded:: 3.7
 
         gridspec_kw : dict, optional
             Dictionary with keywords passed to the `.GridSpec` constructor used
@@ -1868,6 +1905,8 @@ def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False,
         """
         subplot_kw = subplot_kw or {}
         gridspec_kw = dict(gridspec_kw or {})
+        per_subplot_kw = per_subplot_kw or {}
+
         if height_ratios is not None:
             if 'height_ratios' in gridspec_kw:
                 raise ValueError("'height_ratios' must not be defined both as "
@@ -1882,6 +1921,14 @@ def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False,
         # special-case string input
         if isinstance(mosaic, str):
             mosaic = self._normalize_grid_string(mosaic)
+            per_subplot_kw = {
+                tuple(k): v for k, v in per_subplot_kw.items()
+            }
+
+        per_subplot_kw = self._norm_per_subplot_kw(
+            per_subplot_kw
+        )
+
         # Only accept strict bools to allow a possible future API expansion.
         _api.check_isinstance(bool, sharex=sharex, sharey=sharey)
 
@@ -2011,7 +2058,11 @@ def _do_layout(gs, mosaic, unique_ids, nested):
                         raise ValueError(f"There are duplicate keys {name} "
                                          f"in the layout\n{mosaic!r}")
                     ax = self.add_subplot(
-                        gs[slc], **{'label': str(name), **subplot_kw}
+                        gs[slc], **{
+                            'label': str(name),
+                            **subplot_kw,
+                            **per_subplot_kw.get(name, {})
+                        }
                     )
                     output[name] = ax
                 elif method == 'nested':
@@ -2048,7 +2099,11 @@ def _do_layout(gs, mosaic, unique_ids, nested):
             if sharey:
                 ax.sharey(ax0)
                 ax._label_outer_yaxis(check_patch=True)
-
+        if extra := set(per_subplot_kw) - set(ret):
+            raise ValueError(
+                f"The keys {extra} are in *per_subplot_kw* "
+                "but not in the mosaic."
+            )
         return ret
 
     def _set_artist_props(self, a):
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -1479,7 +1479,8 @@ def subplots(nrows=1, ncols=1, *, sharex=False, sharey=False, squeeze=True,
 
 def subplot_mosaic(mosaic, *, sharex=False, sharey=False,
                    width_ratios=None, height_ratios=None, empty_sentinel='.',
-                   subplot_kw=None, gridspec_kw=None, **fig_kw):
+                   subplot_kw=None, gridspec_kw=None,
+                   per_subplot_kw=None, **fig_kw):
     """
     Build a layout of Axes based on ASCII art or nested lists.
 
@@ -1576,7 +1577,8 @@ def subplot_mosaic(mosaic, *, sharex=False, sharey=False,
         mosaic, sharex=sharex, sharey=sharey,
         height_ratios=height_ratios, width_ratios=width_ratios,
         subplot_kw=subplot_kw, gridspec_kw=gridspec_kw,
-        empty_sentinel=empty_sentinel
+        empty_sentinel=empty_sentinel,
+        per_subplot_kw=per_subplot_kw,
     )
     return fig, ax_dict
 
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -848,7 +848,12 @@ def test_animated_with_canvas_change(fig_test, fig_ref):
 class TestSubplotMosaic:
     @check_figures_equal(extensions=["png"])
     @pytest.mark.parametrize(
-        "x", [[["A", "A", "B"], ["C", "D", "B"]], [[1, 1, 2], [3, 4, 2]]]
+        "x", [
+            [["A", "A", "B"], ["C", "D", "B"]],
+            [[1, 1, 2], [3, 4, 2]],
+            (("A", "A", "B"), ("C", "D", "B")),
+            ((1, 1, 2), (3, 4, 2))
+        ]
     )
     def test_basic(self, fig_test, fig_ref, x):
         grid_axes = fig_test.subplot_mosaic(x)
@@ -998,6 +1003,10 @@ def test_fail_list_of_str(self):
             plt.subplot_mosaic(['foo', 'bar'])
         with pytest.raises(ValueError, match='must be 2D'):
             plt.subplot_mosaic(['foo'])
+        with pytest.raises(ValueError, match='must be 2D'):
+            plt.subplot_mosaic([['foo', ('bar',)]])
+        with pytest.raises(ValueError, match='must be 2D'):
+            plt.subplot_mosaic([['a', 'b'], [('a', 'b'), 'c']])
 
     @check_figures_equal(extensions=["png"])
     @pytest.mark.parametrize("subplot_kw", [{}, {"projection": "polar"}, None])
@@ -1011,8 +1020,26 @@ def test_subplot_kw(self, fig_test, fig_ref, subplot_kw):
 
         axB = fig_ref.add_subplot(gs[0, 1], **subplot_kw)
 
+    @check_figures_equal(extensions=["png"])
+    @pytest.mark.parametrize("multi_value", ['BC', tuple('BC')])
+    def test_per_subplot_kw(self, fig_test, fig_ref, multi_value):
+        x = 'AB;CD'
+        grid_axes = fig_test.subplot_mosaic(
+            x,
+            subplot_kw={'facecolor': 'red'},
+            per_subplot_kw={
+                'D': {'facecolor': 'blue'},
+                multi_value: {'facecolor': 'green'},
+            }
+        )
+
+        gs = fig_ref.add_gridspec(2, 2)
+        for color, spec in zip(['red', 'green', 'green', 'blue'], gs):
+            fig_ref.add_subplot(spec, facecolor=color)
+
     def test_string_parser(self):
         normalize = Figure._normalize_grid_string
+
         assert normalize('ABC') == [['A', 'B', 'C']]
         assert normalize('AB;CC') == [['A', 'B'], ['C', 'C']]
         assert normalize('AB;CC;DE') == [['A', 'B'], ['C', 'C'], ['D', 'E']]
@@ -1029,6 +1056,25 @@ def test_string_parser(self):
                          DE
                          """) == [['A', 'B'], ['C', 'C'], ['D', 'E']]
 
+    def test_per_subplot_kw_expander(self):
+        normalize = Figure._norm_per_subplot_kw
+        assert normalize({"A": {}, "B": {}}) == {"A": {}, "B": {}}
+        assert normalize({("A", "B"): {}}) == {"A": {}, "B": {}}
+        with pytest.raises(
+                ValueError, match=f'The key {"B"!r} appears multiple times'
+        ):
+            normalize({("A", "B"): {}, "B": {}})
+        with pytest.raises(
+                ValueError, match=f'The key {"B"!r} appears multiple times'
+        ):
+            normalize({"B": {}, ("A", "B"): {}})
+
+    def test_extra_per_subplot_kw(self):
+        with pytest.raises(
+                ValueError, match=f'The keys {set("B")!r} are in'
+        ):
+            Figure().subplot_mosaic("A", per_subplot_kw={"B": {}})
+
     @check_figures_equal(extensions=["png"])
     @pytest.mark.parametrize("str_pattern",
                              ["AAA\nBBB", "\nAAA\nBBB\n", "ABC\nDEF"]
diff --git a/tutorials/provisional/mosaic.py b/tutorials/provisional/mosaic.py
@@ -202,8 +202,8 @@ def identify_axes(ax_dict, fontsize=48):
 # empty sentinel with the string shorthand because it may be stripped
 # while processing the input.
 #
-# Controlling mosaic and subplot creation
-# =======================================
+# Controlling mosaic creation
+# ===========================
 #
 # This feature is built on top of `.gridspec` and you can pass the
 # keyword arguments through to the underlying `.gridspec.GridSpec`
@@ -278,15 +278,72 @@ def identify_axes(ax_dict, fontsize=48):
 
 
 ###############################################################################
+# Controlling subplot creation
+# ============================
+
 # We can also pass through arguments used to create the subplots
-# (again, the same as `.Figure.subplots`).
+# (again, the same as `.Figure.subplots`) which will apply to all
+# of the Axes created.
 
 
 axd = plt.figure(constrained_layout=True).subplot_mosaic(
     "AB", subplot_kw={"projection": "polar"}
 )
 identify_axes(axd)
 
+###############################################################################
+# Per-Axes subplot keyword arguments
+# ----------------------------------
+#
+#
+# If you need to control the parameters passed to each subplot individually use
+# *per_subplot_kw* to pass a mapping between the Axes identifiers (or
+# tuples of Axes identifiers) to dictionaries of keywords to be passed.
+#
+# .. versionadded:: 3.7
+#
+
+
+fig, axd = plt.subplot_mosaic(
+    "AB;CD",
+    per_subplot_kw={
+        "A": {"projection": "polar"},
+        ("C", "D"): {"xscale": "log"}
+    },
+)
+identify_axes(axd)
+
+###############################################################################
+# If the layout is specified with the string short-hand, then we know the
+# Axes labels will be one character and can unambiguously interpret longer
+# strings in *per_subplot_kw* to specify a set of Axes to apply the
+# keywords to:
+
+
+fig, axd = plt.subplot_mosaic(
+    "AB;CD",
+    per_subplot_kw={
+        "AD": {"projection": "polar"},
+        "BC": {"facecolor": ".9"}
+    },
+)
+identify_axes(axd)
+
+###############################################################################
+# if *subplot_kw* and *per_subplot_kw* are used together, then they are
+# merged with *per_subplot_kw* taking priority:
+
+
+axd = plt.figure(constrained_layout=True).subplot_mosaic(
+    "AB;CD",
+    subplot_kw={"facecolor": "xkcd:tangerine"},
+    per_subplot_kw={
+        "B": {"facecolor": "xkcd:water blue"},
+        "D": {"projection": "polar", "facecolor": "w"},
+    }
+)
+identify_axes(axd)
+
 
 ###############################################################################
 # Nested list input
