Add plot_layout

has2k1 · has2k1 · commit 40d4ff61b2d4 · 2025-09-02T23:50:57.000+03:00
diff --git a/doc/_quartodoc.yml b/doc/_quartodoc.yml
@@ -541,6 +541,7 @@ quartodoc:
         - Stack
         - Beside
         - plot_spacer
+        - plot_layout
 
     - title: Options
       desc: |
diff --git a/doc/changelog.qmd b/doc/changelog.qmd
@@ -2,7 +2,7 @@
 title: Changelog
 ---
 
-## v0.15.1
+## v0.16.0
 (not-yet-released)
 
 ### New
@@ -12,6 +12,9 @@ title: Changelog
 - Added [](:class:`~plotnine.composition.Compose.show`), making it possible to output/show 
   plot composition objects when they are not the last object in a cell.
 
+- Added [](:class:`~plotnine.composition.plot_layout`) with which you can customise the
+  layout of plots in composition.
+
 ### Enhancements
 
 - In a jupyter environment, the output when the plot is the last in a cell is
diff --git a/plotnine/_mpl/layout_manager/_layout_tree.py b/plotnine/_mpl/layout_manager/_layout_tree.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import abc
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import cached_property
 from typing import TYPE_CHECKING
 
@@ -72,7 +72,18 @@ class LayoutTree:
     Each composition is a tree or subtree
     """
 
-    gridspec: p9GridSpec
+    cmp: Compose
+    """
+    Composition that this tree represents
+    """
+
+    nodes: list[LayoutSpaces | LayoutTree]
+    """
+    The spaces or tree of spaces in the composition that the tree
+    represents.
+    """
+
+    gridspec: p9GridSpec = field(init=False)
     """
     Gridspec of the composition
 
@@ -86,11 +97,8 @@ class LayoutTree:
     LayoutSpaces.
     """
 
-    nodes: list[LayoutSpaces | LayoutTree]
-    """
-    The spaces or tree of spaces in the composition that the tree
-    represents.
-    """
+    def __post_init__(self):
+        self.gridspec = self.cmp.gridspec
 
     @staticmethod
     def create(
@@ -122,9 +130,9 @@ def create(
                 nodes.append(LayoutTree.create(item, lookup_spaces))
 
         if isinstance(cmp, Beside):
-            return ColumnsTree(cmp.gridspec, nodes)
+            return ColumnsTree(cmp, nodes)
         else:
-            return RowsTree(cmp.gridspec, nodes)
+            return RowsTree(cmp, nodes)
 
     @cached_property
     def sub_compositions(self) -> list[LayoutTree]:
@@ -310,13 +318,15 @@ def resize(self):
         """
         Resize the widths of gridspec so that panels have equal widths
         """
-        # The new width of each panel is the average width of all
-        # the panels plus all the space to the left and right
-        # of the panels.
+        # The new width of each panel is the average width (scaled by
+        # the plot_layout widths) of all the panels plus all the space
+        # to the left and right of the panels.
+        widths = np.array(self.cmp._plot_layout.widths)
         plot_widths = np.array(self.plot_widths)
         panel_widths = np.array(self.panel_widths)
         non_panel_space = plot_widths - panel_widths
-        new_plot_widths = panel_widths.mean() + non_panel_space
+        scaled_panel_widths = panel_widths.mean() * widths
+        new_plot_widths = scaled_panel_widths + non_panel_space
         width_ratios = new_plot_widths / new_plot_widths.min()
         self.gridspec.set_width_ratios(width_ratios)
         self.resize_sub_compositions()
@@ -472,12 +482,15 @@ def resize(self):
 
         This method resizes (recursively) the contained compositions
         """
-        # The new height of each panel is the average width of all
-        # the panels plus all the space above and below the panels.
+        # The new height of each panel is the average width (scaled by
+        # the plot_layout heights) of all the panels plus all the space
+        # above and below the panels.
+        heights = np.array(self.cmp._plot_layout.heights)
         plot_heights = np.array(self.plot_heights)
         panel_heights = np.array(self.panel_heights)
         non_panel_space = plot_heights - panel_heights
