Add design parameter to plot_layout

doc/changelog.qmd

@@ -10,6 +10,9 @@ title: Changelog
 - Added [](:class:`~plotnine.composition.plot_layout`) with which you can customise the
   layout of plots in composition.
 
+- [](:class:`~plotnine.composition.plot_layout`) gained a `design` parameter for
+  patchwork-style text-grid layouts.
+
 - You can now pass a sequence of horizontal and vertical alignment (ha & va) values in
   element_text for colorbars.
 

plotnine/_mpl/layout_manager/_grid.py

@@ -1,6 +1,7 @@
 from collections.abc import Iterator
 from dataclasses import InitVar, dataclass
 from typing import (
+    Callable,
     Generic,
     Literal,
     Sequence,
@@ -12,6 +13,9 @@
 
 T = TypeVar("T")
 
+Rect = tuple[int, int, int, int]
+"""Inclusive (r0, r1, c0, c1) rectangle in grid coordinates."""
+
 
 @dataclass
 class Grid(Generic[T]):
@@ -92,3 +96,183 @@ def iter_cols(self) -> Iterator[list[T | None]]:
         n = self._grid.shape[1]
         for col in range(n):
             yield self[:, col]
+
+    def reduce_cols(
+        self,
+        fn: Callable[[T], float],
+        default: float,
+    ) -> list[float]:
+        """
+        One value per column: the largest `fn(item)` in that column
+
+        Parameters
+        ----------
+        fn
+            Mapping from an item to the numeric value being compared.
+        default
+            Value used for columns whose cells are all None.
+
+        Returns
+        -------
+        out
+            One value per column, in left-to-right order.
+        """
+        out: list[float] = []
+        for c in range(self._grid.shape[1]):
+            items = [n for n in self[:, c] if n is not None]
+            out.append(max(fn(n) for n in items) if items else default)
+        return out
+
+    def reduce_rows(
+        self,
+        fn: Callable[[T], float],
+        default: float,
+    ) -> list[float]:
+        """
+        One value per row: the largest `fn(item)` in that row
+
+        Parameters
+        ----------
+        fn
+            Mapping from an item to the numeric value being compared.
+        default
+            Value used for rows whose cells are all None.
+
+        Returns
+        -------
+        out
+            One value per row, in top-to-bottom order.
+        """
+        out: list[float] = []
+        for r in range(self._grid.shape[0]):
+            items = [n for n in self[r, :] if n is not None]
+            out.append(max(fn(n) for n in items) if items else default)
+        return out
+
+    def items_on_edge(
+        self,
+        side: Literal["top", "bottom", "left", "right"],
+        idx: int,
+    ) -> list[T]:
+        """
+        Items whose `side` edge sits at row/col `idx`
+
+        In a grid where no item spans more than one cell, an item in
+        row `r` has both its top and bottom edges at row `r`, and
+        analogously for columns; so all four sides return the same
+        items for that row or column.
+
+        Parameters
+        ----------
+        side
+            Which edge of an item to match: `"top"` and `"bottom"`
+            select by row, `"left"` and `"right"` select by column.
+        idx
+            Row index when `side` is `"top"` or `"bottom"`; column
+            index when `side` is `"left"` or `"right"`.
+
+        Returns
+        -------
+        out
+            The matching items, with None cells filtered out.
+        """
+        cells = self[idx, :] if side in ("top", "bottom") else self[:, idx]
+        return [n for n in cells if n is not None]
+
+
+class DesignGrid(Grid[T]):
+    """
+    Grid where items span rectangular regions
+
+    Each item is associated with an inclusive rectangle
+    `(r0, r1, c0, c1)` and placed at every cell within it; this
+    keeps base-class `__getitem__`, `iter_rows`, and `iter_cols`
+    working as in `Grid`. The reductions are overridden to be
+    span-aware: an item spanning multiple columns contributes its
+    measurement divided by its column span to each column it
+    covers (and analogously for rows).
+
+    Parameters
+    ----------
+    nrow
+        Number of rows in the grid.
+    ncol
+        Number of columns in the grid.
+    items
+        Items to place. One per rectangle, in the order rectangles
+        appear in `rects`.
+    rects
+        Inclusive `(r0, r1, c0, c1)` rectangle for each item.
+        Trusted: overlap and shape are not validated here.
+    """
+
+    # Bypass Grid's dataclass __init__ — rectangle expansion is a
+    # different placement scheme than row/col-major.
+    def __init__(
+        self,
+        nrow: int,
+        ncol: int,
+        items: Sequence[T],
+        rects: Sequence[Rect],
+    ):
+        if len(items) != len(rects):
+            raise ValueError(
+                f"Got {len(items)} items but {len(rects)} rectangles"
+            )
+        self._grid = np.empty((nrow, ncol), dtype=object)
+        self._items: list[T] = list(items)
+        self._rects: list[Rect] = list(rects)
+        # Place each item at every cell of its rectangle so the base
+        # class's __getitem__ / iter_rows / iter_cols keep working.
+        for item, (r0, r1, c0, c1) in zip(self._items, self._rects):
+            self._grid[r0 : r1 + 1, c0 : c1 + 1] = item
+
+    def reduce_cols(
+        self,
+        fn: Callable[[T], float],
+        default: float,
+    ) -> list[float]:
+        # An item spanning multiple columns shares its measurement
+        # across the columns it covers: fn(item) / colspan goes into
+        # each. Then per-column max as in Grid.reduce_cols.
+        out: list[float] = []
+        for c in range(self._grid.shape[1]):
+            contribs = [
+                fn(item) / (c1 - c0 + 1)
+                for item, (_, _, c0, c1) in zip(self._items, self._rects)
+                if c0 <= c <= c1
+            ]
+            out.append(max(contribs) if contribs else default)
+        return out
+
+    def reduce_rows(
+        self,
+        fn: Callable[[T], float],
+        default: float,
+    ) -> list[float]:
+        # Mirror of reduce_cols: fn(item) / rowspan into each row the
+        # item covers, then per-row max.
+        out: list[float] = []
+        for r in range(self._grid.shape[0]):
+            contribs = [
+                fn(item) / (r1 - r0 + 1)
+                for item, (r0, r1, _, _) in zip(self._items, self._rects)
+                if r0 <= r <= r1
+            ]
+            out.append(max(contribs) if contribs else default)
+        return out
+
+    def items_on_edge(
+        self,
+        side: Literal["top", "bottom", "left", "right"],
+        idx: int,
+    ) -> list[T]:
+        # An item's top/bottom edge is its r0/r1; left/right is c0/c1.
+        # Match the requested edge to idx exactly — not "the item is
+        # present at row/col idx", which would include spanned cells.
+        out: list[T] = []
+        for item, (r0, r1, c0, c1) in zip(self._items, self._rects):
+            edge = {"top": r0, "bottom": r1, "left": c0, "right": c1}[side]
+            if edge == idx:
+                out.append(item)
+        return out

plotnine/_mpl/layout_manager/_layout_tree.py

@@ -124,19 +124,19 @@ class LayoutTree:
     represents.
     """
 
+    grid: Grid["Node"]
+    """
+    Per-cell layout of `nodes`. `Grid` for compositions without a
+    design; `DesignGrid` when `plot_layout(design=...)` is used.
+    """
+
     sub_gridspec: p9GridSpec = field(init=False, repr=False)
     """
     Gridspec (nxn) that contains the composed items
     """
 
     def __post_init__(self):
         self.sub_gridspec = self.cmp._sub_gridspec
-        self.grid = Grid["Node"](
-            self.nrow,
-            self.ncol,
-            self.nodes,
-            order="row_major" if self.cmp.layout.byrow else "col_major",
-        )
 
     @property
     def ncol(self) -> int:
@@ -172,7 +172,18 @@ def create(cmp: Compose) -> LayoutTree:
             else:
                 nodes.append(LayoutTree.create(item))
 
-        return LayoutTree(cmp, nodes)
+        if (spec := getattr(cmp, "_design_spec", None)) is not None:
+            grid = spec.make_grid(nodes)
+        else:
+            order = "row_major" if cmp.layout.byrow else "col_major"
+            grid = Grid["Node"](
+                cast("int", cmp.layout.nrow),
+                cast("int", cmp.layout.ncol),
+                nodes,
+                order=order,
+            )
+
+        return LayoutTree(cmp, nodes, grid)
 
     @cached_property
     def sub_compositions(self) -> list[LayoutTree]:
@@ -331,14 +342,8 @@ def panel_widths(self) -> Sequence[float]:
         For each column, the representative number for the panel width
         is the maximum width among all panels in the column.
         """
-        # This method is used after aligning the panels. Therefore, the
-        # wides panel_width (i.e. max()) is the good representative width
-        # of the column.
         w = self.plot_width / self.ncol
-        return [
-            max(node.panel_width for node in col if node) if any(col) else w
-            for col in self.grid.iter_cols()
-        ]
+        return self.grid.reduce_cols(lambda n: n.panel_width, default=w)
 
     @property
     def panel_heights(self) -> Sequence[float]:
@@ -349,10 +354,7 @@ def panel_heights(self) -> Sequence[float]:
         is the maximum height among all panels in the row.
         """
         h = self.plot_height / self.nrow
-        return [
-            max([node.panel_height for node in row if node]) if any(row) else h
-            for row in self.grid.iter_rows()
-        ]
+        return self.grid.reduce_rows(lambda n: n.panel_height, default=h)
 
     @property
     def plot_widths(self) -> Sequence[float]:
@@ -363,10 +365,10 @@ def plot_widths(self) -> Sequence[float]:
         the widest plot.
         """
         w = self.sub_gridspec.width / self.ncol
-        return [
-            max([node.plot_width if node else w for node in col])
-            for col in self.grid.iter_cols()
-        ]
+        return self.grid.reduce_cols(
+            lambda n: max(n.plot_width, w),
+            default=w,
+        )
 
     @property
     def plot_heights(self) -> Sequence[float]:
@@ -377,10 +379,10 @@ def plot_heights(self) -> Sequence[float]:
         the tallest plot.
         """
         h = self.sub_gridspec.height / self.nrow
-        return [
-            max([node.plot_height if node else h for node in row])
-            for row in self.grid.iter_rows()
-        ]
+        return self.grid.reduce_rows(
+            lambda n: max(n.plot_height, h),
+            default=h,
+        )
 
     @property
     def panel_width_ratios(self) -> Sequence[float]:
@@ -408,7 +410,7 @@ def bottom_spaces_in_row(self, r: int) -> list[bottom_space]:
         bottom_spaces in the bottom row of that composition.
         """
         spaces: list[bottom_space] = []
-        for node in self.grid[r, :]:
+        for node in self.grid.items_on_edge("bottom", r):
             if isinstance(node, PlotSideSpaces):
                 spaces.append(node.b)
             elif isinstance(node, LayoutTree):
@@ -423,7 +425,7 @@ def top_spaces_in_row(self, r: int) -> list[top_space]:
         top_spaces in the top row of that composition.
         """
         spaces: list[top_space] = []
-        for node in self.grid[r, :]:
+        for node in self.grid.items_on_edge("top", r):
             if isinstance(node, PlotSideSpaces):
                 spaces.append(node.t)
             elif isinstance(node, LayoutTree):
@@ -438,7 +440,7 @@ def left_spaces_in_col(self, c: int) -> list[left_space]:
         left_spaces in the left most column of that composition.
         """
         spaces: list[left_space] = []
-        for node in self.grid[:, c]:
+        for node in self.grid.items_on_edge("left", c):
             if isinstance(node, PlotSideSpaces):
                 spaces.append(node.l)
             elif isinstance(node, LayoutTree):
@@ -453,7 +455,7 @@ def right_spaces_in_col(self, c: int) -> list[right_space]:
         right_spaces in the right most column of that composition.
         """
         spaces: list[right_space] = []
-        for node in self.grid[:, c]:
+        for node in self.grid.items_on_edge("right", c):
             if isinstance(node, PlotSideSpaces):
                 spaces.append(node.r)
             elif isinstance(node, LayoutTree):

plotnine/composition/_compose.py

@@ -33,6 +33,7 @@
     from plotnine._mpl.layout_manager._composition_side_space import (
         CompositionSideSpaces,
     )
+    from plotnine.composition._design import DesignSpec
     from plotnine.composition._guide_area import guide_area
     from plotnine.ggplot import PlotAddable, ggplot
     from plotnine.typing import FigureFormat, MimeBundle
@@ -114,6 +115,11 @@ class Compose:
     plot_layout's theme parameter affects this gridspec.
     """
 
+    _design_spec: DesignSpec | None = None
+    """
+    Parsed `plot_layout(design=...)`. `None` when no design is set.
+    """
+
     _sub_gridspec: p9GridSpec
     """
     Gridspec (nxn) that contains the composed [ggplot | Compose] items
@@ -584,7 +590,15 @@ def _generate_gridspecs(self, figure: p9Figure, container_gs: p9GridSpec):
         # "subplot" in the grid. The SubplotSpec is the handle for the
         # area in the grid; it allows us to put a plot or a nested
         # composion in that area.
-        for item, subplot_spec in zip(self, self._sub_gridspec):
+        # With plot_layout(design=...), each item gets a SubplotSpec
+        # sliced from its rectangle (potentially spanning multiple cells)
+        # instead of one cell per item.
+        if (spec := getattr(self, "_design_spec", None)) is not None:
+            pairs = list(zip(self, spec.get_subplotspecs(self._sub_gridspec)))
+        else:
+            pairs = list(zip(self, self._sub_gridspec))
+
+        for item, subplot_spec in pairs:
             # This container gs will contain a plot or a composition,
             # i.e. it will be assigned to one of:
             #    1. ggplot._gridspec