Merge pull request #27952 from AnsonTran/align-titles

rcomer · web-flow · commit 23ea9098c510 · 2024-04-04T20:29:07.000+01:00
ENH: Align titles
diff --git a/doc/api/figure_api.rst b/doc/api/figure_api.rst
@@ -71,6 +71,7 @@ Annotating
    Figure.align_labels
    Figure.align_xlabels
    Figure.align_ylabels
+   Figure.align_titles
    Figure.autofmt_xdate
 
 
@@ -264,6 +265,7 @@ Annotating
    SubFigure.align_labels
    SubFigure.align_xlabels
    SubFigure.align_ylabels
+   SubFigure.align_titles
 
 Adding and getting Artists
 --------------------------
diff --git a/doc/users/next_whats_new/figure_align_titles.rst b/doc/users/next_whats_new/figure_align_titles.rst
@@ -0,0 +1,7 @@
+subplot titles can now be automatically aligned
+-----------------------------------------------
+
+Subplot axes titles can be misaligned vertically if tick labels or
+xlabels are placed at the top of one subplot. The new method on the
+`.Figure` class: `.Figure.align_titles` will now align the titles
+vertically.
diff --git a/galleries/examples/subplots_axes_and_figures/align_labels_demo.py b/galleries/examples/subplots_axes_and_figures/align_labels_demo.py
@@ -1,37 +1,43 @@
 """
-===============
-Aligning Labels
-===============
+==========================
+Aligning Labels and Titles
+==========================
 
-Aligning xlabel and ylabel using `.Figure.align_xlabels` and
-`.Figure.align_ylabels`
+Aligning xlabel, ylabel, and title using `.Figure.align_xlabels`,
+`.Figure.align_ylabels`, and `.Figure.align_titles`.
 
-`.Figure.align_labels` wraps these two functions.
+`.Figure.align_labels` wraps the x and y label functions.
 
 Note that the xlabel "XLabel1 1" would normally be much closer to the
-x-axis, and "YLabel1 0" would be much closer to the y-axis of their
-respective axes.
+x-axis, "YLabel0 0" would be much closer to the y-axis, and title
+"Title0 0" would be much closer to the top of their respective axes.
 """
 import matplotlib.pyplot as plt
 import numpy as np
 
-import matplotlib.gridspec as gridspec
+fig, axs = plt.subplots(2, 2, layout='constrained')
 
-fig = plt.figure(tight_layout=True)
-gs = gridspec.GridSpec(2, 2)
-
-ax = fig.add_subplot(gs[0, :])
+ax = axs[0][0]
 ax.plot(np.arange(0, 1e6, 1000))
-ax.set_ylabel('YLabel0')
-ax.set_xlabel('XLabel0')
+ax.set_title('Title0 0')
+ax.set_ylabel('YLabel0 0')
+
+ax = axs[0][1]
+ax.plot(np.arange(1., 0., -0.1) * 2000., np.arange(1., 0., -0.1))
+ax.set_title('Title0 1')
+ax.xaxis.tick_top()
+ax.tick_params(axis='x', rotation=55)
+
 
 for i in range(2):
-    ax = fig.add_subplot(gs[1, i])
+    ax = axs[1][i]
     ax.plot(np.arange(1., 0., -0.1) * 2000., np.arange(1., 0., -0.1))
     ax.set_ylabel('YLabel1 %d' % i)
     ax.set_xlabel('XLabel1 %d' % i)
     if i == 0:
         ax.tick_params(axis='x', rotation=55)
+
 fig.align_labels()  # same as fig.align_xlabels(); fig.align_ylabels()
+fig.align_titles()
 
 plt.show()
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2985,8 +2985,13 @@ def _update_title_position(self, renderer):
 
         titles = (self.title, self._left_title, self._right_title)
 
-        # Need to check all our twins too, and all the children as well.
-        axs = self._twinned_axes.get_siblings(self) + self.child_axes
+        # Need to check all our twins too, aligned axes, and all the children
+        # as well.
+        axs = set()
+        axs.update(self.child_axes)
+        axs.update(self._twinned_axes.get_siblings(self))
+        axs.update(self.figure._align_label_groups['title'].get_siblings(self))
+
         for ax in self.child_axes:  # Child positions must be updated first.
             locator = ax.get_axes_locator()
             ax.apply_aspect(locator(self, renderer) if locator else None)
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -132,10 +132,14 @@ def __init__(self, **kwargs):
         self._supxlabel = None
         self._supylabel = None
 
-        # groupers to keep track of x and y labels we want to align.
-        # see self.align_xlabels and self.align_ylabels and
-        # axis._get_tick_boxes_siblings
-        self._align_label_groups = {"x": cbook.Grouper(), "y": cbook.Grouper()}
+        # groupers to keep track of x, y labels and title we want to align.
+        # see self.align_xlabels, self.align_ylabels,
+        # self.align_titles, and axis._get_tick_boxes_siblings
+        self._align_label_groups = {
+            "x": cbook.Grouper(),
+            "y": cbook.Grouper(),
+            "title": cbook.Grouper()
+        }
 
         self._localaxes = []  # track all Axes
         self.artists = []
