ENH: Add subfigure_mosaic for creating SubFigure layouts

lib/matplotlib/figure.py

@@ -1896,12 +1896,12 @@ def _normalize_grid_string(layout):
             layout = inspect.cleandoc(layout)
             return [list(ln) for ln in layout.strip('\n').split('\n')]
 
-    def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False,
-                       width_ratios=None, height_ratios=None,
-                       empty_sentinel='.',
-                       subplot_kw=None, per_subplot_kw=None, gridspec_kw=None):
+    def _do_mosaic(self, mosaic, *, make_child, sharex=False, sharey=False,
+                   width_ratios=None, height_ratios=None,
+                   empty_sentinel='.',
+                   subplot_kw=None, per_subplot_kw=None, gridspec_kw=None):
         """
-        Build a layout of Axes based on ASCII art or nested lists.
+        Private method to build a layout based on ASCII art or nested lists.
 
         This is a helper function to build complex GridSpec layouts visually.
 
@@ -1949,9 +1949,10 @@ def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False,
 
         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
-            units behave as for `subplots`.  If False, each subplot's x- or
-            y-axis will be independent.
+            among all created children that support sharing (i.e., Axes objects).
+            For SubFigures created by subfigure_mosaic, these parameters are not
+            applicable since SubFigures are containers that don't have axes at
+            creation time.
 
         width_ratios : array-like of length *ncols*, optional
             Defines the relative widths of the columns. Each column gets a
@@ -2043,7 +2044,7 @@ def _make_array(inp):
 
             Returns
             -------
-            2D object array
+            2D object arraywam
             """
             r0, *rest = inp
             if isinstance(r0, str):
@@ -2156,14 +2157,14 @@ def _do_layout(gs, mosaic, unique_ids, nested):
                     if name in output:
                         raise ValueError(f"There are duplicate keys {name} "
                                          f"in the layout\n{mosaic!r}")
