Refactor plot utils

xarray/plot/facetgrid.py

@@ -9,8 +9,8 @@
 from ..core.formatting import format_item
 from ..core.pycompat import getargspec
 from .utils import (
-    _determine_cmap_params, _infer_xy_labels, import_matplotlib_pyplot,
-    label_from_attrs)
+    _determine_cmap_params, _infer_line_data, _infer_xy_labels,
+    import_matplotlib_pyplot, label_from_attrs)
 
 # Overrides axes.labelsize, xtick.major.size, ytick.major.size
 # from mpl.rcParams
@@ -221,95 +221,75 @@ def map_dataarray(self, func, x, y, **kwargs):
 
         """
 
-        cmapkw = kwargs.get('cmap')
-        colorskw = kwargs.get('colors')
-        cbar_kwargs = kwargs.pop('cbar_kwargs', {})
-        cbar_kwargs = {} if cbar_kwargs is None else dict(cbar_kwargs)
+        if func.__name__ == 'line':
+            add_legend = kwargs.pop('add_legend', True)
+            kwargs['add_legend'] = False
+            func_kwargs = kwargs.copy()
+            func_kwargs['_labels'] = False
 
-        if kwargs.get('cbar_ax', None) is not None:
-            raise ValueError('cbar_ax not supported by FacetGrid.')
+        else:
+            cmapkw = kwargs.get('cmap')
+            colorskw = kwargs.get('colors')
+            cbar_kwargs = kwargs.pop('cbar_kwargs', {})
+            cbar_kwargs = {} if cbar_kwargs is None else dict(cbar_kwargs)
+
+            if kwargs.get('cbar_ax', None) is not None:
+                raise ValueError('cbar_ax not supported by FacetGrid.')
 
-        # colors is mutually exclusive with cmap
-        if cmapkw and colorskw:
-            raise ValueError("Can't specify both cmap and colors.")
+            # colors is mutually exclusive with cmap
+            if cmapkw and colorskw:
+                raise ValueError("Can't specify both cmap and colors.")
 
-        # These should be consistent with xarray.plot._plot2d
-        cmap_kwargs = {'plot_data': self.data.values,
-                       # MPL default
-                       'levels': 7 if 'contour' in func.__name__ else None,
-                       'filled': func.__name__ != 'contour',
-                       }
+            # These should be consistent with xarray.plot._plot2d
+            cmap_kwargs = {'plot_data': self.data.values,
+                           # MPL default
+                           'levels': 7 if 'contour' in func.__name__ else None,
+                           'filled': func.__name__ != 'contour',
+                           }
 
-        cmap_args = getargspec(_determine_cmap_params).args
-        cmap_kwargs.update((a, kwargs[a]) for a in cmap_args if a in kwargs)
+            cmap_args = getargspec(_determine_cmap_params).args
+            cmap_kwargs.update((a, kwargs[a])
+                               for a in cmap_args if a in kwargs)
 
-        cmap_params = _determine_cmap_params(**cmap_kwargs)
+            cmap_params = _determine_cmap_params(**cmap_kwargs)
 
-        if colorskw is not None:
-            cmap_params['cmap'] = None
+            if colorskw is not None:
+                cmap_params['cmap'] = None
 
-        # Order is important
-        func_kwargs = kwargs.copy()
-        func_kwargs.update(cmap_params)
-        func_kwargs.update({'add_colorbar': False, 'add_labels': False})
+            # Order is important
+            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]], x=x, y=y,
-            imshow=func.__name__ == 'imshow', rgb=kwargs.get('rgb', None))
+            # Get x, y labels for the first subplot
+            x, y = _infer_xy_labels(
+                darray=self.data.loc[self.name_dicts.flat[0]], x=x, y=y,
+                imshow=func.__name__ == 'imshow', rgb=kwargs.get('rgb', None))
 
         for d, ax in zip(self.name_dicts.flat, self.axes.flat):
             # None is the sentinel value
             if d is not None:
                 subset = self.data.loc[d]
-                mappable = func(subset, x, y, ax=ax, **func_kwargs)
+                mappable = func(subset, x=x, y=y, ax=ax, **func_kwargs)
                 self._mappables.append(mappable)
 
-        self._cmap_extend = cmap_params.get('extend')
-        self._finalize_grid(x, y)
-
-        if kwargs.get('add_colorbar', True):
-            self.add_colorbar(**cbar_kwargs)
-
-        return self
-
-    def map_dataarray_line(self, x=None, y=None, hue=None, **kwargs):
-        """
-        Apply a line plot to a 2d facet subset of the data.
-
-        Parameters
-        ----------
-        x, y, hue: string
-            dimension names for the axes and hues of each facet
-
-        Returns
-        -------
-        self : FacetGrid object
-
-        """
-        from .plot import line, _infer_line_data
-
-        add_legend = kwargs.pop('add_legend', True)
-        kwargs['add_legend'] = False
+        if func.__name__ == 'line':
+            _, _, hueplt, xlabel, ylabel, huelabel = _infer_line_data(
+                darray=self.data.loc[self.name_dicts.flat[0]],
+                x=x, y=y, hue=func_kwargs['hue'])
 
-        for d, ax in zip(self.name_dicts.flat, self.axes.flat):
-            # None is the sentinel value
-            if d is not None:
-                subset = self.data.loc[d]
-                mappable = line(subset, x=x, y=y, hue=hue,
-                                ax=ax, _labels=False,
-                                **kwargs)
-                self._mappables.append(mappable)
-        _, _, hueplt, xlabel, ylabel, huelabel = _infer_line_data(
-            darray=self.data.loc[self.name_dicts.flat[0]],
-            x=x, y=y, hue=hue)
+            self._hue_var = hueplt
+            self._hue_label = huelabel
+            self._finalize_grid(xlabel, ylabel)
 
-        self._hue_var = hueplt
-        self._hue_label = huelabel
-        self._finalize_grid(xlabel, ylabel)
+            if add_legend and hueplt is not None and huelabel is not None:
+                self.add_legend()
+        else:
+            self._cmap_extend = cmap_params.get('extend')
+            self._finalize_grid(x, y)
 
-        if add_legend and hueplt is not None and huelabel is not None:
-            self.add_legend()
+            if kwargs.get('add_colorbar', True):
+                self.add_colorbar(**cbar_kwargs)
 
         return self
 
@@ -522,3 +502,29 @@ def map(self, func, *args, **kwargs):
         self._finalize_grid(*args[:2])
 
         return self
+
+
+def _easy_facetgrid(data, plotfunc, x=None, y=None, row=None,
+                    col=None, col_wrap=None, sharex=True, sharey=True,
+                    aspect=None, size=None, subplot_kws=None, **kwargs):
+    """
+    Convenience method to call xarray.plot.FacetGrid from 2d plotting methods
+
+    kwargs are the arguments to 2d plotting method
+    """
+    ax = kwargs.pop('ax', None)
+    figsize = kwargs.pop('figsize', None)
+    if ax is not None:
+        raise ValueError("Can't use axes when making faceted plots.")
+    if aspect is None:
+        aspect = 1
+    if size is None:
+        size = 3
+    elif figsize is not None:
+        raise ValueError('cannot provide both `figsize` and `size` arguments')
+
+    g = FacetGrid(data=data, col=col, row=row, col_wrap=col_wrap,
+                  sharex=sharex, sharey=sharey, figsize=figsize,
+                  aspect=aspect, size=size, subplot_kws=subplot_kws)
+
+    return g.map_dataarray(plotfunc, x, y, **kwargs)