Reimplement resizing of plot compositions

Author: has2k1

plotnine/_mpl/gridspec.py

@@ -243,6 +243,20 @@ def bbox(self):
         """
         return TransformedBbox(self.bbox_relative, self.figure.transSubfigure)
 
+    @property
+    def width(self) -> float:
+        """
+        Width of bbox in figure space
+        """
+        return self.bbox_relative.width
+
+    @property
+    def height(self) -> float:
+        """
+        Height of bbox in figure space
+        """
+        return self.bbox_relative.height
+
     def to_transform(self) -> Transform:
         """
         Return transform of this gridspec

plotnine/_mpl/layout_manager/_layout_tree.py

@@ -258,55 +258,53 @@ def right_most_spaces(self) -> list[right_spaces]:
     @property
     def panel_width(self) -> float:
         """
-        A representative width for panels of the nodes
+        A width of all panels in this composition
         """
         return sum(self.panel_widths)
 
     @property
     def panel_height(self) -> float:
         """
-        A representative height for panels of the nodes
+        A height of all panels in this composition
         """
         return sum(self.panel_heights)
 
     @property
     def plot_width(self) -> float:
         """
-        A representative for width for plots of the nodes
+        A width of all plots in this tree/composition
         """
-        return sum(self.plot_widths)
+        return self.gridspec.width
 
     @property
     def plot_height(self) -> float:
         """
-        A representative for height for plots of the nodes
+        A height of all plots in this tree/composition
         """
-        return sum(self.plot_heights)
+        return self.gridspec.height
 
     @property
     def panel_widths(self) -> Sequence[float]:
         """
         Widths [figure space] of the panels along horizontal dimension
-
-        For each column, the effective panel width is the width of the
-        shortest panel.
         """
-        n = self.ncol
+        # 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 if node else 1 / n for node in col])
+            max(node.panel_width for node in col if node) if any(col) else w
             for col in self.grid.iter_cols()
         ]
 
     @property
     def panel_heights(self) -> Sequence[float]:
         """
         Heights [figure space] of the panels along vertical dimension
-
-        For each row, the representative height is that of the shortest panel.
         """
-        n = self.nrow
+        h = self.plot_height / self.nrow
         return [
-            max([node.panel_height if node else 1 / n for node in row])
+            max([node.panel_height for node in row if node]) if any(row) else h
             for row in self.grid.iter_rows()
         ]
 
@@ -315,11 +313,11 @@ def plot_widths(self) -> Sequence[float]:
         """
         Widths [figure space] of the plots along horizontal dimension
 
-        For each column, the representative width is that of the shortest plot.
+        For each column, the representative width is that of the widest plot.
         """
-        n = self.ncol
+        w = self.gridspec.width / self.ncol
         return [
-            max([node.plot_width if node else 1 / n for node in col])
+            max([node.plot_width if node else w for node in col])
             for col in self.grid.iter_cols()
         ]
 
@@ -328,14 +326,32 @@ def plot_heights(self) -> Sequence[float]:
         """
         Heights [figure space] of the plots along vertical dimension
 
-        For each row, the representative height is that of the shortest plot.
+        For each row, the representative height is that of the tallest plot.
         """
-        n = self.nrow
+        h = self.gridspec.height / self.nrow
         return [
-            max([node.plot_height if node else 1 / n for node in row])
+            max([node.plot_height if node else h for node in row])
             for row in self.grid.iter_rows()
         ]
 
+    @property
+    def panel_width_ratios(self) -> Sequence[float]:
+        """
+        The relative widths of the panels in the composition
+
+        These are normalised to have a mean = 1.
+        """
+        return cast("Sequence[float]", self.cmp._layout.widths)
+
+    @property
+    def panel_height_ratios(self) -> Sequence[float]:
+        """
+        The relative heights of the panels in the composition
+
+        These are normalised to have a mean = 1.
+        """
+        return cast("Sequence[float]", self.cmp._layout.heights)
+
     def bottom_spaces_in_row(self, r: int) -> list[bottom_spaces]:
         spaces: list[bottom_spaces] = []
         for node in self.grid[r, :]:
@@ -509,21 +525,22 @@ def align_axis_titles(self):
             tree.align_axis_titles()
 
     def resize_widths(self):
-        n = self.ncol
-        resize_ratios = np.array(self.cmp._layout.widths)
-        base_panel_widths = np.ones(n) * self.panel_width
-        scaled_panel_widths = base_panel_widths * resize_ratios
+        # The scaling calcuation to get the new panel width is
+        # straight-forward because the ratios have a mean of 1.
+        # So the multiplication preserves the total panel width.
+        new_panel_widths = np.mean(self.panel_widths) * np.array(
+            self.panel_width_ratios
+        )
         non_panel_space = np.array(self.plot_widths) - self.panel_widths
