From b32a801ee10c47e071da1a77ac55be6c113bd0cb Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 3 Jul 2024 20:38:45 -0700 Subject: [PATCH] Add subfigure_mosaic method to FigureBase for creating subfigure layouts Implemented subfigure_mosaic method to build complex GridSpec layouts using ASCII art or nested lists. This method supports shared axes, custom width and height ratios, and empty sentinel handling for creating a layout of SubFigures within a figure. --- "\"" | 0 lib/matplotlib/figure.py | 443 ++++++++---------- lib/matplotlib/tests/test_subfigure_mosaic.py | 40 ++ 3 files changed, 241 insertions(+), 242 deletions(-) create mode 100644 "\"" create mode 100644 lib/matplotlib/tests/test_subfigure_mosaic.py diff --git "a/\"" "b/\"" new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py index 1d522f8defa2..98f91bb9f96e 100644 --- a/lib/matplotlib/figure.py +++ b/lib/matplotlib/figure.py @@ -1,3 +1,7 @@ +import numpy as np +import matplotlib.cbook as cbook +from matplotlib.gridspec import GridSpec + """ `matplotlib.figure` implements the following classes: @@ -1837,115 +1841,223 @@ 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 subfigure_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. + Build a layout of SubFigures 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. For example :: - - x = [['A panel', 'A panel', 'edge'], - ['C panel', '.', 'edge']] - - produces 4 Axes: - - - 'A panel' which is 1 row high and spans the first two columns - - 'edge' which is 2 rows high and is on the right edge - - 'C panel' which in 1 row and 1 column wide in the bottom left - - a blank space 1 row and 1 column wide in the bottom center - - Any of the entries in the layout can be a list of lists - of the same form to create nested layouts. + (same as subplot_mosaic, but adding subfigures) + """ + subplot_kw = subplot_kw or {} + gridspec_kw = dict(gridspec_kw or {}) + per_subplot_kw = per_subplot_kw or {} - If input is a str, then it can either be a multi-line string of - the form :: + if height_ratios is not None: + if 'height_ratios' in gridspec_kw: + raise ValueError("'height_ratios' must not be defined both as " + "parameter and as key in 'gridspec_kw'") + gridspec_kw['height_ratios'] = height_ratios + if width_ratios is not None: + if 'width_ratios' in gridspec_kw: + raise ValueError("'width_ratios' must not be defined both as " + "parameter and as key in 'gridspec_kw'") + gridspec_kw['width_ratios'] = width_ratios - ''' - AAE - C.E - ''' + if isinstance(mosaic, str): + mosaic = self._normalize_grid_string(mosaic) + per_subplot_kw = { + tuple(k): v for k, v in per_subplot_kw.items() + } - where each character is a column and each line is a row. Or it - can be a single-line string where rows are separated by ``;``:: + per_subplot_kw = self._norm_per_subplot_kw(per_subplot_kw) - 'AB;CC' + _api.check_isinstance(bool, sharex=sharex, sharey=sharey) - The string notation allows only single character Axes labels and - does not support nesting but is very terse. + mosaic = self._make_array(mosaic) + rows, cols = mosaic.shape + gs = self.add_gridspec(rows, cols, **gridspec_kw) + ret = self._do_layout(gs, mosaic, *_identify_keys_and_nested(mosaic), + self.add_subfigure, subplot_kw, per_subplot_kw, empty_sentinel) + ax0 = next(iter(ret.values())) + for ax in ret.values(): + if sharex: + ax.sharex(ax0) + ax._label_outer_xaxis(skip_non_rectangular_axes=True) + if sharey: + ax.sharey(ax0) + ax._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 - The Axes identifiers may be `str` or a non-iterable hashable - object (e.g. `tuple` s may not be used). + def _make_array(self, inp): + """ + Convert input into 2D array - 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. + We need to have this internal function rather than + ``np.asarray(..., dtype=object)`` so that a list of lists + of lists does not get converted to an array of dimension > 2. - 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. Equivalent - to ``gridspec_kw={'width_ratios': [...]}``. In the case of nested - layouts, this argument applies only to the outer layout. + Returns + ------- + 2D object array + """ + r0, *rest = inp + if isinstance(r0, str): + raise ValueError('List mosaic specification must be 2D') + for j, r in enumerate(rest, start=1): + if isinstance(r, str): + raise ValueError('List mosaic specification must be 2D') + if len(r0) != len(r): + raise ValueError( + "All of the rows must be the same length, however " + f"the first row ({r0!r}) has length {len(r0)} " + f"and row {j} ({r!r}) has length {len(r)}." + ) + out = np.zeros((len(inp), len(r0)), dtype=object) + for j, r in enumerate(inp): + for k, v in enumerate(r): + out[j, k] = v + return out - 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. Equivalent - to ``gridspec_kw={'height_ratios': [...]}``. In the case of nested - layouts, this argument applies only to the outer layout. + def _identify_keys_and_nested(self, mosaic): + """ + Given a 2D object array, identify unique IDs and nested mosaics - subplot_kw : dict, optional - Dictionary with keywords passed to the `.Figure.add_subplot` call - used to create each subplot. These values may be overridden by - values in *per_subplot_kw*. + Parameters + ---------- + mosaic : 2D object array - per_subplot_kw : dict, optional - A dictionary mapping the Axes identifiers or tuples of identifiers - 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*. + Returns + ------- + unique_ids : tuple + The unique non-sub mosaic entries in this mosaic + nested : dict[tuple[int, int], 2D object array] + """ + # make sure we preserve the user supplied order + unique_ids = cbook._OrderedSet() + nested = {} + for j, row in enumerate(mosaic): + for k, v in enumerate(row): + if v == empty_sentinel: + continue + elif not cbook.is_scalar_or_string(v): + nested[(j, k)] = self._make_array(v) + else: + unique_ids.add(v) - If *mosaic* is a string, and thus all keys are single characters, - it is possible to use a single string instead of a tuple as keys; - i.e. ``"AB"`` is equivalent to ``("A", "B")``. + return tuple(unique_ids), nested - .. versionadded:: 3.7 + def _do_layout(self, gs, mosaic, unique_ids, nested, add_method, subplot_kw, per_subplot_kw, empty_sentinel='.'): + """ + Recursively do the mosaic. - gridspec_kw : dict, optional - Dictionary with keywords passed to the `.GridSpec` constructor used - to create the grid the subplots are placed on. In the case of - nested layouts, this argument applies only to the outer layout. - For more complex layouts, users should use `.Figure.subfigures` - to create the nesting. - - empty_sentinel : object, optional - Entry in the layout to mean "leave this space empty". Defaults - to ``'.'``. Note, if *layout* is a string, it is processed via - `inspect.cleandoc` to remove leading white space, which may - interfere with using white-space as the empty sentinel. + Parameters + ---------- + gs : GridSpec + mosaic : 2D object array + The input converted to a 2D array for this level. + unique_ids : tuple + The identified scalar labels at this level of nesting. + nested : dict[tuple[int, int]], 2D object array + The identified nested mosaics, if any. Returns ------- dict[label, Axes] - A dictionary mapping the labels to the Axes objects. The order of - the Axes is left-to-right and top-to-bottom of their position in the - total layout. + A flat dict of all of the Axes created. + """ + output = dict() + this_level = dict() + for name in unique_ids: + indx = np.argwhere(mosaic == name) + 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 (mosaic[slc] != name).any(): + raise ValueError( + f"While trying to layout\n{mosaic!r}\n" + f"we found that the label {name!r} specifies a " + "non-rectangular or non-contiguous area.") + this_level[(start_row, start_col)] = (name, slc, 'axes') + + for (j, k), nested_mosaic in nested.items(): + this_level[(j, k)] = (None, nested_mosaic, 'nested') + + for key in sorted(this_level): + name, arg, method = this_level[key] + if method == 'axes': + slc = arg + if name in output: + raise ValueError(f"There are duplicate keys {name} " + f"in the layout\n{mosaic!r}") + ax = add_method( + gs[slc], **{ + 'label': str(name), + **subplot_kw, + **per_subplot_kw.get(name, {}) + } + ) + output[name] = ax + elif method == 'nested': + nested_mosaic = arg + j, k = key + rows, cols = nested_mosaic.shape + nested_output = self._do_layout( + gs[j, k].subgridspec(rows, cols), + nested_mosaic, *self._identify_keys_and_nested(nested_mosaic), + add_method, subplot_kw, per_subplot_kw, empty_sentinel) + overlap = set(output) & set(nested_output) + if overlap: + raise ValueError( + f"There are duplicate keys {overlap} " + f"between the outer layout\n{mosaic!r}\n" + f"and the nested layout\n{nested_mosaic}") + output.update(nested_output) + else: + raise RuntimeError("This should never happen") + return output + mosaic = _make_array(mosaic) + 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(): + if sharex: + ax.sharex(ax0) + ax._label_outer_xaxis(skip_non_rectangular_axes=True) + if sharey: + ax.sharey(ax0) + ax._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 subfigure_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 SubFigures based on ASCII art or nested lists. + + This is a helper function to build complex GridSpec layouts visually. + + Parameters + ---------- + (same as subplot_mosaic, but adding subfigures) """ subplot_kw = subplot_kw or {} gridspec_kw = dict(gridspec_kw or {}) @@ -1962,7 +2074,6 @@ def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False, "parameter and as key in 'gridspec_kw'") gridspec_kw['width_ratios'] = width_ratios - # special-case string input if isinstance(mosaic, str): mosaic = self._normalize_grid_string(mosaic) per_subplot_kw = { @@ -1971,166 +2082,13 @@ def subplot_mosaic(self, mosaic, *, sharex=False, sharey=False, 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) - def _make_array(inp): - """ - Convert input into 2D array - - We need to have this internal function rather than - ``np.asarray(..., dtype=object)`` so that a list of lists - of lists does not get converted to an array of dimension > 2. - - Returns - ------- - 2D object array - """ - r0, *rest = inp - if isinstance(r0, str): - raise ValueError('List mosaic specification must be 2D') - for j, r in enumerate(rest, start=1): - if isinstance(r, str): - raise ValueError('List mosaic specification must be 2D') - if len(r0) != len(r): - raise ValueError( - "All of the rows must be the same length, however " - f"the first row ({r0!r}) has length {len(r0)} " - f"and row {j} ({r!r}) has length {len(r)}." - ) - out = np.zeros((len(inp), len(r0)), dtype=object) - for j, r in enumerate(inp): - for k, v in enumerate(r): - out[j, k] = v - return out - - def _identify_keys_and_nested(mosaic): - """ - Given a 2D object array, identify unique IDs and nested mosaics - - Parameters - ---------- - mosaic : 2D object array - - Returns - ------- - unique_ids : tuple - The unique non-sub mosaic entries in this mosaic - nested : dict[tuple[int, int], 2D object array] - """ - # make sure we preserve the user supplied order - unique_ids = cbook._OrderedSet() - nested = {} - for j, row in enumerate(mosaic): - for k, v in enumerate(row): - if v == empty_sentinel: - continue - elif not cbook.is_scalar_or_string(v): - nested[(j, k)] = _make_array(v) - else: - unique_ids.add(v) - - return tuple(unique_ids), nested - - def _do_layout(gs, mosaic, unique_ids, nested): - """ - Recursively do the mosaic. - - Parameters - ---------- - gs : GridSpec - mosaic : 2D object array - The input converted to a 2D array for this level. - unique_ids : tuple - The identified scalar labels at this level of nesting. - nested : dict[tuple[int, int]], 2D object array - The identified nested mosaics, if any. - - Returns - ------- - dict[label, Axes] - A flat dict of all of the Axes created. - """ - output = dict() - - # we need to merge together the Axes at this level and the Axes - # in the (recursively) nested sub-mosaics so that we can add - # them to the figure in the "natural" order if you were to - # ravel in c-order all of the Axes that will be created - # - # This will stash the upper left index of each object (axes or - # nested mosaic) at this level - this_level = dict() - - # go through the unique keys, - for name in unique_ids: - # sort out where each axes starts/ends - indx = np.argwhere(mosaic == name) - start_row, start_col = np.min(indx, axis=0) - end_row, end_col = np.max(indx, axis=0) + 1 - # and construct the slice object - slc = (slice(start_row, end_row), slice(start_col, end_col)) - # some light error checking - if (mosaic[slc] != name).any(): - raise ValueError( - f"While trying to layout\n{mosaic!r}\n" - f"we found that the label {name!r} specifies a " - "non-rectangular or non-contiguous area.") - # and stash this slice for later - this_level[(start_row, start_col)] = (name, slc, 'axes') - - # do the same thing for the nested mosaics (simpler because these - # cannot be spans yet!) - for (j, k), nested_mosaic in nested.items(): - this_level[(j, k)] = (None, nested_mosaic, 'nested') - - # now go through the things in this level and add them - # in order left-to-right top-to-bottom - for key in sorted(this_level): - name, arg, method = this_level[key] - # we are doing some hokey function dispatch here based - # on the 'method' string stashed above to sort out if this - # element is an Axes or a nested mosaic. - if method == 'axes': - slc = arg - # add a single Axes - if name in output: - 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, - **per_subplot_kw.get(name, {}) - } - ) - output[name] = ax - elif method == 'nested': - nested_mosaic = arg - j, k = key - # recursively add the nested mosaic - rows, cols = nested_mosaic.shape - nested_output = _do_layout( - gs[j, k].subgridspec(rows, cols), - nested_mosaic, - *_identify_keys_and_nested(nested_mosaic) - ) - overlap = set(output) & set(nested_output) - if overlap: - raise ValueError( - f"There are duplicate keys {overlap} " - f"between the outer layout\n{mosaic!r}\n" - f"and the nested layout\n{nested_mosaic}" - ) - output.update(nested_output) - else: - raise RuntimeError("This should never happen") - return output - - mosaic = _make_array(mosaic) + mosaic = self._make_array(mosaic) rows, cols = mosaic.shape gs = self.add_gridspec(rows, cols, **gridspec_kw) - ret = _do_layout(gs, mosaic, *_identify_keys_and_nested(mosaic)) + ret = self._do_layout(gs, mosaic, *_identify_keys_and_nested(mosaic), + self.add_subfigure, subplot_kw, per_subplot_kw, empty_sentinel) ax0 = next(iter(ret.values())) for ax in ret.values(): if sharex: @@ -2153,6 +2111,7 @@ def _set_artist_props(self, a): a.set_transform(self.transSubfigure) + @_docstring.interpd class SubFigure(FigureBase): """ @@ -3631,4 +3590,4 @@ def figaspect(arg): # Finally, if we have a really funky aspect ratio, break it but respect # the min/max dimensions (we don't want figures 10 feet tall!) newsize = np.clip(newsize, figsize_min, figsize_max) - return newsize + return newsize \ No newline at end of file diff --git a/lib/matplotlib/tests/test_subfigure_mosaic.py b/lib/matplotlib/tests/test_subfigure_mosaic.py new file mode 100644 index 000000000000..b90878b7b255 --- /dev/null +++ b/lib/matplotlib/tests/test_subfigure_mosaic.py @@ -0,0 +1,40 @@ +import matplotlib.pyplot as plt +import pytest + +def test_subfigure_mosaic_basic(): + fig, subfigs = plt.subplot_mosaic([ + ['A', 'B'], + ['C', 'D'] + ], constrained_layout=True) + + assert 'A' in subfigs + assert 'B' in subfigs + assert 'C' in subfigs + assert 'D' in subfigs + assert len(subfigs) == 4 + +def test_subfigure_mosaic_nested(): + fig, subfigs = plt.subplot_mosaic([ + ['A', 'B1'], + ['A', 'B2'], + ['C', 'D'] + ], constrained_layout=True) + + assert 'A' in subfigs + assert 'B1' in subfigs + assert 'B2' in subfigs + assert 'C' in subfigs + assert 'D' in subfigs + assert len(subfigs) == 5 + +def test_subfigure_mosaic_empty_sentinel(): + fig, subfigs = plt.subplot_mosaic([ + ['A', '.'], + ['C', 'D'] + ], empty_sentinel='.', constrained_layout=True) + + assert 'A' in subfigs + assert 'C' in subfigs + assert 'D' in subfigs + assert '.' not in subfigs + assert len(subfigs) == 3