Repository: pydata/xarray

Issue:
New Feature: Add FacetGrid functionality to plot.line()
When data is more than 2 dimensional, `plot()` creates a `FacetGrid` where each facet gets plots two dimensional data. I would find it very useful to have similar functionality for line plots. 

Let's define three `DataArray`s, of 2,3 and 4 dimensions, named `d2`, `d3` and `d4`.

```python
d4=xr.DataArray(np.random.randn(3,4,5,6),
                dims=['x','y','z','w'],
                coords=[range(3),range(4),range(5),range(6)])
d3=d4[...,0].drop('w')
d2=d3[...,0].drop('z')

d4.plot(row='x', col='y')   #creates a 2d facet grid, each facet is a pcolormesh
d3.plot(col='x')           #creates a 1d facet grid, each facet is a pcolormesh
plt.show()
```

It would be very useful to have a similar interface for `line` plots too. That is, I would expect
```python
d2.plot.line(row='y', x='x')
```
to create a row of panels, each having one line. Similarly, I would expect
```python
d3.plot.line(col='y', x='x')
#or alternatively
d3.plot.line(col='y', hue='z')
```
to create a row of panels, each having a series of lines, with `hue` corresponding to `z` and so forth.  Finally, 
```python
d4.plot.line(row='z',col='y', x='x')
#or alternatively
d4.plot.line(row='z',col='y', hue='w')
```
should create a 2d matrix of panels, each having a series of lines, with `hue` corresponding to `w`.

This seems to me like a natural generalization of the behaviour of `plot()`, and I found myself coding manually loops to do that a few times.


Discussion:
yohai: Example code to demonstrate functionality:

```
import xarray as xr
import numpy as np
import pandas as pd
np.random.seed(0)

#set up data:
tm=pd.date_range('2000-01-01',end='2000-01-02',freq='H')
d4=xr.DataArray(np.random.randn(len(tm),4,3,3),
                dims=['time', 'y', 'z', 'w'],
                coords=[tm,['foo', 'bar', 'foobar', 'barfoo'], range(3), range(3)],
                name='foo')
d3 = d4[..., 0].drop('w')
d2 = d3[..., 0].drop('z')

#generate a faceted 4-dimensional plot:
d4.plot(row='w', col='y', hue='z')
#or equivalently:
d4.plot.line(row='w', col='y', hue='z')
#or equivalently:
d4.plot.line(row='w', col='y', x='time')

#also works for plotting along y axis:
d3.plot.line(col='y', y='time')
```
shoyer: I'm really excited about this!

Could you kindly paste in examples of what the generated plots look like in your comment? (GitHub will upload an image when you paste it in.)
yohai: was simpler than I thought. @shoyer I added the plots.
jhamman: @yohai - this looks great. Can we add some basic unit tests covering this functionality? 
yohai: @jhamman @shoyer Sorry for being a noob, but what exactly do you mean by "add some basic unit tests"? you mean writing new functions that will verify that stuff works as it should, or that the new addition does not mess up previous plotting behavior?
fmaussion: > writing new functions that will verify that stuff works as it should

