Align axis titles in a composition

Author: has2k1

plotnine/_mpl/layout_manager/_layout_tree.py

@@ -118,6 +118,7 @@ def harmonise(self):
         """
         Align and resize plots in composition to look good
         """
+        self.align_axis_titles()
         self.align()
         self.resize()
 
@@ -147,6 +148,20 @@ def align_sub_compositions(self):
         for tree in self.sub_compositions:
             tree.align()
 
+    @abc.abstractmethod
+    def align_axis_titles(self):
+        """
+        Align the axis titles along the composing dimension
+
+        Since the alignment value used to for this purpose is one of
+        the fields in the _side_space, it affects the space created
+        for the panel.
+
+        We could align the titles within self.align but we would have
+        to store the value outside the _side_space and pick it up when
+        setting the position of the texts!
+        """
+
     def resize_sub_compositions(self):
         """
         Resize panels in the compositions contained in this one
@@ -365,6 +380,50 @@ def top_tags_align(self) -> bool:
         arr = np.array(self.top_tag_heights)
         return all(arr == arr[0])
 
+    @property
+    def left_axis_titles_align(self) -> bool:
+        """
+        Return True if the left axis titles align
+        """
+        arr = np.array(self.left_axis_title_clearances)
+        return all(arr == arr[0])
+
+    @property
+    def bottom_axis_titles_align(self) -> bool:
+        """
+        Return True if the bottom axis titles align
+        """
+        arr = np.array(self.bottom_axis_title_clearances)
+        return all(arr == arr[0])
+
+    @cached_property
+    @abc.abstractmethod
+    def left_axis_title_clearance(self) -> float:
+        """
+        Distance between the left y-axis title and the panel
+        """
+
+    @cached_property
+    @abc.abstractmethod
+    def bottom_axis_title_clearance(self) -> float:
+        """
+        Distance between the left x-axis title and the panel
+        """
+
+    @cached_property
+    def left_axis_title_clearances(self) -> list[float]:
+        """
+        Distances between the left y-axis titles and the panels
+        """
+        return [node.left_axis_title_clearance for node in self.nodes]
+
+    @cached_property
+    def bottom_axis_title_clearances(self) -> list[float]:
+        """
+        Distances between the bottom x-axis titles and the panels
+        """
+        return [node.bottom_axis_title_clearance for node in self.nodes]
+
     @abc.abstractmethod
     def set_left_margin_alignment(self, value: float):
         """
@@ -429,6 +488,22 @@ def set_top_tag_alignment(self, value: float):
         In figure dimenstions
         """
 
+    @abc.abstractmethod
+    def set_left_axis_title_alignment(self, value: float):
+        """
+        Set the space to align the left axis titles in this composition
+
+        In figure dimenstions
+        """
+
+    @abc.abstractmethod
+    def set_bottom_axis_title_alignment(self, value: float):
+        """
+        Set the space to align the bottom axis titles in this composition
+
+        In figure dimenstions
+        """
+
 
 @dataclass
 class ColumnsTree(LayoutTree):
@@ -458,6 +533,11 @@ def align(self):
         self.align_panel_bottoms()
         self.align_sub_compositions()
 
+    def align_axis_titles(self):
+        self.align_bottom_axis_titles()
+        for tree in self.sub_compositions:
+            tree.align_axis_titles()
+
     def resize(self):
         """
         Resize the widths of gridspec so that panels have equal widths
@@ -552,6 +632,23 @@ def align_top_tags(self):
             else:
                 item.set_top_tag_alignment(value)
 
+    def align_bottom_axis_titles(self):
+        if self.bottom_axis_titles_align:
+            pass
+
+        values = max(self.bottom_axis_title_clearances) - np.array(
+            self.bottom_axis_title_clearances
+        )
+        # We ignore 0 values since they can undo values
+        # set to align this composition with an outer one.
+        for item, value in zip(self.nodes, values):
+            if value == 0:
+                continue
+            if isinstance(item, LayoutSpaces):
+                item.b.axis_title_alignment = value
+            else:
+                item.set_bottom_axis_title_alignment(value)
+
     @cached_property
     def panel_lefts(self):
         left_item = self.nodes[0]
