Add byrow parameter to plot_layout

Author: has2k1

plotnine/_mpl/gridspec.py

@@ -11,15 +11,16 @@
     from matplotlib.figure import SubplotParams
     from matplotlib.gridspec import GridSpecBase
 
+from matplotlib.gridspec import SubplotSpec
 from matplotlib.transforms import Bbox, BboxTransformTo, TransformedBbox
 
 if TYPE_CHECKING:
     from matplotlib.figure import Figure
-    from matplotlib.gridspec import SubplotSpec
     from matplotlib.patches import Rectangle
     from matplotlib.transforms import Transform
 
     from plotnine._mpl.layout_manager._spaces import GridSpecParams
+    from plotnine.composition._plot_layout import plot_layout
 
 
 class p9GridSpec(GridSpecBase):
@@ -60,10 +61,12 @@ def __init__(
         *,
         width_ratios=None,
         height_ratios=None,
+        byrow: bool = True,
         nest_into: SubplotSpec | None = None,
     ):
         self.figure = figure
         self._nested_gridspecs = []
+        self.byrow = byrow
 
         super().__init__(
             nrows,
@@ -91,6 +94,34 @@ def __init__(
             hspace=0,
         )
 
+    @staticmethod
+    def from_layout(
+        layout: plot_layout,
+        figure: Figure,
+        *,
+        nest_into: SubplotSpec | None = None,
+    ) -> p9GridSpec:
+        """
+        Create gridspec from a plot_layout instance
+        """
+        return p9GridSpec(
+            layout.nrow,
+            layout.ncol,
+            figure,
+            byrow=True if layout.byrow is None else layout.byrow,
+            nest_into=nest_into,
+        )
+
+    def __iter__(self):
+        from itertools import product
+
+        if self.byrow:
+            for r, c in product(range(self.nrows), range(self.ncols)):
+                yield SubplotSpec(self, r * self.ncols + c)
+        else:
+            for c, r in product(range(self.ncols), range(self.nrows)):
+                yield SubplotSpec(self, r * self.ncols + c)
+
     @property
     def patch(self) -> Rectangle:
         """

plotnine/_mpl/layout_manager/_grid.py

@@ -39,13 +39,13 @@ def __post_init__(
                 self[r, c] = item
                 c += 1
                 if c >= ncol:
-                    c, r = 0, r + 1
+                    r, c = r + 1, 0
         else:
             for item in items:
                 self[r, c] = item
                 r += 1
                 if r >= nrow:
-                    c, r = r + 1, 0
+                    r, c = 0, c + 1
 
     @overload
     def __getitem__(self, index: tuple[int, int]) -> T | None: ...

plotnine/_mpl/layout_manager/_layout_tree.py

@@ -127,7 +127,12 @@ class LayoutTree:
 
     def __post_init__(self):
         self.gridspec = self.cmp.gridspec
-        self.grid = Grid["Node"](self.nrow, self.ncol, self.nodes)
+        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:

plotnine/composition/_compose.py

@@ -365,9 +365,8 @@ def _create_gridspec(self, figure, nest_into):
         from plotnine._mpl.gridspec import p9GridSpec
 
         self.layout._setup(self)
-
-        self.gridspec = p9GridSpec(
-            self.nrow, self.ncol, figure, nest_into=nest_into
+        self.gridspec = p9GridSpec.from_layout(
+            self.layout, figure=figure, nest_into=nest_into
         )
 
     def _setup(self) -> Figure:
@@ -404,7 +403,7 @@ def _make_plotspecs(
             # "subplot" in the grid. The SubplotSpec is the handle that
             # allows us to set it up for a plot or to nest another gridspec
             # in it.
-            for item, subplot_spec in zip(cmp, cmp.gridspec):  # pyright: ignore[reportArgumentType]
+            for item, subplot_spec in zip(cmp, cmp.gridspec):
                 if isinstance(item, ggplot):
                     yield plotspec(
                         item,

plotnine/composition/_plot_layout.py

@@ -16,24 +16,31 @@ class plot_layout(ComposeAddable):
     Customise the layout of plots in a composition
     """
 
-    widths: Sequence[float] | None = None
+    nrow: int | None = None
     """
-    Relative widths of each column
+    Number of rows
     """
 
-    heights: Sequence[float] | None = None
+    ncol: int | None = None
     """
-    Relative heights of each column
+    Number of columns
     """
 
-    nrow: int | None = None
+    byrow: bool | None = None
     """
-    Number of rows
+    How to place plots into the grid.
+    If None or True, they are placed row by row, left to right.
+    If False, they are placed column by column, top to bottom.
     """
 
-    ncol: int | None = None
+    widths: Sequence[float] | None = None
     """
-    Number of columns
+    Relative widths of each column
+    """
+
+    heights: Sequence[float] | None = None
+    """
+    Relative heights of each column
     """
 
     _cmp: Compose = field(init=False, repr=False)
@@ -84,6 +91,10 @@ def _setup(self, cmp: Compose):
 
         nrow, ncol = self.nrow, self.ncol
 
+        # byrow
+        if self.byrow is None:
+            self.byrow = True
+
         # setup widths & heights
         ws, hs = self.widths, self.heights
         if ws is None:
@@ -111,6 +122,8 @@ def update(self, other: plot_layout):
             self.ncol = other.ncol
         if other.nrow:
             self.nrow = other.nrow
+        if other.byrow is not None:
+            self.byrow = other.byrow
 
 
 def repeat(seq: Sequence[float], n: int) -> list[float]:

tests/baseline_images/test_plot_composition/plot_layout_byrow.png

@@ -318,3 +318,13 @@ def test_add_into_stack_error():
         (c1 + p3).draw()
 
     assert "more items than the layout rows" in str(ve.value)
+
+
+def test_plot_layout_byrow():
+    p1 = plot.red
+    p2 = plot.green
+    p3 = plot.blue
+    p4 = plot.yellow
+
+    p = (p1 + p2 + p3 + p4) + plot_layout(nrow=3, byrow=False)
+    assert p == "plot_layout_byrow"