Yes. See [here](https://github.com/pydata/xarray/blob/master/xarray/tests/test_plot.py#L1223) for example tests for FacetGrid. Testing plotting functions is not trivial but we usually test only very
basic functionality. You might find some examples in here.

> or that the new addition does not mess up previous plotting behavior?

This is usually covered by continuous integration, which already ran the tests on your current PR
(the "all checks have passed" green button above). So on this side you are safe for now, unless
our previous tests weren't efficient enough ;-)
yohai: Thanks @shoyer. I fixed everything except for the coverall part, which is weird. The lines it marks as untested are the function `_line_facetgrid`, which the main function that performs all of the testing. I suspect that it simply doesn't run my new tests for some reason. When I run `pytest` locally these lines do get tested (I verified by adding `print` commands there).
Should I add decorators like `pytest.mark.something`for coverall to notice these tests?
shoyer: > Thanks @shoyer. I fixed everything except for the coverall part, which is weird. The lines it marks as untested are the function _line_facetgrid, which the main function that performs all of the testing. I suspect that it simply doesn't run my new tests for some reason. When I run pytest locally these lines do get tested (I verified by adding print commands there).
Should I add decorators like pytest.mark.somethingfor coverall to notice these tests?

The specific problematic lines are those marked red in coveralls, e.g.,
https://coveralls.io/builds/17253745/source?filename=xarray%2Fplot%2Fplot.py#L101
https://coveralls.io/builds/17253745/source?filename=xarray%2Fplot%2Fplot.py#L176
https://coveralls.io/builds/17253745/source?filename=xarray%2Fplot%2Fplot.py#L291

These are all cases where an error is raised.
yohai: OK, somehow the rebasing thing broke everything and I am at a loss. @shoyer, @dcherian please guide me on how to proceed. Sorry.
shoyer: It looks like the merge went OK, but you have a couple of tests that are failing. Click on the "Travis-CI" link above next to the red "X" to see details:
https://travis-ci.org/pydata/xarray/jobs/386973465

Patch:
diff --git a/doc/plotting.rst b/doc/plotting.rst
index fa364d48..54fa2f57 100644
--- a/doc/plotting.rst
+++ b/doc/plotting.rst
@@ -208,7 +208,11 @@ It is required to explicitly specify either
 2. ``hue``: the dimension you want to represent by multiple lines.
 
 Thus, we could have made the previous plot by specifying ``hue='lat'`` instead of ``x='time'``.
-If required, the automatic legend can be turned off using ``add_legend=False``.
+If required, the automatic legend can be turned off using ``add_legend=False``. Alternatively,
+``hue`` can be passed directly to :py:func:`xarray.plot` as `air.isel(lon=10, lat=[19,21,22]).plot(hue='lat')`.
+
+
+
 
 Dimension along y-axis
 ~~~~~~~~~~~~~~~~~~~~~~
@@ -218,7 +222,7 @@ It is also possible to make line plots such that the data are on the x-axis and
 .. ipython:: python
 
     @savefig plotting_example_xy_kwarg.png
-    air.isel(time=10, lon=[10, 11]).plot.line(y='lat', hue='lon')
+    air.isel(time=10, lon=[10, 11]).plot(y='lat', hue='lon')
 
 Changing Axes Direction
 -----------------------
@@ -452,6 +456,13 @@ arguments to the xarray plotting methods/functions. This returns a
     @savefig plot_facet_dataarray.png
     g_simple = t.plot(x='lon', y='lat', col='time', col_wrap=3)
 
+Faceting also works for line plots.
+
+.. ipython:: python
+
+    @savefig plot_facet_dataarray_line.png
+    g_simple_line = t.isel(lat=slice(0,None,4)).plot(x='lon', hue='lat', col='time', col_wrap=3)
+
 4 dimensional
 ~~~~~~~~~~~~~
 
diff --git a/xarray/plot/facetgrid.py b/xarray/plot/facetgrid.py
index 361b4726..771f0879 100644
--- a/xarray/plot/facetgrid.py
+++ b/xarray/plot/facetgrid.py
@@ -5,7 +5,6 @@ import itertools
 import warnings
 
 import numpy as np
-
 from ..core.formatting import format_item
 from ..core.pycompat import getargspec
 from .utils import (
@@ -267,6 +266,44 @@ class FacetGrid(object):
 
         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
+
+        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)
+        _, _, _, xlabel, ylabel, huelabel = _infer_line_data(
+            darray=self.data.loc[self.name_dicts.flat[0]],
+            x=x, y=y, hue=hue)
+
+        self._finalize_grid(xlabel, ylabel)
+
+        if add_legend and huelabel:
+            self.add_line_legend(huelabel)
+
+        return self
+
     def _finalize_grid(self, *axlabels):
         """Finalize the annotations and layout."""
         self.set_axis_labels(*axlabels)
@@ -277,6 +314,34 @@ class FacetGrid(object):
             if namedict is None:
                 ax.set_visible(False)
 
+    def add_line_legend(self, huelabel):
+        figlegend = self.fig.legend(
+            handles=self._mappables[-1],
+            labels=list(self.data.coords[huelabel].values),
+            title=huelabel,
+            loc="center right")
+
+        # Draw the plot to set the bounding boxes correctly
+        self.fig.draw(self.fig.canvas.get_renderer())
+
+        # Calculate and set the new width of the figure so the legend fits
+        legend_width = figlegend.get_window_extent().width / self.fig.dpi
+        figure_width = self.fig.get_figwidth()
+        self.fig.set_figwidth(figure_width + legend_width)
+
+        # Draw the plot again to get the new transformations
+        self.fig.draw(self.fig.canvas.get_renderer())
+
+        # Now calculate how much space we need on the right side
+        legend_width = figlegend.get_window_extent().width / self.fig.dpi
+        space_needed = legend_width / (figure_width + legend_width) + 0.02
+        # margin = .01
+        # _space_needed = margin + space_needed
+        right = 1 - space_needed
+
+        # Place the subplot axes to give space for the legend
+        self.fig.subplots_adjust(right=right)
+
     def add_colorbar(self, **kwargs):
         """Draw a colorbar
         """
diff --git a/xarray/plot/plot.py b/xarray/plot/plot.py
index f49ec5c5..6322fc09 100644
--- a/xarray/plot/plot.py
+++ b/xarray/plot/plot.py
@@ -83,8 +83,32 @@ def _easy_facetgrid(darray, plotfunc, x, y, row=None, col=None,
     return g.map_dataarray(plotfunc, x, y, **kwargs)
 
 
-def plot(darray, row=None, col=None, col_wrap=None, ax=None, rtol=0.01,
-         subplot_kws=None, **kwargs):
+def _line_facetgrid(darray, row=None, col=None, hue=None,
+                    col_wrap=None, sharex=True, sharey=True, aspect=None,
+                    size=None, subplot_kws=None, **kwargs):
+    """
+    Convenience method to call xarray.plot.FacetGrid for line plots
+    kwargs are the arguments to pyplot.plot()
+    """
+    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=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)
+
+
+def plot(darray, row=None, col=None, col_wrap=None, ax=None, hue=None,
+         rtol=0.01, subplot_kws=None, **kwargs):
     """
     Default plot of DataArray using matplotlib.pyplot.
 
@@ -106,6 +130,8 @@ def plot(darray, row=None, col=None, col_wrap=None, ax=None, rtol=0.01,
         If passed, make row faceted plots on this dimension name
     col : string, optional
         If passed, make column faceted plots on this dimension name
+    hue : string, optional
+        If passed, make faceted line plots with hue on this dimension name
     col_wrap : integer, optional
         Use together with ``col`` to wrap faceted plots
     ax : matplotlib axes, optional
@@ -129,26 +155,28 @@ def plot(darray, row=None, col=None, col_wrap=None, ax=None, rtol=0.01,
     plot_dims = set(darray.dims)
     plot_dims.discard(row)
     plot_dims.discard(col)
+    plot_dims.discard(hue)
 
     ndims = len(plot_dims)
 
-    error_msg = ('Only 2d plots are supported for facets in xarray. '
+    error_msg = ('Only 1d and 2d plots are supported for facets in xarray. '
                  'See the package `Seaborn` for more options.')
 
-    if ndims == 1:
+    if ndims in [1, 2]:
         if row or col:
-            raise ValueError(error_msg)
-        plotfunc = line
-    elif ndims == 2:
-        # Only 2d can FacetGrid
-        kwargs['row'] = row
-        kwargs['col'] = col
-        kwargs['col_wrap'] = col_wrap
-        kwargs['subplot_kws'] = subplot_kws
-
-        plotfunc = pcolormesh
+            kwargs['row'] = row
+            kwargs['col'] = col
+            kwargs['col_wrap'] = col_wrap
+            kwargs['subplot_kws'] = subplot_kws
+        if ndims == 1:
+            plotfunc = line
+            kwargs['hue'] = hue
+        elif ndims == 2:
+            if hue:
+                raise ValueError('hue is not compatible with 2d data')
+            plotfunc = pcolormesh
     else:
-        if row or col:
+        if row or col or hue:
             raise ValueError(error_msg)
         plotfunc = hist
 
@@ -157,6 +185,61 @@ def plot(darray, row=None, col=None, col_wrap=None, ax=None, rtol=0.01,
     return plotfunc(darray, **kwargs)
 
 
+def _infer_line_data(darray, x, y, hue):
+    error_msg = ('must be either None or one of ({0:s})'
+                 .format(', '.join([repr(dd) for dd in darray.dims])))
+    ndims = len(darray.dims)
+
+    if x is not None and x not in darray.dims:
+        raise ValueError('x ' + error_msg)
+
+    if y is not None and y not in darray.dims:
+        raise ValueError('y ' + error_msg)
+
+    if x is not None and y is not None:
+        raise ValueError('You cannot specify both x and y kwargs'
+                         'for line plots.')
+
+    if ndims == 1:
+        dim, = darray.dims  # get the only dimension name
+        huename = None
+        hueplt = None
+        huelabel = ''
+
+        if (x is None and y is None) or x == dim:
+            xplt = darray.coords[dim]
+            yplt = darray
+
+        else:
+            yplt = darray.coords[dim]
+            xplt = darray
+
+    else:
+        if x is None and y is None and hue is None:
+            raise ValueError('For 2D inputs, please'
+                             'specify either hue, x or y.')
+
+        if y is None:
+            xname, huename = _infer_xy_labels(darray=darray, x=x, y=hue)
+            yname = darray.name
+            xplt = darray.coords[xname]
+            yplt = darray.transpose(xname, huename)
+
+        else:
+            yname, huename = _infer_xy_labels(darray=darray, x=y, y=hue)
+            xname = darray.name
+            xplt = darray.transpose(yname, huename)
+            yplt = darray.coords[yname]
+
+        hueplt = darray.coords[huename]
+        huelabel = label_from_attrs(darray[huename])
+
+    xlabel = label_from_attrs(xplt)
+    ylabel = label_from_attrs(yplt)
+
+    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):
@@ -182,8 +265,7 @@ def line(darray, *args, **kwargs):
         Axis on which to plot this figure. By default, use the current axis.
         Mutually exclusive with ``size`` and ``figsize``.
     hue : string, optional
-        Coordinate for which you want multiple lines plotted
-        (2D DataArrays only).
+        Coordinate for which you want multiple lines plotted.
     x, y : string, optional
         Coordinates for x, y axis. Only one of these may be specified.
         The other coordinate plots values from the DataArray on which this
@@ -201,6 +283,15 @@ def line(darray, *args, **kwargs):
 
     """
 
+    # Handle facetgrids first
+    row = kwargs.pop('row', None)
+    col = kwargs.pop('col', None)
+    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)
     if ndims > 2:
         raise ValueError('Line plots are for 1- or 2-dimensional DataArrays. '
@@ -218,70 +309,28 @@ def line(darray, *args, **kwargs):
     xincrease = kwargs.pop('xincrease', True)
     yincrease = kwargs.pop('yincrease', True)
     add_legend = kwargs.pop('add_legend', True)
+    _labels = kwargs.pop('_labels', True)
 
     ax = get_axis(figsize, size, aspect, ax)
-
-    error_msg = ('must be either None or one of ({0:s})'
-                 .format(', '.join([repr(dd) for dd in darray.dims])))
-
-    if x is not None and x not in darray.dims:
-        raise ValueError('x ' + error_msg)
-
-    if y is not None and y not in darray.dims:
-        raise ValueError('y ' + error_msg)
-
-    if x is not None and y is not None:
-        raise ValueError('You cannot specify both x and y kwargs'
-                         'for line plots.')
-
-    if ndims == 1:
-        dim, = darray.dims  # get the only dimension name
-
-        if (x is None and y is None) or x == dim:
-            xplt = darray.coords[dim]
-            yplt = darray
-
-        else:
-            yplt = darray.coords[dim]
-            xplt = darray
-
-    else:
-        if x is None and y is None and hue is None:
-            raise ValueError('For 2D inputs, please specify either hue or x.')
-
-        if y is None:
-            xlabel, huelabel = _infer_xy_labels(darray=darray, x=x, y=hue)
-            ylabel = darray.name
-            xplt = darray.coords[xlabel]
-            yplt = darray.transpose(xlabel, huelabel)
-
-        else:
-            ylabel, huelabel = _infer_xy_labels(darray=darray, x=y, y=hue)
-            xlabel = darray.name
-            xplt = darray.transpose(ylabel, huelabel)
-            yplt = darray.coords[ylabel]
-
-        huecoords = darray[huelabel]
-        huelabel = label_from_attrs(huecoords)
-
-    xlabel = label_from_attrs(xplt)
-    ylabel = label_from_attrs(yplt)
+    xplt, yplt, hueplt, xlabel, ylabel, huelabel = \
+        _infer_line_data(darray, x, y, hue)
 
     _ensure_plottable(xplt)
 
     primitive = ax.plot(xplt, yplt, *args, **kwargs)
 
-    if xlabel is not None:
-        ax.set_xlabel(xlabel)
+    if _labels:
+        if xlabel is not None:
+            ax.set_xlabel(xlabel)
 
-    if ylabel is not None:
-        ax.set_ylabel(ylabel)
+        if ylabel is not None:
+            ax.set_ylabel(ylabel)
 
-    ax.set_title(darray._title_for_slice())
+        ax.set_title(darray._title_for_slice())
 
     if darray.ndim == 2 and add_legend:
         ax.legend(handles=primitive,
-                  labels=list(huecoords.values),
+                  labels=list(hueplt.values),
                   title=huelabel)
 
     # Rotate dates on xlabels


Test patch:
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
index db1fb2fd..cdb515ba 100644
--- a/xarray/tests/test_plot.py
+++ b/xarray/tests/test_plot.py
@@ -1521,6 +1521,72 @@ class TestFacetGrid4d(PlotTestCase):
             assert substring_in_axes(label, ax)
 
 
+class TestFacetedLinePlots(PlotTestCase):
+    def setUp(self):
+        self.darray = DataArray(np.random.randn(10, 6, 3, 4),
+                                dims=['hue', 'x', 'col', 'row'],
+                                coords=[range(10), range(6),
+                                        range(3), ['A', 'B', 'C', 'C++']],
+                                name='Cornelius Ortega the 1st')
+
+    def test_facetgrid_shape(self):
+        g = self.darray.plot(row='row', col='col', hue='hue')
+        assert g.axes.shape == (len(self.darray.row), len(self.darray.col))
+
+        g = self.darray.plot(row='col', col='row', hue='hue')
+        assert g.axes.shape == (len(self.darray.col), len(self.darray.row))
+
+    def test_default_labels(self):
+        g = self.darray.plot(row='row', col='col', hue='hue')
+        # Rightmost column should be labeled
+        for label, ax in zip(self.darray.coords['row'].values, g.axes[:, -1]):
+            assert substring_in_axes(label, ax)
+
+        # Top row should be labeled
+        for label, ax in zip(self.darray.coords['col'].values, g.axes[0, :]):
+            assert substring_in_axes(str(label), ax)
+
+        # Leftmost column should have array name
+        for ax in g.axes[:, 0]:
+            assert substring_in_axes(self.darray.name, ax)
+
+    def test_test_empty_cell(self):
+        g = self.darray.isel(row=1).drop('row').plot(col='col',
+                                                     hue='hue',
+                                                     col_wrap=2)
+        bottomright = g.axes[-1, -1]
+        assert not bottomright.has_data()
+        assert not bottomright.get_visible()
+
+    def test_set_axis_labels(self):
+        g = self.darray.plot(row='row', col='col', hue='hue')
+        g.set_axis_labels('longitude', 'latitude')
+        alltxt = text_in_fig()
+
+        assert 'longitude' in alltxt
+        assert 'latitude' in alltxt
+
+    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.darray.plot.line(row='row', col='col',
+                                  x='x', ax=plt.axes())
+
+    def test_figsize_and_size(self):
+        with pytest.raises(ValueError):
+            self.darray.plot.line(row='row', col='col',
+                                  x='x', size=3, figsize=4)
+
+    def test_wrong_num_of_dimensions(self):
+        with pytest.raises(ValueError):
+            self.darray.plot(row='row', hue='hue')
+            self.darray.plot.line(row='row', hue='hue')
+
+
 class TestDatetimePlot(PlotTestCase):
     def setUp(self):
         '''