Add Facetgrid.row_labels & Facetgrid.col_labels (#3597)

* Add Facetgrid.row_labels & Facetgrid.col_labels

This allows labels to be changed later.

* Update docs.

Author: dcherian

doc/api.rst

@@ -625,7 +625,25 @@ Plotting
    plot.imshow
    plot.line
    plot.pcolormesh
+
+Faceting
+--------
+.. autosummary::
+   :toctree: generated/
+
    plot.FacetGrid
+   plot.FacetGrid.add_colorbar
+   plot.FacetGrid.add_legend
+   plot.FacetGrid.map
+   plot.FacetGrid.map_dataarray
+   plot.FacetGrid.map_dataarray_line
+   plot.FacetGrid.map_dataset
+   plot.FacetGrid.set_axis_labels
+   plot.FacetGrid.set_ticks
+   plot.FacetGrid.set_titles
+   plot.FacetGrid.set_xlabels
+   plot.FacetGrid.set_ylabels
+
 
 Testing
 =======

doc/plotting.rst

@@ -487,6 +487,7 @@ Faceting here refers to splitting an array along one or two dimensions and
 plotting each group.
 xarray's basic plotting is useful for plotting two dimensional arrays. What
 about three or four dimensional arrays? That's where facets become helpful.
+The general approach to plotting here is called “small multiples”, where the same kind of plot is repeated multiple times, and the specific use of small multiples to display the same relationship conditioned on one ore more other variables is often called a “trellis plot”.
 
 Consider the temperature data set. There are 4 observations per day for two
 years which makes for 2920 values along the time dimension.
@@ -572,8 +573,9 @@ Faceted plotting supports other arguments common to xarray 2d plots.
  FacetGrid Objects
 ===================
 
-:py:class:`xarray.plot.FacetGrid` is used to control the behavior of the
-multiple plots.
+The object returned, ``g`` in the above examples, is a :py:class:`~xarray.plot.FacetGrid`` object
+that links a :py:class:`DataArray` to a matplotlib figure with a particular structure.
+This object can be used to control the behavior of the multiple plots.
 It borrows an API and code from `Seaborn's FacetGrid
 <http://seaborn.pydata.org/tutorial/axis_grids.html>`_.
 The structure is contained within the ``axes`` and ``name_dicts``
@@ -609,6 +611,13 @@ they have been plotted.
     @savefig plot_facet_iterator.png
     plt.draw()
 
+
+:py:class:`~xarray.FacetGrid` objects have methods that let you customize the automatically generated
+axis labels, axis ticks and plot titles. See :py:meth:`~xarray.plot.FacetGrid.set_titles`,
+:py:meth:`~xarray.plot.FacetGrid.set_xlabels`, :py:meth:`~xarray.plot.FacetGrid.set_ylabels` and
+:py:meth:`~xarray.plot.FacetGrid.set_ticks` for more information.
+Plotting functions can be applied to each subset of the data by calling :py:meth:`~xarray.plot.FacetGrid.map_dataarray` or to each subplot by calling :py:meth:`FacetGrid.map`.
+
 TODO: add an example of using the ``map`` method to plot dataset variables
 (e.g., with ``plt.quiver``).
 

doc/whats-new.rst

@@ -39,6 +39,12 @@ Bug fixes
   <https://github.com/spencerkclark>`_.
 - Fix plotting with transposed 2D non-dimensional coordinates. (:issue:`3138`, :pull:`3441`)
   By `Deepak Cherian <https://github.com/dcherian>`_.
+- :py:meth:`~xarray.plot.FacetGrid.set_titles` can now replace existing row titles of a
+  :py:class:`~xarray.plot.FacetGrid` plot. In addition :py:class:`~xarray.plot.FacetGrid` gained
+  two new attributes: :py:attr:`~xarray.plot.FacetGrid.col_labels` and
+  :py:attr:`~xarray.plot.FacetGrid.row_labels` contain matplotlib Text handles for both column and
+  row labels. These can be used to manually change the labels.
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 - Fix issue with Dask-backed datasets raising a ``KeyError`` on some computations involving ``map_blocks`` (:pull:`3598`)
   By `Tom Augspurger <https://github.com/TomAugspurger>`_.
 

xarray/plot/facetgrid.py

@@ -61,6 +61,10 @@ class FacetGrid:
     axes : numpy object array
         Contains axes in corresponding position, as returned from
         plt.subplots
+    col_labels : list
+        list of :class:`matplotlib.text.Text` instances corresponding to column titles.
+    row_labels : list
+        list of :class:`matplotlib.text.Text` instances corresponding to row titles.
     fig : matplotlib.Figure
         The figure containing all the axes
     name_dicts : numpy object array
@@ -200,6 +204,8 @@ def __init__(
         self._ncol = ncol
         self._col_var = col
         self._col_wrap = col_wrap
+        self.row_labels = [None] * nrow
+        self.col_labels = [None] * ncol
         self._x_var = None
         self._y_var = None
         self._cmap_extend = None
@@ -482,22 +488,32 @@ def set_titles(self, template="{coord} = {value}", maxchar=30, size=None, **kwar
                     ax.set_title(title, size=size, **kwargs)
         else:
             # The row titles on the right edge of the grid
-            for ax, row_name in zip(self.axes[:, -1], self.row_names):
+            for index, (ax, row_name, handle) in enumerate(
+                zip(self.axes[:, -1], self.row_names, self.row_labels)
+            ):
                 title = nicetitle(coord=self._row_var, value=row_name, maxchar=maxchar)
-                ax.annotate(
-                    title,
-                    xy=(1.02, 0.5),
-                    xycoords="axes fraction",
-                    rotation=270,
-                    ha="left",
-                    va="center",
-                    **kwargs,
-                )
+                if not handle:
+                    self.row_labels[index] = ax.annotate(
+                        title,
+                        xy=(1.02, 0.5),
+                        xycoords="axes fraction",
+                        rotation=270,
+                        ha="left",
+                        va="center",
+                        **kwargs,
+                    )
+                else:
+                    handle.set_text(title)
 
             # The column titles on the top row
-            for ax, col_name in zip(self.axes[0, :], self.col_names):
+            for index, (ax, col_name, handle) in enumerate(
+                zip(self.axes[0, :], self.col_names, self.col_labels)
+            ):
                 title = nicetitle(coord=self._col_var, value=col_name, maxchar=maxchar)
-                ax.set_title(title, size=size, **kwargs)
+                if not handle:
+                    self.col_labels[index] = ax.set_title(title, size=size, **kwargs)
+                else:
+                    handle.set_text(title)
 
         return self
 

xarray/tests/test_plot.py

@@ -62,6 +62,15 @@ def substring_in_axes(substring, ax):
     return False
 
 
+def substring_not_in_axes(substring, ax):
+    """
+    Return True if a substring is not found anywhere in an axes
+    """
+    alltxt = {t.get_text() for t in ax.findobj(mpl.text.Text)}
+    check = [(substring not in txt) for txt in alltxt]
+    return all(check)
+
+
 def easy_array(shape, start=0, stop=1):
     """
     Make an array with desired shape using np.linspace
@@ -1776,6 +1785,18 @@ def test_default_labels(self):
         for label, ax in zip(self.darray.coords["col"].values, g.axes[0, :]):
             assert substring_in_axes(label, ax)
 
+        # ensure that row & col labels can be changed
+        g.set_titles("abc={value}")
+        for label, ax in zip(self.darray.coords["row"].values, g.axes[:, -1]):
+            assert substring_in_axes(f"abc={label}", ax)
+            # previous labels were "row=row0" etc.
+            assert substring_not_in_axes("row=", ax)
+
+        for label, ax in zip(self.darray.coords["col"].values, g.axes[0, :]):
+            assert substring_in_axes(f"abc={label}", ax)
+            # previous labels were "col=row0" etc.
+            assert substring_not_in_axes("col=", ax)
+
 
 @pytest.mark.filterwarnings("ignore:tight_layout cannot")
 class TestFacetedLinePlotsLegend(PlotTestCase):