Simplify FacetGrid.map_dataarary

Author: shoyer

xray/plot/facetgrid.py

@@ -1,5 +1,6 @@
 from __future__ import division
 
+import inspect
 import warnings
 import itertools
 import functools
@@ -208,46 +209,26 @@ 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)
+        defaults.update(cmap_params)
 
         # Get x, y labels for the first subplot
         x, y = _infer_xy_labels(darray=self.data.loc[self.name_dicts.flat[0]],

xray/plot/plot.py

@@ -408,7 +408,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

@@ -270,15 +270,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
@@ -289,7 +289,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)
@@ -306,7 +306,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(