ENH: Scatter plots of one variable vs another

xarray/core/dataset.py

@@ -33,6 +33,7 @@
     Frozen, SortedKeysDict, either_dict_or_kwargs, decode_numpy_dict_values,
     ensure_us_time_resolution, hashable, maybe_wrap_array)
 from .variable import IndexVariable, Variable, as_variable, broadcast_variables
+from ..plot.plot import dataset_scatter
 
 # list of attributes of pd.DatetimeIndex that are ndarrays of time info
 _DATETIMEINDEX_COMPONENTS = ['year', 'month', 'day', 'hour', 'minute',
@@ -3592,6 +3593,13 @@ def real(self):
     def imag(self):
         return self._unary_op(lambda x: x.imag, keep_attrs=True)(self)
 
+    def scatter(self,  x=None, y=None, hue=None, col=None, row=None,
+                    col_wrap=None, sharex=True, sharey=True, aspect=None,
+                    size=None, subplot_kws=None,  add_legend=True, **kwargs):
+        lc = locals()
+        ds = lc.pop('self')
+        return dataset_scatter(ds=ds, **lc)
+
     def filter_by_attrs(self, **kwargs):
         """Returns a ``Dataset`` with variables that match specific conditions.
 

xarray/plot/facetgrid.py

@@ -266,7 +266,7 @@ def map_dataarray(self, func, x, y, **kwargs):
 
         return self
 
-    def map_dataarray_line(self, x=None, y=None, hue=None, **kwargs):
+    def map_dataarray_line(self, plotfunc, x=None, y=None, hue=None, **kwargs):
         """
         Apply a line plot to a 2d facet subset of the data.
 
@@ -280,7 +280,8 @@ def map_dataarray_line(self, x=None, y=None, hue=None, **kwargs):
         self : FacetGrid object
 
         """
-        from .plot import line, _infer_line_data
+        from .plot import (_infer_line_data, _infer_scatter_data,
+                          line, dataset_scatter)
 
         add_legend = kwargs.pop('add_legend', True)
         kwargs['add_legend'] = False
@@ -289,13 +290,19 @@ def map_dataarray_line(self, x=None, y=None, hue=None, **kwargs):
             # None is the sentinel value
             if d is not None:
                 subset = self.data.loc[d]
-                mappable = line(subset, x=x, y=y, hue=hue,
+                mappable = plotfunc(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)
+
+        if plotfunc == line:
+            _, _, hueplt, xlabel, ylabel, huelabel = _infer_line_data(
+                darray=self.data.loc[self.name_dicts.flat[0]],
+                x=x, y=y, hue=hue)
+        elif plotfunc == dataset_scatter:
+            _, _, hueplt, xlabel, ylabel, huelabel = _infer_scatter_data(
+                ds=self.data.loc[self.name_dicts.flat[0]],
+                x=x, y=y, hue=hue)
 
         self._hue_var = hueplt
         self._hue_label = huelabel

xarray/plot/plot.py

@@ -104,7 +104,7 @@ def _line_facetgrid(darray, row=None, col=None, hue=None,
     g = FacetGrid(data=darray, 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_line(hue=hue, **kwargs)
+    return g.map_dataarray_line(line, hue=hue, **kwargs)
 
 
 def plot(darray, row=None, col=None, col_wrap=None, ax=None, hue=None,
@@ -240,6 +240,47 @@ def _infer_line_data(darray, x, y, hue):
     return xplt, yplt, hueplt, xlabel, ylabel, huelabel
 
 
+def _infer_scatter_data(ds, x, y, hue):
+    dvars = set(ds.data_vars.keys())
+    error_msg = (' must be either one of ({0:s})'
+                 .format(', '.join(dvars)))
+
+    if x not in dvars:
+        raise ValueError(x + error_msg)
+
+    if y not in dvars:
+        raise ValueError(y + error_msg)
+
+    dims = ds[x].dims
+    if ds[y].dims != dims:
+        raise ValueError('{} and {} must have the same dimensions.'
+                         ''.format(x, y))
+
+    if hue is not None and hue not in dims:
+        raise ValueError(hue + 'must be either one of ({0:s})'
+                         ''.format(', '.join(dims)))
+
+    dims = set(dims)
+    if hue:
+        dims.remove(hue)
+        xplt = ds[x].stack(stackdim=dims).transpose('stackdim', hue).values
+        yplt = ds[y].stack(stackdim=dims).transpose('stackdim', hue).values
+    else:
+        xplt = ds[x].values.flatten()
+        yplt = ds[y].values.flatten()
+
+    if hue:
+        hueplt = ds[x].coords[hue]
+        huelabel = label_from_attrs(ds[x][hue])
+    else:
+        hueplt = None
+        huelabel = None
+
+    xlabel = label_from_attrs(ds[x])
+    ylabel = label_from_attrs(ds[y])
+    return xplt, yplt, hueplt, xlabel, ylabel, huelabel
+
+
 # This function signature should not change so that it can use
 # matplotlib format strings
 def line(darray, *args, **kwargs):
@@ -289,6 +330,7 @@ def line(darray, *args, **kwargs):
     if row or col:
         allargs = locals().copy()
         allargs.update(allargs.pop('kwargs'))
+        allargs.update(allargs.pop('args'))
         return _line_facetgrid(**allargs)
 
     ndims = len(darray.dims)
@@ -309,8 +351,6 @@ def line(darray, *args, **kwargs):
     yincrease = kwargs.pop('yincrease', True)
     add_legend = kwargs.pop('add_legend', True)
     _labels = kwargs.pop('_labels', True)
-    if args is ():
-        args = kwargs.pop('args', ())
 
     ax = get_axis(figsize, size, aspect, ax)
     xplt, yplt, hueplt, xlabel, ylabel, huelabel = \
@@ -873,7 +913,7 @@ def _is_monotonic(coord, axis=0):
         return np.all(delta_pos) or np.all(delta_neg)
 
 
-def _infer_interval_breaks(coord, axis=0, check_monotonic=False):
+def _infer_interval_breaks(coord, axis=0):
     """
     >>> _infer_interval_breaks(np.arange(5))
     array([-0.5,  0.5,  1.5,  2.5,  3.5,  4.5])
@@ -883,7 +923,7 @@ def _infer_interval_breaks(coord, axis=0, check_monotonic=False):
     """
     coord = np.asarray(coord)
 
-    if check_monotonic and not _is_monotonic(coord, axis=axis):
+    if not _is_monotonic(coord, axis=axis):
         raise ValueError("The input coordinate is not sorted in increasing "
                          "order along axis %d. This can lead to unexpected "
                          "results. Consider calling the `sortby` method on "
@@ -922,8 +962,8 @@ def pcolormesh(x, y, z, ax, infer_intervals=None, **kwargs):
 
     if infer_intervals:
         if len(x.shape) == 1:
-            x = _infer_interval_breaks(x, check_monotonic=True)
-            y = _infer_interval_breaks(y, check_monotonic=True)
+            x = _infer_interval_breaks(x)
+            y = _infer_interval_breaks(y)
         else:
             # we have to infer the intervals on both axes
             x = _infer_interval_breaks(x, axis=1)
@@ -941,3 +981,44 @@ def pcolormesh(x, y, z, ax, infer_intervals=None, **kwargs):
         ax.set_ylim(y[0], y[-1])
 
     return primitive
+
+
+def dataset_scatter(ds, x=None, y=None, hue=None, col=None, row=None,
+                    col_wrap=None, sharex=True, sharey=True, aspect=None,
+                    size=None, subplot_kws=None,  add_legend=True, **kwargs):
+    if col or row:
+        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=ds, 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_line(x=x, y=y, hue=hue, plotfunc=dataset_scatter, **kwargs)
+
+    xplt, yplt, hueplt, xlabel, ylabel, huelabel = _infer_scatter_data(ds, x, y, hue)
+
+    figsize = kwargs.pop('figsize', None)
+    ax = kwargs.pop('ax', None)
+    ax = get_axis(figsize, size, aspect, ax)
+    primitive = ax.plot(xplt, yplt, '.')
+
+    if kwargs.get('_labels', True):
+        if xlabel is not None:
+            ax.set_xlabel(xlabel)
+
+        if ylabel is not None:
+            ax.set_ylabel(ylabel)
+
+    if add_legend and huelabel is not None:
+        ax.legend(handles=primitive,
+                  labels=list(hueplt.values),
+                  title=huelabel)
+    return primitive