Add plot_layout

Author: has2k1

doc/_quartodoc.yml

@@ -541,6 +541,7 @@ quartodoc:
         - Stack
         - Beside
         - plot_spacer
+        - plot_layout
 
     - title: Options
       desc: |

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

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,14 +318,16 @@ 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.
-        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
-        width_ratios = new_plot_widths / new_plot_widths.min()
+        # The new width of each panel is the total panel area scaled
+        # by the factor given in plot_layout.
+        # The new width of the plot includes the space not taken up
+        # by the panels.
+        factors = np.array(self.cmp._plot_layout.widths)
+        _totals = np.ones(len(self.panel_widths)) * sum(self.panel_widths)
+        scaled_panel_widths = _totals * factors
+        non_panel_space = np.array(self.plot_widths) - self.panel_widths
+        new_plot_widths = scaled_panel_widths + non_panel_space
+        width_ratios = new_plot_widths / new_plot_widths.max()
         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.
-        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
+        # The new height of each panel is the total panel area scaled
+        # by the factor given in plot_layout.
+        # The new height of the plot includes the space not taken up
+        # by the panels.
+        factors = np.array(self.cmp._plot_layout.heights)
+        _totals = np.ones(len(self.panel_heights)) * sum(self.panel_heights)
+        scaled_panel_heights = _totals * factors
+        non_panel_space = np.array(self.plot_heights) - self.panel_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()

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",
 )

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

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
+
+        - Ensure that the widths and heights are set and normalised to unit sum
+        """
+        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 = normalise(ws)
+        self.heights = normalise(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 normalise(seq) -> list[float]:
+    total = sum(seq)
+    return [x / total for x in seq]

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