-        new_plot_widths = scaled_panel_widths + non_panel_space
+        new_plot_widths = new_panel_widths + non_panel_space
         width_ratios = new_plot_widths / new_plot_widths.max()
         self.gridspec.set_width_ratios(width_ratios)
 
     def resize_heights(self):
-        n = self.nrow
-        resize_ratios = np.array(self.cmp._layout.heights)
-        base_panel_heights = np.ones(n) * self.panel_height
-        scaled_panel_heights = base_panel_heights * resize_ratios
+        new_panel_heights = np.mean(self.panel_heights) * np.array(
+            self.panel_height_ratios
+        )
         non_panel_space = np.array(self.plot_heights) - self.panel_heights
-        new_plot_heights = scaled_panel_heights + non_panel_space
+        new_plot_heights = new_panel_heights + non_panel_space
         height_ratios = new_plot_heights / new_plot_heights.max()
         self.gridspec.set_height_ratios(height_ratios)

plotnine/_mpl/layout_manager/_spaces.py

@@ -924,14 +924,14 @@ def plot_width(self) -> float:
         """
         Width [figure dimensions] of the whole plot
         """
-        return self.plot._gridspec.bbox_relative.width
+        return float(self.plot._gridspec.width)
 
     @property
     def plot_height(self) -> float:
         """
         Height [figure dimensions] of the whole plot
         """
-        return self.plot._gridspec.bbox_relative.height
+        return float(self.plot._gridspec.height)
 
     @property
     def panel_width(self) -> float:

plotnine/composition/_plot_layout.py

@@ -86,6 +86,11 @@ def repeat(seq: Sequence[float], n: int) -> list[float]:
     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]
+def normalise(seq: Sequence[float]) -> list[float]:
+    """
+    Normalise seq so that the mean is 1
+    """
+    mean = sum(seq) / len(seq)
+    if mean == 0:
+        raise ValueError("Cannot rescale: mean is zero")
+    return [x / mean for x in seq]

tests/baseline_images/test_plot_composition/plot_layout_extra_col_width.png

@@ -1,4 +1,11 @@
-from plotnine import element_text, facet_grid, facet_wrap, theme, theme_gray
+from plotnine import (
+    element_text,
+    facet_grid,
+    facet_wrap,
+    labs,
+    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
@@ -179,3 +186,68 @@ def test_plot_layout_nested_resize():
 
     p = (((p1 | p2) + ws) / ((p3 | p4) + ws)) + hs
     assert p == "plot_layout_nested_resize"
+
+
+def test_plot_layout_extra_cols():
+    p1 = plot.red
+    p2 = plot.green
+    p3 = plot.blue
+    p = (p1 | p2 | p3) + plot_layout(ncol=5)
+    assert p == "plot_layout_extra_cols"
+
+
+def test_plot_layout_extra_col_width():
+    # An extra column is extactly panel_width wide (no margin)
+    # By stacking two rows, where one has an extra column, we can
+    # confirm the size.
+    p1 = plot.red
+    p2 = plot.green
+    p3 = plot.blue
+    p4 = plot.yellow + labs(y="") + theme(plot_margin=0)
+
+    c1 = (p1 | p2 | p3) + plot_layout(ncol=4)
+    c2 = p1 | p2 | p3 | p4
+    p = c1 / c2
+    assert p == "plot_layout_extra_col_width"
+
+
+def test_plot_layout_extra_rows():
+    p1 = plot.red
+    p2 = plot.green
+    p3 = plot.blue
+    p = (p1 / p2 / p3) + plot_layout(nrow=5)
+    assert p == "plot_layout_extra_rows"
+
+
+def test_plot_layout_extra_row_width():
+    # An extra row is extactly panel_width wide (no margin)
+    # By stacking two rows, where one has an extra row, we can
+    # confirm the size.
+    p1 = plot.red
+    p2 = plot.green
+    p3 = plot.blue
+    p4 = plot.yellow + labs(x="", title="") + theme(plot_margin=0)
+
+    c1 = (p1 / p2 / p3) + plot_layout(nrow=4)
+    c2 = p1 / p2 / p3 / p4
+    p = c1 | c2
+    assert p == "plot_layout_extra_row_width"
+
+
+def test_wrap_complicated():
+    p1 = plot.red
+    p2 = plot.green
+    p3 = plot.blue + g.points
+    p4 = plot.yellow
+    p5 = plot.cyan
+    p6 = plot.orange
+    p7 = plot.purple + g.cols
+
+    c1 = (p1 + p2 + p3) + plot_layout(ncol=2)
+    c2 = (p4 + p5 + p6 + p7) + plot_layout(ncol=4, widths=[1, 2, 4, 8])
+
+    # The top composition has two rows and the bottom one has one.
+    # With [1, 2] height ratios, the panels in each row should have
+    # the same height.
+    p = (c1 / c2) + plot_layout(heights=[2, 1])
+    assert p == "wrap_complicated"