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_PlotMethods
 
 # list of attributes of pd.DatetimeIndex that are ndarrays of time info
 _DATETIMEINDEX_COMPONENTS = ['year', 'month', 'day', 'hour', 'minute',
@@ -3592,6 +3593,16 @@ def real(self):
     def imag(self):
         return self._unary_op(lambda x: x.imag, keep_attrs=True)(self)
 
+    @property
+    def plot(self):
+        """
+        Access plotting functions. Use it as a namespace to use
+        xarray.plot functions as Dataset methods
+        >>> ds.plot.scatter(...)  # equivalent to xarray.plot.scatter(ds,...)
+
+        """
+        return _Dataset_PlotMethods(self)
+
     def filter_by_attrs(self, **kwargs):
         """Returns a ``Dataset`` with variables that match specific conditions.
 

xarray/plot/facetgrid.py

@@ -280,7 +280,7 @@ 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, line
 
         add_legend = kwargs.pop('add_legend', True)
         kwargs['add_legend'] = False
@@ -293,9 +293,10 @@ def map_dataarray_line(self, x=None, y=None, hue=None, **kwargs):
                                 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)
+                darray=self.data.loc[self.name_dicts.flat[0]],
+                x=x, y=y, hue=hue)
 
         self._hue_var = hueplt
         self._hue_label = huelabel
@@ -306,6 +307,35 @@ def map_dataarray_line(self, x=None, y=None, hue=None, **kwargs):
 
         return self
 
+    def map_scatter(self, x=None, y=None, hue=None, discrete_legend=False,
+                    add_legend=None, **kwargs):
+        from .plot import _infer_scatter_meta_data, scatter
+
+        kwargs['add_legend'] = False
+        kwargs['discrete_legend'] = discrete_legend
+        meta_data = _infer_scatter_meta_data(self.data, x, y, hue,
+                                             add_legend, discrete_legend)
+        kwargs['_meta_data'] = meta_data
+        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 = scatter(subset, x=x, y=y, hue=hue,
+                                   ax=ax, **kwargs)
+                self._mappables.append(mappable)
+
+        self._finalize_grid(meta_data['xlabel'], meta_data['ylabel'])
+
+        if hue and meta_data['add_legend']:
+            self._hue_label = meta_data.pop('hue_label', None)
+            if meta_data['discrete_legend']:
+                self._hue_var = meta_data['hue_values']
+                self.add_legend()
+            else:
+                self.add_colorbar(label=self._hue_label)
+
+        return self
+
     def _finalize_grid(self, *axlabels):
         """Finalize the annotations and layout."""
         self.set_axis_labels(*axlabels)

xarray/plot/plot.py

@@ -204,7 +204,7 @@ def _infer_line_data(darray, x, y, hue):
         dim, = darray.dims  # get the only dimension name
         huename = None
         hueplt = None
-        huelabel = ''
+        hue_label = ''
 
         if (x is None and y is None) or x == dim:
             xplt = darray.coords[dim]
@@ -232,12 +232,84 @@ def _infer_line_data(darray, x, y, hue):
             yplt = darray.coords[yname]
 
         hueplt = darray.coords[huename]
-        huelabel = label_from_attrs(darray[huename])
+        hue_label = label_from_attrs(darray[huename])
 
     xlabel = label_from_attrs(xplt)
     ylabel = label_from_attrs(yplt)
 
-    return xplt, yplt, hueplt, xlabel, ylabel, huelabel
+    return xplt, yplt, hueplt, xlabel, ylabel, hue_label
+
+
+def _ensure_numeric(arr):
+    numpy_types = [np.floating, np.integer]
+    return _valid_numpy_subdtype(arr, numpy_types)
+
+
+def _infer_scatter_meta_data(ds, x, y, hue, add_legend, discrete_legend):
+    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)
+
+    if hue and add_legend is None:
+        add_legend = True
+    if add_legend and not hue:
+            raise ValueError('hue must be speicifed for generating a lengend')
+
+    if hue and not _ensure_numeric(ds[hue].values):
+        if discrete_legend is None:
+            discrete_legend = True
+        elif discrete_legend is False:
+            raise TypeError('Cannot create a colorbar for a non numeric'
+                            ' coordinate')
+
+    dims = ds[x].dims
+    if ds[y].dims != dims:
+        raise ValueError('{} and {} must have the same dimensions.'
+                         ''.format(x, y))
+
+    dims_coords = set(list(ds.coords) + list(ds.dims))
+    if hue is not None and hue not in dims_coords:
+        raise ValueError(hue + ' must be either one of ({0:s})'
+                               ''.format(', '.join(dims_coords)))
+
+    if hue:
+        hue_label = label_from_attrs(ds.coords[hue])
+    else:
+        hue_label = None
+
+    return {'add_legend': add_legend,
+            'discrete_legend': discrete_legend,
+            'hue_label': hue_label,
+            'xlabel': label_from_attrs(ds[x]),
+            'ylabel': label_from_attrs(ds[y]),
+            'hue_values': ds[x].coords[hue] if discrete_legend else None}
+
+
+def _infer_scatter_data(ds, x, y, hue, discrete_legend):
+    dims = set(ds[x].dims)
+    if discrete_legend:
+        dims.remove(hue)
+        xplt = ds[x].stack(stackdim=dims).transpose('stackdim', hue).values
+        yplt = ds[y].stack(stackdim=dims).transpose('stackdim', hue).values
+        return {'x': xplt, 'y': yplt}
+    else:
+        data = {'x': ds[x].values.flatten(),
+                'y': ds[y].values.flatten(),
+                'color': None}
+        if hue:
+            # this is a hack to make a dataarray of the shape of ds[x] whose
+            # values are the coordinate hue. There's probably a better way
+            color = ds[x]
+            color[:] = 0
+            color += ds.coords[hue]
+            data['color'] = color.values.flatten()
+        return data
 
 
 # This function signature should not change so that it can use
@@ -313,7 +385,7 @@ def line(darray, *args, **kwargs):
         args = kwargs.pop('args', ())
 
     ax = get_axis(figsize, size, aspect, ax)
-    xplt, yplt, hueplt, xlabel, ylabel, huelabel = \
+    xplt, yplt, hueplt, xlabel, ylabel, hue_label = \
         _infer_line_data(darray, x, y, hue)
 
     _ensure_plottable(xplt)
@@ -332,7 +404,7 @@ def line(darray, *args, **kwargs):
     if darray.ndim == 2 and add_legend:
         ax.legend(handles=primitive,
                   labels=list(hueplt.values),
-                  title=huelabel)
+                  title=hue_label)
 
     # Rotate dates on xlabels
     # Do this without calling autofmt_xdate so that x-axes ticks
@@ -941,3 +1013,82 @@ def pcolormesh(x, y, z, ax, infer_intervals=None, **kwargs):
         ax.set_ylim(y[0], y[-1])
 
     return primitive
+
+
+def scatter(ds, x, y, hue=None, col=None, row=None,
+            col_wrap=None, sharex=True, sharey=True, aspect=None,
+            size=None, subplot_kws=None, add_legend=None,
+            discrete_legend=None, **kwargs):
+
+    if kwargs.get('_meta_data', None):
+        discrete_legend = kwargs['_meta_data']['discrete_legend']
+    else:
+        meta_data = _infer_scatter_meta_data(ds, x, y, hue,
+                                             add_legend, discrete_legend)
+        discrete_legend = meta_data['discrete_legend']
+        add_legend = meta_data['add_legend']
+
+    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_scatter(x=x, y=y, hue=hue, add_legend=add_legend,
+                             discrete_legend=discrete_legend, **kwargs)
+
+    data = _infer_scatter_data(ds, x, y, hue, discrete_legend)
+
+    figsize = kwargs.pop('figsize', None)
+    ax = kwargs.pop('ax', None)
+    ax = get_axis(figsize, size, aspect, ax)
+    if discrete_legend:
+        primitive = ax.plot(data['x'], data['y'], '.')
+    else:
+        primitive = ax.scatter(data['x'], data['y'], c=data['color'])
+    if '_meta_data' in kwargs:  # if this was called from map_scatter,
+        return primitive        # finish here. Else, make labels
+
+    if meta_data.get('xlabel', None):
+        ax.set_xlabel(meta_data.get('xlabel'))
+
+    if meta_data.get('ylabel', None):
+        ax.set_ylabel(meta_data.get('ylabel'))
+    if add_legend and discrete_legend:
+        ax.legend(handles=primitive,
+                  labels=list(meta_data['hue_values'].values),
+                  title=meta_data.get('hue_label', None))
+    if add_legend and not discrete_legend:
+        cbar = ax.figure.colorbar(primitive)
+        if meta_data.get('hue_label', None):
+            cbar.ax.set_ylabel(meta_data.get('hue_label'))
+
+    return primitive
+
+
+class _Dataset_PlotMethods(object):
+    """
+    Enables use of xarray.plot functions as attributes on a Dataset.
+    For example, Dataset.plot.scatter
+    """
+
+    def __init__(self, dataset):
+        self._ds = dataset
+
+    def __call__(self, *args, **kwargs):
+        raise ValueError('Dataset.plot cannot be called directly. Use'
+                         'an explicit plot method, e.g. ds.plot.scatter(...)')
+
+    @functools.wraps(scatter)
+    def scatter(self, *args, **kwargs):
+        return scatter(self._ds, *args, **kwargs)

xarray/tests/test_plot.py

@@ -8,7 +8,7 @@
 import pytest
 
 import xarray.plot as xplt
-from xarray import DataArray
+from xarray import DataArray, Dataset
 from xarray.coding.times import _import_cftime
 from xarray.plot.plot import _infer_interval_breaks
 from xarray.plot.utils import (
@@ -1622,6 +1622,86 @@ def test_wrong_num_of_dimensions(self):
             self.darray.plot.line(row='row', hue='hue')
 
 
+class TestScatterPlots(PlotTestCase):
+    def setUp(self):
+        das = [DataArray(np.random.randn(3, 3, 4, 4),
+                          dims=['x', 'row', 'col', 'hue'],
+                          coords=[range(k) for k in [3, 3, 4, 4]])
+                for _ in [1, 2]]
+        ds = Dataset({'A': das[0], 'B': das[1]})
+        ds.hue.name = 'huename'
+        ds.hue.attrs['units'] = 'hunits'
+        ds.x.attrs['units'] = 'xunits'
+        ds.col.attrs['units'] = 'colunits'
+        ds.row.attrs['units'] = 'rowunits'
+        ds.A.attrs['units'] = 'Aunits'
+        ds.B.attrs['units'] = 'Bunits'
+        self.ds = ds
+
+    def test_facetgrid_shape(self):
+        g = self.ds.plot.scatter(x='A', y='B', row='row', col='col')
+        assert g.axes.shape == (len(self.ds.row), len(self.ds.col))
+
+        g = self.ds.plot.scatter(x='A', y='B', row='col', col='row')
+        assert g.axes.shape == (len(self.ds.col), len(self.ds.row))
+
+    def test_default_labels(self):
+        g = self.ds.plot.scatter('A', 'B', row='row', col='col', hue='hue')
+        # Rightmost column should be labeled
+        for label, ax in zip(self.ds.coords['row'].values, g.axes[:, -1]):
+            assert substring_in_axes(label, ax)
+
+        # Top row should be labeled
+        for label, ax in zip(self.ds.coords['col'].values, g.axes[0, :]):
+            assert substring_in_axes(str(label), ax)
+
+        # Bottom row should have name of x array name and units
+        for ax in g.axes[-1, :]:
+            assert ax.get_xlabel() == 'A [Aunits]'
+
+        # Leftmost column should have name of y array name and units
+        for ax in g.axes[:, 0]:
+            assert ax.get_ylabel() == 'B [Bunits]'
+
+    def test_both_x_and_y(self):
+        with pytest.raises(ValueError):
+            self.darray.plot.line(row='row', col='col',
+                                  x='x', y='hue')
+
+    def test_axes_in_faceted_plot(self):
+        with pytest.raises(ValueError):
+            self.ds.plot.scatter(x='A', y='B', row='row', ax=plt.axes())
+
+    def test_figsize_and_size(self):
+        with pytest.raises(ValueError):
+            self.ds.plot.scatter(x='A', y='B', row='row', size=3, figsize=4)
+
+    def test_bad_args(self):
+        with pytest.raises(ValueError):
+            self.ds.plot.scatter(x='A', y='B', add_legend=True)
+            self.ds.plot.scatter(x='A', y='The Spanish Inquisition')
+            self.ds.plot.scatter(x='The Spanish Inquisition', y='B')
+
+    def test_non_numeric_legened(self):
+        self.ds['hue'] = pd.date_range('2000-01-01', periods=4)
+        lines = self.ds.plot.scatter(x='A', y='B', hue='hue')
+        # should make a discrete legend
+        assert lines[0].axes.legend_ is not None
+        # and raise an error if explicitly not allowed to do so
+        with pytest.raises(ValueError):
+            self.ds.plot.scatter(x='A', y='B', hue='hue',
+                                 discrete_legend=False)
+
+    def test_add_legened_by_default(self):
+        sc = self.ds.plot.scatter(x='A', y='B', hue='hue')
+        assert len(sc.figure.axes) == 2
+
+    def test_not_same_dimensions(self):
+        self.ds['A'] = self.ds['A'].isel(x=0)
+        with pytest.raises(ValueError):
+            self.ds.plot.scatter(x='A', y='B')
+
+
 class TestDatetimePlot(PlotTestCase):
     def setUp(self):
         '''