-                    ax = self.add_subplot(
+                    child = make_child(
                         gs[slc], **{
                             'label': str(name),
                             **subplot_kw,
                             **per_subplot_kw.get(name, {})
                         }
                     )
-                    output[name] = ax
+                    output[name] = child
                 elif method == 'nested':
                     nested_mosaic = arg
                     j, k = key
@@ -2190,21 +2191,165 @@ def _do_layout(gs, mosaic, unique_ids, nested):
         rows, cols = mosaic.shape
         gs = self.add_gridspec(rows, cols, **gridspec_kw)
         ret = _do_layout(gs, mosaic, *_identify_keys_and_nested(mosaic))
-        ax0 = next(iter(ret.values()))
-        for ax in ret.values():
+        first_child = next(iter(ret.values()))
+        for child in ret.values():
             if sharex:
-                ax.sharex(ax0)
-                ax._label_outer_xaxis(skip_non_rectangular_axes=True)
+                child.sharex(first_child)
+                child._label_outer_xaxis(skip_non_rectangular_axes=True)
             if sharey:
-                ax.sharey(ax0)
-                ax._label_outer_yaxis(skip_non_rectangular_axes=True)
+                child.sharey(first_child)
+                child._label_outer_yaxis(skip_non_rectangular_axes=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 subplot_mosaic(self, mosaic, *, sharex=False, sharey=False,
+                       width_ratios=None, height_ratios=None,
+                       empty_sentinel='.',
+                       subplot_kw=None, per_subplot_kw=None, gridspec_kw=None):
+        """
+        Build a layout of Axes based on ASCII art or nested lists.
+
+        This is a helper function to build complex GridSpec layouts visually.
+
+        See :ref:`mosaic`
+        for an example and full API documentation
+
+        Parameters
+        ----------
+        mosaic : list of list of {hashable or nested} or str
+            A visual layout of how you want your Axes to be arranged
+            labeled as strings.
+        sharex, sharey : bool, default: False
+            If True, the x-axis (*sharex*) or y-axis (*sharey*) will be shared
+            among all subplots.
+        width_ratios : array-like of length *ncols*, optional
+            Defines the relative widths of the columns.
+        height_ratios : array-like of length *nrows*, optional
+            Defines the relative heights of the rows.
+        empty_sentinel : object, optional
+            Entry in the layout to mean "leave this space empty".
+            Defaults to ``'.'``.
+        subplot_kw : dict, optional
+            Dictionary with keywords passed to the `.Figure.add_subplot` call
+            used to create each subplot.
+        per_subplot_kw : dict, optional
+            A dictionary mapping the Axes identifiers to a dictionary of
+            keyword arguments to be passed to the `.Figure.add_subplot` call
+            used to create each subplot.
+        gridspec_kw : dict, optional
+            Dictionary with keywords passed to the `.GridSpec` constructor
+            used to create the grid the subplots are placed on.
+
+        Returns
+        -------
+        dict[label, Axes]
+            A dictionary mapping the labels to the Axes objects.
+        """
+        return self._do_mosaic(
+            mosaic,
+            make_child=self.add_subplot,
+            sharex=sharex,
+            sharey=sharey,
+            width_ratios=width_ratios,
+            height_ratios=height_ratios,
+            empty_sentinel=empty_sentinel,
+            subplot_kw=subplot_kw,
+            per_subplot_kw=per_subplot_kw,
+            gridspec_kw=gridspec_kw
+        )
+
+    def subfigure_mosaic(self, mosaic, *,
+                         width_ratios=None, height_ratios=None,
+                         empty_sentinel='.',
+                         per_subfigure_kw=None, gridspec_kw=None, **kwargs):
+        """
+        Build a layout of SubFigures based on ASCII art or nested lists.
+
+        This method is similar to `.subplot_mosaic`, but creates `.SubFigure`
+        objects instead of `.Axes` objects. This is useful for creating complex
+        layouts where each subfigure can contain its own subplots, colorbars,
+        and other elements.
+
+        See :ref:`mosaic` for an example and full API documentation.
+
+        Parameters
+        ----------
+        mosaic : list of list of {hashable or nested} or str
+            A visual layout of how you want your SubFigures to be arranged,
+            labeled as strings. The same syntax as `.subplot_mosaic` is supported.
+
+        width_ratios : array-like of length *ncols*, optional
+            Defines the relative widths of the columns. Each column gets a
+            relative width of ``width_ratios[i] / sum(width_ratios)``.
+            If not given, all columns will have the same width.
+
+        height_ratios : array-like of length *nrows*, optional
+            Defines the relative heights of the rows. Each row gets a
+            relative height of ``height_ratios[i] / sum(height_ratios)``.
+            If not given, all rows will have the same height.
+
+        empty_sentinel : object, optional
+            Entry in the layout to mean "leave this space empty". Defaults
+            to ``'.'``.
+
+        per_subfigure_kw : dict, optional
+            A dictionary mapping the SubFigure identifiers to a dictionary of
+            keyword arguments to be passed to the `.Figure.add_subfigure` call
+            used to create each subfigure. The values in these dictionaries
+            have precedence over values in *kwargs*.
+
+        gridspec_kw : dict, optional
+            Dictionary with keywords passed to the `.GridSpec` constructor used
+            to create the grid the subfigures are placed on.
+
+        **kwargs
+            Additional keyword arguments passed to `.Figure.add_subfigure` for
+            all subfigures.
+
+        Returns
+        -------
+        dict[label, SubFigure]
+            A dictionary mapping the labels to the `.SubFigure` objects.
+
+        See Also
+        --------
+        Figure.subplot_mosaic : Create a mosaic of Axes.
+        Figure.subfigures : Create a regular grid of SubFigures.
+        Figure.add_subfigure : Add a single SubFigure.
+
+        Notes
+        -----
+        Unlike `.subplot_mosaic`, this method does not support the *sharex* and
+        *sharey* parameters. While axes across different subfigures can technically
+        share limits, this would need to be configured after the axes are created
+        within each subfigure, not at the subfigure creation time.
+
+        Examples
+        --------
+        >>> fig = plt.figure()
+        >>> subfigs = fig.subfigure_mosaic("AB;CD")
+        >>> subfigs['A'].subplots(2, 2)
+        >>> subfigs['B'].subplots(1, 3)
+        >>> subfigs['C'].subplots()
+        >>> subfigs['D'].subplots(3, 1)
+        """
+        return self._do_mosaic(
+            mosaic,
+            make_child=self.add_subfigure,
+            sharex=False,  # Not applicable for SubFigures at creation time
+            sharey=False,  # Not applicable for SubFigures at creation time
+            width_ratios=width_ratios,
+            height_ratios=height_ratios,
+            empty_sentinel=empty_sentinel,
+            subplot_kw=kwargs,  # Pass kwargs as subplot_kw
+            per_subplot_kw=per_subfigure_kw,
+            gridspec_kw=gridspec_kw
+        )
+
     def _set_artist_props(self, a):
         if a != self:
             a.set_figure(self)

lib/matplotlib/tests/test_figure.py

@@ -1308,6 +1308,123 @@ def test_share_all(self):
                    for ax in ax_dict.values())
 
 
+class TestSubfigureMosaic:
+    """Test suite for subfigure_mosaic functionality."""
+
+    def test_basic(self):
+        """Test basic subfigure_mosaic functionality."""
+        fig = plt.figure(layout='constrained')
+
+        # Create a simple 2x2 mosaic
+        mosaic = [["A", "B"], ["C", "D"]]
+        subfigs = fig.subfigure_mosaic(mosaic)
+
+        # Check we got 4 subfigures with correct keys
+        assert len(subfigs) == 4
+        assert set(subfigs.keys()) == {"A", "B", "C", "D"}
+
+        # Check they're all SubFigure objects
+        for sf in subfigs.values():
+            assert isinstance(sf, mpl.figure.SubFigure)
+
+    def test_string_input(self):
+        """Test string input format for subfigure mosaic."""
+        fig = plt.figure(layout='constrained')
+
+        mosaic = "AB\nCD"
+        subfigs = fig.subfigure_mosaic(mosaic)
+
+        assert len(subfigs) == 4
+        assert set(subfigs.keys()) == {"A", "B", "C", "D"}
+
+    def test_empty_cells(self):
+        """Test that empty cells work with the default '.' sentinel."""
+        fig = plt.figure(layout='constrained')
+
+        # Create mosaic with empty cells
+        mosaic = [["A", ".", "B"], ["C", "C", "."]]
+        subfigs = fig.subfigure_mosaic(mosaic)
+
+        # Should only have 3 subfigures (no '.')
+        assert len(subfigs) == 3
+        assert set(subfigs.keys()) == {"A", "B", "C"}
+
+    def test_width_height_ratios(self):
+        """Test that width and height ratios work."""
+        fig = plt.figure(layout='constrained')
+
+        mosaic = [["A", "B"], ["C", "D"]]
+        width_ratios = [1, 2]
+        height_ratios = [2, 1]
+
+        subfigs = fig.subfigure_mosaic(mosaic,
+                                      width_ratios=width_ratios,
+                                      height_ratios=height_ratios)
+
+        # Just check we got the right subfigures
+        # Actual ratio testing would require checking geometry
+        assert len(subfigs) == 4
+        assert set(subfigs.keys()) == {"A", "B", "C", "D"}
+
+    def test_nested_subfigure_mosaic(self):
+        """Test creating a mosaic within a subfigure from another mosaic.
+
+        This tests the behavior discussed in the issue:
+        - subfigure_mosaic creates FLAT subfigures (all direct children of caller)
+        - But users CAN create hierarchies by calling subfigure_mosaic on a SubFigure
+
+        Expected structure after all calls:
+        Figure
+        ├── SubFigure A
+        │   ├── SubFigure C
+        │   ├── SubFigure D
+        │   ├── SubFigure E
+        │   └── SubFigure F
+        └── SubFigure B
+
+        Return values:
+        - top_subfigs = {"A": SubFigure_A, "B": SubFigure_B}
+        - nested_subfigs = {"C": SubFigure_C, "D": SubFigure_D, 
+                           "E": SubFigure_E, "F": SubFigure_F}
+        """
+        fig = plt.figure(layout='constrained')
+
+        # Step 1: Create top-level subfigures using mosaic "AB"
+        # This should create 2 subfigures as direct children of fig
+        top_subfigs = fig.subfigure_mosaic("AB")
+
+        # Verify return structure
+        assert list(top_subfigs.keys()) == ["A", "B"]
+        assert all(isinstance(sf, mpl.figure.SubFigure) for sf in top_subfigs.values())
+
+        # Verify these are direct children of the figure
+        assert len(fig.subfigs) == 2
+        assert top_subfigs["A"] in fig.subfigs
+        assert top_subfigs["B"] in fig.subfigs
+
+        # Step 2: Create a mosaic WITHIN subfigure A using "CD\nEF"
+        # This creates a 2x2 grid inside subfigure A:
+        #   C D
+        #   E F
+        nested_subfigs = top_subfigs["A"].subfigure_mosaic("CD\nEF")
+
+        # Verify return structure
+        assert list(nested_subfigs.keys()) == ["C", "D", "E", "F"]
+        assert all(isinstance(sf, mpl.figure.SubFigure) for sf in nested_subfigs.values())
+
+        # Verify the nested subfigures are children of A, not the main figure
+        assert len(top_subfigs["A"].subfigs) == 4
+        for key, sf in nested_subfigs.items():
+            assert sf in top_subfigs["A"].subfigs  # IS a child of A
+            assert sf not in fig.subfigs            # NOT a child of fig
+
+        # The main figure should still only have 2 direct children (A and B)
+        assert len(fig.subfigs) == 2
+
+        # Verify B has no children
+        assert len(top_subfigs["B"].subfigs) == 0
+
+
 def test_reused_gridspec():
     """Test that these all use the same gridspec"""
     fig = plt.figure()