-        new_plot_heights = panel_heights.mean() + non_panel_space
+        scaled_panel_heights = panel_heights.mean() * heights
+        new_plot_heights = scaled_panel_heights + non_panel_space
         height_ratios = new_plot_heights / new_plot_heights.max()
         self.gridspec.set_height_ratios(height_ratios)
         self.resize_sub_compositions()
diff --git a/plotnine/composition/__init__.py b/plotnine/composition/__init__.py
@@ -1,11 +1,13 @@
 from ._beside import Beside
 from ._compose import Compose
+from ._plot_layout import plot_layout
 from ._plot_spacer import plot_spacer
 from ._stack import Stack
 
 __all__ = (
     "Compose",
     "Stack",
     "Beside",
+    "plot_layout",
     "plot_spacer",
 )
diff --git a/plotnine/composition/_compose.py b/plotnine/composition/_compose.py
@@ -6,6 +6,9 @@
 from io import BytesIO
 from typing import TYPE_CHECKING, overload
 
+from plotnine.composition._plot_layout import plot_layout
+from plotnine.composition._types import ComposeAddable
+
 from .._utils.context import plot_composition_context
 from .._utils.ipython import (
     get_ipython,
@@ -85,7 +88,12 @@ class Compose:
 
     items: list[ggplot | Compose]
     """
-    The objects to be arranged (composed).
+    The objects to be arranged (composed)
+    """
+
+    _plot_layout: plot_layout = field(init=False, repr=False)
+    """
+    The instance of plot_layout added to the composition
     """
 
     # These are created in the _create_figure method
@@ -130,7 +138,10 @@ def __truediv__(self, rhs: ggplot | Compose) -> Compose:
         Add rhs as a row
         """
 
-    def __add__(self, rhs: ggplot | Compose | PlotAddable) -> Compose:
+    def __add__(
+        self,
+        rhs: ggplot | Compose | PlotAddable | ComposeAddable,
+    ) -> Compose:
         """
         Add rhs to the composition
 
@@ -141,10 +152,13 @@ def __add__(self, rhs: ggplot | Compose | PlotAddable) -> Compose:
         """
         from plotnine import ggplot
 
-        if not isinstance(rhs, (ggplot, Compose)):
-            cmp = deepcopy(self)
-            cmp.last_plot = cmp.last_plot + rhs
-            return cmp
+        self = deepcopy(self)
+
+        if isinstance(rhs, ComposeAddable):
+            return rhs.__radd__(self)
+        elif not isinstance(rhs, (ggplot, Compose)):
+            self.last_plot = self.last_plot + rhs
+            return self
 
         t1, t2 = type(self).__name__, type(rhs).__name__
         msg = f"unsupported operand type(s) for +: '{t1}' and '{t2}'"
@@ -181,14 +195,12 @@ def __and__(self, rhs: PlotAddable) -> Compose:
         """
         self = deepcopy(self)
 
-        def add_other(cmp: Compose):
-            for i, item in enumerate(cmp):
-                if isinstance(item, Compose):
-                    add_other(item)
-                else:
-                    cmp[i] = item + copy(rhs)
+        for i, item in enumerate(self):
+            if isinstance(item, Compose):
+                self[i] = item & rhs
+            else:
+                item += copy(rhs)
 
-        add_other(self)
         return self
 
     def __mul__(self, rhs: PlotAddable) -> Compose:
@@ -204,9 +216,9 @@ def __mul__(self, rhs: PlotAddable) -> Compose:
 
         self = deepcopy(self)
 
-        for i, item in enumerate(self):
+        for item in self:
             if isinstance(item, ggplot):
-                self[i] = item + copy(rhs)
+                item += copy(rhs)
 
         return self
 
@@ -338,6 +350,11 @@ def _create_gridspec(self, figure, nest_into):
             self.nrow, self.ncol, figure, nest_into=nest_into
         )
 
