Merge pull request #613 from shoyer/facet-dataset

Ensure xplt.FacetGrid works on Dataset objects

Author: shoyer

doc/plotting.rst

@@ -404,8 +404,9 @@ attributes, both 2d Numpy object arrays.
 
     g.name_dicts
 
-It's possible to select the :py:class:`xray.DataArray` corresponding to the FacetGrid
-through the ``name_dicts``.
+It's possible to select the :py:class:`xray.DataArray` or
+:py:class:`xray.Dataset` corresponding to the FacetGrid through the
+``name_dicts``.
 
 .. ipython:: python
 
@@ -427,6 +428,8 @@ they have been plotted.
     @savefig plot_facet_iterator.png height=12in
     plt.show()
 
+TODO: add an example of using the ``map`` method to plot dataset variables
+(e.g., with ``plt.quiver``).
 
 Maps
 ----

xray/plot/facetgrid.py

@@ -1,5 +1,6 @@
 from __future__ import division
 
+import inspect
 import warnings
 import itertools
 import functools
@@ -176,9 +177,16 @@ def __init__(self, data, col=None, row=None, col_wrap=None,
         self._col_wrap = col_wrap
         self._x_var = None
         self._y_var = None
+        self._cmap_extend = None
         self._mappables = []
 
-        self.set_titles()
+    @property
+    def _left_axes(self):
+        return self.axes[:, 0]
+
+    @property
+    def _bottom_axes(self):
+        return self.axes[-1, :]
 
     def map_dataarray(self, func, x, y, **kwargs):
         """
@@ -201,46 +209,22 @@ def map_dataarray(self, func, x, y, **kwargs):
         self : FacetGrid object
 
         """
-
         # These should be consistent with xray.plot._plot2d
         cmap_kwargs = {'plot_data': self.data.values,
-                       'vmin': None,
-                       'vmax': None,
-                       'cmap': None,
-                       'center': None,
-                       'robust': False,
-                       'extend': None,
                        # MPL default
                        'levels': 7 if 'contour' in func.__name__ else None,
                        'filled': func.__name__ != 'contour',
                        }
 
-        # Allow kwargs to override these defaults
-        # Remove cmap_kwargs from kwargs for now, we will add them back later
-        for param in list(kwargs):
-            if param in cmap_kwargs:
-                cmap_kwargs[param] = kwargs.pop(param)
+        cmap_args = inspect.getargspec(_determine_cmap_params).args
+        cmap_kwargs.update((a, kwargs[a]) for a in cmap_args if a in kwargs)
 
-        # colormap inference has to happen here since all the data in
-        # self.data is required to make the right choice
         cmap_params = _determine_cmap_params(**cmap_kwargs)
 
-        if 'contour' in func.__name__:
-            # extend is a keyword argument only for contour and contourf, but
-            # passing it to the colorbar is sufficient for imshow and
-            # pcolormesh
-            kwargs['extend'] = cmap_params['extend']
-            kwargs['levels'] = cmap_params['levels']
-
-        defaults = {
-            'add_colorbar': False,
-            'add_labels': False,
-            'norm': cmap_params.pop('cnorm'),
-        }
-
         # Order is important
-        defaults.update(cmap_params)
-        defaults.update(kwargs)
+        func_kwargs = kwargs.copy()
+        func_kwargs.update(cmap_params)
+        func_kwargs.update({'add_colorbar': False, 'add_labels': False})
 
         # Get x, y labels for the first subplot
         x, y = _infer_xy_labels(darray=self.data.loc[self.name_dicts.flat[0]],
@@ -250,40 +234,67 @@ def map_dataarray(self, func, x, y, **kwargs):
             # None is the sentinel value
             if d is not None:
                 subset = self.data.loc[d]
-                self._mappables.append(func(subset, x, y, ax=ax, **defaults))
+                mappable = func(subset, x, y, ax=ax, **func_kwargs)
+                self._mappables.append(mappable)
 
-        # Left side labels
-        for ax in self.axes[:, 0]:
-            ax.set_ylabel(y)
+        self._cmap_extend = cmap_params.get('extend')
+        self._finalize_grid(x, y)
 
-        # Bottom labels
-        for ax in self.axes[-1, :]:
-            ax.set_xlabel(x)
+        if kwargs.get('add_colorbar', True):
+            self.add_colorbar()
+
+        return self
 
+    def _finalize_grid(self, *axlabels):
+        """Finalize the annotations and layout."""
+        self.set_axis_labels(*axlabels)
+        self.set_titles()
         self.fig.tight_layout()
 
-        if self._single_group:
-            for d, ax in zip(self.name_dicts.flat, self.axes.flat):
-                if d is None:
-                    ax.set_visible(False)
+        for ax, namedict in zip(self.axes.flat, self.name_dicts.flat):
+            if namedict is None:
+                ax.set_visible(False)
 
-        # colorbar
-        if kwargs.get('add_colorbar', True):
-            cbar = self.fig.colorbar(self._mappables[-1],
-                                     ax=list(self.axes.flat),
-                                     extend=cmap_params['extend'])
+    def add_colorbar(self, **kwargs):
+        """Draw a colorbar
+        """
+        kwargs = kwargs.copy()
+        if self._cmap_extend is not None:
+            kwargs.setdefault('extend', self._cmap_extend)
+        if getattr(self.data, 'name', None) is not None:
+            kwargs.setdefault('label', self.data.name)
+        self.fig.colorbar(self._mappables[-1], ax=list(self.axes.flat),
+                          **kwargs)
+        return self
 
-            if self.data.name:
-                cbar.set_label(self.data.name, rotation=90,
-                               verticalalignment='bottom')
+    def set_axis_labels(self, x_var=None, y_var=None):
+        """Set axis labels on the left column and bottom row of the grid."""
+        if x_var is not None:
+            self._x_var = x_var
+            self.set_xlabels(x_var)
+        if y_var is not None:
+            self._y_var = y_var
+            self.set_ylabels(y_var)
+        return self
 
-        self._x_var = x
-        self._y_var = y
+    def set_xlabels(self, label=None, **kwargs):
+        """Label the x axis on the bottom row of the grid."""
+        if label is None:
+            label = self._x_var
+        for ax in self._bottom_axes:
+            ax.set_xlabel(label, **kwargs)
+        return self
 
+    def set_ylabels(self, label=None, **kwargs):
+        """Label the y axis on the left column of the grid."""
+        if label is None:
+            label = self._y_var
+        for ax in self._left_axes:
+            ax.set_ylabel(label, **kwargs)
         return self
 
     def set_titles(self, template="{coord} = {value}", maxchar=30,
-                   fontsize=_FONTSIZE, **kwargs):
+                   **kwargs):
         """
         Draw titles either above each facet or on the grid margins.
 
@@ -293,8 +304,6 @@ def set_titles(self, template="{coord} = {value}", maxchar=30,
             Template for plot titles containing {coord} and {value}
         maxchar : int
             Truncate titles at maxchar
-        fontsize : string or int
-            Passed to matplotlib.text
         kwargs : keyword args
             additional arguments to matplotlib.text
 
@@ -303,8 +312,9 @@ def set_titles(self, template="{coord} = {value}", maxchar=30,
         self: FacetGrid object
 
         """
+        import matplotlib as mpl
 
-        kwargs['fontsize'] = fontsize
+        kwargs["size"] = kwargs.pop("size", mpl.rcParams["axes.labelsize"])
 
         nicetitle = functools.partial(_nicetitle, maxchar=maxchar,
                                       template=template)
@@ -394,6 +404,10 @@ def map(self, func, *args, **kwargs):
                 data = self.data.loc[namedict]
                 plt.sca(ax)
                 innerargs = [data[a].values for a in args]
-                func(*innerargs, **kwargs)
+                # TODO: is it possible to verify that an artist is mappable?
+                mappable = func(*innerargs, **kwargs)
+                self._mappables.append(mappable)
+
+        self._finalize_grid(*args[:2])
 
         return self

xray/plot/plot.py

@@ -343,10 +343,6 @@ def newplotfunc(darray, x=None, y=None, ax=None, row=None, col=None,
             allargs = locals().copy()
             allargs.update(allargs.pop('kwargs'))
 
-            # Allows use of better FacetGrid defaults
-            assert allargs.pop('add_labels')
-            assert allargs.pop('add_colorbar')
-
             # Need the decorated plotting function
             allargs['plotfunc'] = globals()[plotfunc.__name__]
 
@@ -408,7 +404,7 @@ def newplotfunc(darray, x=None, y=None, ax=None, row=None, col=None,
             kwargs['levels'] = cmap_params['levels']
 
         # This allows the user to pass in a custom norm coming via kwargs
-        kwargs.setdefault('norm', cmap_params['cnorm'])
+        kwargs.setdefault('norm', cmap_params['norm'])
 
         ax, primitive = plotfunc(xval, yval, zval, ax=ax,
                                  cmap=cmap_params['cmap'],

xray/plot/utils.py

@@ -176,7 +176,7 @@ def _determine_cmap_params(plot_data, vmin=None, vmax=None, cmap=None,
         cmap, cnorm = _build_discrete_cmap(cmap, levels, extend, filled)
 
     return dict(vmin=vmin, vmax=vmax, cmap=cmap, extend=extend,
-                levels=levels, cnorm=cnorm)
+                levels=levels, norm=cnorm)
 
 
 def _infer_xy_labels(darray, x, y):

xray/test/test_plot.py

@@ -32,6 +32,11 @@ def text_in_fig():
     return set(alltxt)
 
 
+def find_possible_colorbars():
+    # nb. this function also matches meshes from pcolormesh
+    return plt.gcf().findobj(mpl.collections.QuadMesh)
+
+
 def substring_in_axes(substring, ax):
     '''
     Return True if a substring is found anywhere in an axes
@@ -266,15 +271,15 @@ def test_robust(self):
         self.assertEqual(cmap_params['cmap'].name, 'viridis')
         self.assertEqual(cmap_params['extend'], 'both')
         self.assertIsNone(cmap_params['levels'])
-        self.assertIsNone(cmap_params['cnorm'])
+        self.assertIsNone(cmap_params['norm'])
 
     def test_center(self):
         cmap_params = _determine_cmap_params(self.data, center=0.5)
         self.assertEqual(cmap_params['vmax'] - 0.5, 0.5 - cmap_params['vmin'])
         self.assertEqual(cmap_params['cmap'], 'RdBu_r')
         self.assertEqual(cmap_params['extend'], 'neither')
         self.assertIsNone(cmap_params['levels'])
-        self.assertIsNone(cmap_params['cnorm'])
+        self.assertIsNone(cmap_params['norm'])
 
     def test_integer_levels(self):
         data = self.data + 1
@@ -285,7 +290,7 @@ def test_integer_levels(self):
         self.assertEqual(cmap_params['cmap'].name, 'Blues')
         self.assertEqual(cmap_params['extend'], 'neither')
         self.assertEqual(cmap_params['cmap'].N, 5)
-        self.assertEqual(cmap_params['cnorm'].N, 6)
+        self.assertEqual(cmap_params['norm'].N, 6)
 
         cmap_params = _determine_cmap_params(data, levels=5,
                                              vmin=0.5, vmax=1.5)
@@ -302,7 +307,7 @@ def test_list_levels(self):
         self.assertEqual(cmap_params['vmin'], 0)
         self.assertEqual(cmap_params['vmax'], 5)
         self.assertEqual(cmap_params['cmap'].N, 5)
-        self.assertEqual(cmap_params['cnorm'].N, 6)
+        self.assertEqual(cmap_params['norm'].N, 6)
 
         for wrap_levels in [list, np.array, pd.Index, DataArray]:
             cmap_params = _determine_cmap_params(
@@ -790,9 +795,7 @@ def test_colorbar(self):
             clim = np.array(image.get_clim())
             self.assertTrue(np.allclose(expected, clim))
 
-        # There's only one colorbar
-        cbar = plt.gcf().findobj(mpl.collections.QuadMesh)
-        self.assertEqual(1, len(cbar))
+        self.assertEqual(1, len(find_possible_colorbars()))
 
     def test_empty_cell(self):
         g = xplt.FacetGrid(self.darray, col='z', col_wrap=2)
@@ -885,6 +888,44 @@ def test_map(self):
         self.g.map(plt.contourf, 'x', 'y', Ellipsis)
         self.g.map(lambda: None)
 
+    def test_map_dataset(self):
+        g = xplt.FacetGrid(self.darray.to_dataset(name='foo'), col='z')
+        g.map(plt.contourf, 'x', 'y', 'foo')
+
+        alltxt = text_in_fig()
+        for label in ['x', 'y']:
+            self.assertIn(label, alltxt)
+        # everything has a label
+        self.assertNotIn('None', alltxt)
+
+        # colorbar can't be inferred automatically
+        self.assertNotIn('foo', alltxt)
+        self.assertEqual(0, len(find_possible_colorbars()))
+
+        g.add_colorbar(label='colors!')
+        self.assertIn('colors!', text_in_fig())
+        self.assertEqual(1, len(find_possible_colorbars()))
+
+    def test_set_axis_labels(self):
+        g = self.g.map_dataarray(xplt.contourf, 'x', 'y')
+        g.set_axis_labels('longitude', 'latitude')
+        alltxt = text_in_fig()
+        for label in ['longitude', 'latitude']:
+            self.assertIn(label, alltxt)
+
+    def test_facetgrid_colorbar(self):
+        a = easy_array((10, 15, 4))
+        d = DataArray(a, dims=['y', 'x', 'z'], name='foo')
+
+        d.plot.imshow(x='x', y='y', col='z')
+        self.assertEqual(1, len(find_possible_colorbars()))
+
+        d.plot.imshow(x='x', y='y', col='z', add_colorbar=True)
+        self.assertEqual(1, len(find_possible_colorbars()))
+
+        d.plot.imshow(x='x', y='y', col='z', add_colorbar=False)
+        self.assertEqual(0, len(find_possible_colorbars()))
+
 
 class TestFacetGrid4d(PlotTestCase):