@@ -1293,7 +1297,7 @@ def subplots_adjust(self, left=None, bottom=None, right=None, top=None,
 
     def align_xlabels(self, axs=None):
         """
-        Align the xlabels of subplots in the same subplot column if label
+        Align the xlabels of subplots in the same subplot row if label
         alignment is being done automatically (i.e. the label position is
         not manually set).
 
@@ -1314,6 +1318,7 @@ def align_xlabels(self, axs=None):
         See Also
         --------
         matplotlib.figure.Figure.align_ylabels
+        matplotlib.figure.Figure.align_titles
         matplotlib.figure.Figure.align_labels
 
         Notes
@@ -1375,6 +1380,7 @@ def align_ylabels(self, axs=None):
         See Also
         --------
         matplotlib.figure.Figure.align_xlabels
+        matplotlib.figure.Figure.align_titles
         matplotlib.figure.Figure.align_labels
 
         Notes
@@ -1412,6 +1418,53 @@ def align_ylabels(self, axs=None):
                         # grouper for groups of ylabels to align
                         self._align_label_groups['y'].join(ax, axc)
 
+    def align_titles(self, axs=None):
+        """
+        Align the titles of subplots in the same subplot row if title
+        alignment is being done automatically (i.e. the title position is
+        not manually set).
+
+        Alignment persists for draw events after this is called.
+
+        Parameters
+        ----------
+        axs : list of `~matplotlib.axes.Axes`
+            Optional list of (or ndarray) `~matplotlib.axes.Axes`
+            to align the titles.
+            Default is to align all Axes on the figure.
+
+        See Also
+        --------
+        matplotlib.figure.Figure.align_xlabels
+        matplotlib.figure.Figure.align_ylabels
+        matplotlib.figure.Figure.align_labels
+
+        Notes
+        -----
+        This assumes that ``axs`` are from the same `.GridSpec`, so that
+        their `.SubplotSpec` positions correspond to figure positions.
+
+        Examples
+        --------
+        Example with titles::
+
+            fig, axs = plt.subplots(1, 2)
+            axs[0].set_aspect('equal')
+            axs[0].set_title('Title 0')
+            axs[1].set_title('Title 1')
+            fig.align_titles()
+        """
+        if axs is None:
+            axs = self.axes
+        axs = [ax for ax in np.ravel(axs) if ax.get_subplotspec() is not None]
+        for ax in axs:
+            _log.debug(' Working on: %s', ax.get_title())
+            rowspan = ax.get_subplotspec().rowspan
+            for axc in axs:
+                rowspanc = axc.get_subplotspec().rowspan
+                if (rowspan.start == rowspanc.start):
+                    self._align_label_groups['title'].join(ax, axc)
+
     def align_labels(self, axs=None):
         """
         Align the xlabels and ylabels of subplots with the same subplots
@@ -1430,8 +1483,8 @@ def align_labels(self, axs=None):
         See Also
         --------
         matplotlib.figure.Figure.align_xlabels
-
         matplotlib.figure.Figure.align_ylabels
+        matplotlib.figure.Figure.align_titles
         """
         self.align_xlabels(axs=axs)
         self.align_ylabels(axs=axs)
diff --git a/lib/matplotlib/figure.pyi b/lib/matplotlib/figure.pyi
@@ -161,6 +161,7 @@ class FigureBase(Artist):
     ) -> None: ...
     def align_xlabels(self, axs: Iterable[Axes] | None = ...) -> None: ...
     def align_ylabels(self, axs: Iterable[Axes] | None = ...) -> None: ...
+    def align_titles(self, axs: Iterable[Axes] | None = ...) -> None: ...
     def align_labels(self, axs: Iterable[Axes] | None = ...) -> None: ...
     def add_gridspec(self, nrows: int = ..., ncols: int = ..., **kwargs) -> GridSpec: ...
     @overload
diff --git a/lib/matplotlib/tests/baseline_images/test_figure/figure_align_titles_constrained.png b/lib/matplotlib/tests/baseline_images/test_figure/figure_align_titles_constrained.png
diff --git a/lib/matplotlib/tests/baseline_images/test_figure/figure_align_titles_tight.png b/lib/matplotlib/tests/baseline_images/test_figure/figure_align_titles_tight.png
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -66,6 +66,32 @@ def test_align_labels():
     fig.align_labels()
 
 
+@image_comparison(['figure_align_titles_tight.png',
+                   'figure_align_titles_constrained.png'],
+                  tol=0 if platform.machine() == 'x86_64' else 0.022,
+                  style='mpl20')
+def test_align_titles():
+    for layout in ['tight', 'constrained']:
+        fig, axs = plt.subplots(1, 2, layout=layout, width_ratios=[2, 1])
+
+        ax = axs[0]
+        ax.plot(np.arange(0, 1e6, 1000))
+        ax.set_title('Title0 left', loc='left')
+        ax.set_title('Title0 center', loc='center')
+        ax.set_title('Title0 right', loc='right')
+
+        ax = axs[1]
+        ax.plot(np.arange(0, 1e4, 100))
+        ax.set_title('Title1')
+        ax.set_xlabel('Xlabel0')
+        ax.xaxis.set_label_position("top")
+        ax.xaxis.tick_top()
+        for tick in ax.get_xticklabels():
+            tick.set_rotation(90)
+
+        fig.align_titles()
+
+
 def test_align_labels_stray_axes():
     fig, axs = plt.subplots(2, 2)
     for nn, ax in enumerate(axs.flat):