@@ -620,6 +717,14 @@ def bottom_tag_height(self) -> float:
     def top_tag_height(self) -> float:
         return max(self.top_tag_heights)
 
+    @cached_property
+    def left_axis_title_clearance(self) -> float:
+        return self.left_axis_title_clearances[0]
+
+    @cached_property
+    def bottom_axis_title_clearance(self) -> float:
+        return max(self.bottom_axis_title_clearances)
+
     def set_left_margin_alignment(self, value: float):
         left_item = self.nodes[0]
         if isinstance(left_item, LayoutSpaces):
@@ -662,6 +767,20 @@ def set_top_tag_alignment(self, value: float):
             else:
                 item.set_top_tag_alignment(value)
 
+    def set_bottom_axis_title_alignment(self, value: float):
+        for item in self.nodes:
+            if isinstance(item, LayoutSpaces):
+                item.b.axis_title_alignment = value
+            else:
+                item.set_bottom_axis_title_alignment(value)
+
+    def set_left_axis_title_alignment(self, value: float):
+        left_item = self.nodes[0]
+        if isinstance(left_item, LayoutSpaces):
+            left_item.l.axis_title_alignment = value
+        else:
+            left_item.set_left_axis_title_alignment(value)
+
 
 @dataclass
 class RowsTree(LayoutTree):
@@ -688,6 +807,11 @@ def align(self):
         self.align_panel_rights()
         self.align_sub_compositions()
 
+    def align_axis_titles(self):
+        self.align_left_axis_titles()
+        for tree in self.sub_compositions:
+            tree.align_axis_titles()
+
     def resize(self):
         """
         Resize the heights of gridspec so that panels have equal heights
@@ -804,6 +928,21 @@ def align_right_tags(self):
             else:
                 item.set_right_tag_alignment(value)
 
+    def align_left_axis_titles(self):
+        if self.left_axis_titles_align:
+            pass
+
+        values = max(self.left_axis_title_clearances) - np.array(
+            self.left_axis_title_clearances
+        )
+        for item, value in zip(self.nodes, values):
+            if value == 0:
+                continue
+            if isinstance(item, LayoutSpaces):
+                item.l.axis_title_alignment = value
+            else:
+                item.set_left_axis_title_alignment(value)
+
     @cached_property
     def panel_lefts(self):
         values = []
@@ -872,6 +1011,14 @@ def top_tag_height(self) -> float:
     def bottom_tag_height(self) -> float:
         return self.bottom_tag_heights[-1]
 
+    @cached_property
+    def left_axis_title_clearance(self) -> float:
+        return max(self.left_axis_title_clearances)
+
+    @cached_property
+    def bottom_axis_title_clearance(self) -> float:
+        return self.bottom_axis_title_clearances[-1]
+
     def set_left_margin_alignment(self, value: float):
         for item in self.nodes:
             if isinstance(item, LayoutSpaces):
@@ -913,3 +1060,17 @@ def set_right_tag_alignment(self, value: float):
                 item.r.tag_alignment = value
             else:
                 item.set_right_tag_alignment(value)
+
+    def set_left_axis_title_alignment(self, value: float):
+        for item in self.nodes:
+            if isinstance(item, LayoutSpaces):
+                item.l.axis_title_alignment = value
+            else:
+                item.set_left_axis_title_alignment(value)
+
+    def set_bottom_axis_title_alignment(self, value: float):
+        bottom_item = self.nodes[-1]
+        if isinstance(bottom_item, LayoutSpaces):
+            bottom_item.b.axis_title_alignment = value
+        else:
+            bottom_item.set_bottom_axis_title_alignment(value)

plotnine/_mpl/layout_manager/_spaces.py

@@ -16,6 +16,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, cast
 
+from plotnine.exceptions import PlotnineError
 from plotnine.facets import facet_grid, facet_null, facet_wrap
 
 from ._layout_items import LayoutItems
@@ -233,7 +234,7 @@ def tag_width(self) -> float:
         """
         The width of the tag including the margins
 