+        if not hasattr(self, "_plot_layout"):
+            self._plot_layout = plot_layout()
+
+        self._plot_layout._setup(self.nrow, self.ncol)
+
     def _setup(self) -> Figure:
         """
         Setup this instance for the building process
diff --git a/plotnine/composition/_plot_layout.py b/plotnine/composition/_plot_layout.py
@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from copy import copy
+from dataclasses import dataclass
+from itertools import cycle
+from typing import TYPE_CHECKING, Sequence
+
+from plotnine.composition._types import ComposeAddable
+
+if TYPE_CHECKING:
+    from ._compose import Compose
+
+
+@dataclass(kw_only=True)
+class plot_layout(ComposeAddable):
+    """
+    Customise the layout of plots in a composition
+    """
+
+    widths: Sequence[float] | None = None
+    """
+    Relative widths of each column
+    """
+
+    heights: Sequence[float] | None = None
+    """
+    Relative heights of each column
+    """
+
+    def __radd__(self, cmp: Compose) -> Compose:
+        """
+        Add plot layout to composition
+        """
+        cmp._plot_layout = copy(self)
+        return cmp
+
+    def _setup(self, nrow: int, ncol: int):
+        """
+        Setup default parameters as they are expected by the layout manager
+        """
+        # from mizani.bounds import rescale
+
+        ws, hs = self.widths, self.heights
+        if ws is None:
+            ws = (1 / ncol,) * ncol
+        elif len(ws) != ncol:
+            ws = repeat(ws, ncol)
+
+        if hs is None:
+            hs = (1 / nrow,) * nrow
+        elif len(hs) != nrow:
+            hs = repeat(hs, nrow)
+
+        self.widths = rescale(ws)
+        self.heights = rescale(hs)
+
+
+def repeat(seq: Sequence[float], n: int) -> list[float]:
+    """
+    Ensure returned sequence has n values, repeat as necessary
+    """
+    return [val for _, val in zip(range(n), cycle(seq))]
+
+
+def rescale(seq) -> list[float]:
+    high = max(seq)
+    return [x / high for x in seq]
diff --git a/plotnine/composition/_types.py b/plotnine/composition/_types.py
@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ._compose import Compose
+
+
+class ComposeAddable:
+    """
+    Object that can be added to a ggplot object
+    """
+
+    def __radd__(self, other: Compose) -> Compose:
+        """
+        Add to compose object
+
+        Parameters
+        ----------
+        other :
+            Compose object
+
+        Returns
+        -------
+        :
+            Compose object
+        """
+        return other
diff --git a/tests/test_plot_composition.py b/tests/test_plot_composition.py
@@ -1,6 +1,7 @@
 from plotnine import element_text, facet_grid, facet_wrap, theme, theme_gray
 from plotnine._utils.yippie import geom as g
 from plotnine._utils.yippie import legend, plot, rotate, tag
+from plotnine.composition._plot_layout import plot_layout
 
 
 def test_basic_horizontal_align_resize():
@@ -151,3 +152,30 @@ def test_plus_operator():
     p4 = plot.brown
     p = (p1 | p2 | (p3 / p4)) + theme_gray()
     assert p == "plus_operator"
+
+
+def test_plot_layout_widths():
+    p1 = plot.red
+    p2 = plot.green
+    p = (p1 | p2) + plot_layout(widths=[1, 4])
+    assert p == "plot_layout_widths"
+
+
+def test_plot_layout_heights():
+    p1 = plot.red
+    p2 = plot.green
+    p = (p1 / p2) + plot_layout(heights=[1, 4])
+    assert p == "plot_layout_heights"
+
+
+def test_plot_layout_nested_resize():
+    p1 = plot.red
+    p2 = plot.green
+    p3 = plot.blue
+    p4 = plot.brown
+
+    ws = plot_layout(widths=[1, 4])
+    hs = plot_layout(heights=[1, 3])
+
+    p = (((p1 | p2) + ws) / ((p3 | p4) + ws)) + hs
+    assert p == "plot_layout_nested_resize"