-        The value is zero expect if all these are true:
+        The value is zero except if all these are true:
             - The tag is in the margin `theme(plot_tag_position = "margin")`
             - The tag at one one of the the following locations;
               left, right, topleft, topright, bottomleft or bottomright
@@ -245,13 +246,41 @@ def tag_height(self) -> float:
         """
         The height of the tag including the margins
 
-        The value is zero expect if all these are true:
+        The value is zero except if all these are true:
             - The tag is in the margin `theme(plot_tag_position = "margin")`
             - The tag at one one of the the following locations;
               top, bottom, topleft, topright, bottomleft or bottomright
         """
         return 0
 
+    @property
+    def axis_title_clearance(self) -> float:
+        """
+        The distance between the axis title and the panel
+
+         Figure
+         ----------------------------
+        |         Panel              |
+        |         -----------        |
+        |        |           |       |
+        |        |           |       |
+        |  Y<--->|           |       |
+        |        |           |       |
+        |        |           |       |
+        |         -----------        |
+        |                            |
+         ----------------------------
+
+        We use this value to when aligning axis titles in a
+        plot composition.
+        """
+
+        try:
+            return self.total - self.sum_upto("axis_title_alignment")
+        except AttributeError as err:
+            # There is probably an error in in the layout manager
+            raise PlotnineError("Side has no axis title") from err
+
 
 @dataclass
 class left_spaces(_side_spaces):
@@ -311,6 +340,14 @@ class left_spaces(_side_spaces):
     axis_title_y_margin_left: float = 0
     axis_title_y: float = 0
     axis_title_y_margin_right: float = 0
+    axis_title_alignment: float = 0
+    """
+    Space added to align the axis title with others in a composition
+
+    This value is calculated during the layout process. The amount is
+    the difference between the largest and smallest axis_title_clearance
+    among the items in the composition.
+    """
     axis_text_y_margin_left: float = 0
     axis_text_y: float = 0
     axis_text_y_margin_right: float = 0
@@ -674,6 +711,15 @@ class bottom_spaces(_side_spaces):
     axis_title_x_margin_bottom: float = 0
     axis_title_x: float = 0
     axis_title_x_margin_top: float = 0
+    axis_title_alignment: float = 0
+    """
+    Space added to align the axis title with others in a composition
+
+    This value is calculated during the layout process in a tree structure
+    that has convenient access to the sides/edges of the panels in the
+    composition. It's amount is the difference in height between this axis
+    text (and it's margins) and the tallest axis text (and it's margin).
+    """
     axis_text_x_margin_bottom: float = 0
     axis_text_x: float = 0
     axis_text_x_margin_top: float = 0
@@ -945,6 +991,24 @@ def bottom_tag_height(self) -> float:
         """
         return self.b.tag_height
 
+    @property
+    def left_axis_title_clearance(self) -> float:
+        """
+        Distance between the left y-axis title and the panel
+
+        In figure dimensions.
+        """
+        return self.l.axis_title_clearance
+
+    @property
+    def bottom_axis_title_clearance(self) -> float:
+        """
+        Distance between the bottom x-axis title and the panel
+
+        In figure dimensions.
+        """
+        return self.b.axis_title_clearance
+
     def increase_horizontal_plot_margin(self, dw: float):
         """
         Increase the plot_margin to the right & left of the panels

tests/baseline_images/test_plot_composition/complex_composition.png

@@ -11,7 +11,7 @@
 from matplotlib.testing.compare import compare_images
 
 from plotnine import ggplot, theme
-from plotnine.composition import Arrange
+from plotnine.composition import Arrange, Beside, Stack
 from plotnine.themes.theme import DEFAULT_RCPARAMS
 
 TOLERANCE = 2  # Default tolerance for the tests
@@ -271,4 +271,8 @@ def composition_equals(cmp: Arrange, name: str) -> bool:
     return not err
 
 
+# Note that, dataclass subclasses have their own __eq__ and not that of the
+# parent class.
 Arrange.__eq__ = composition_equals  # pyright: ignore[reportAttributeAccessIssue]
+Beside.__eq__ = composition_equals  # pyright: ignore[reportAttributeAccessIssue]
+Stack.__eq__ = composition_equals  # pyright: ignore[reportAttributeAccessIssue]