Merge branch 'master' into plot

README.rst

@@ -115,5 +115,5 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 
-xray includes portions of pandas and numpy. Their licenses are included in the
-licenses directory.
+xray includes portions of pandas, NumPy and Seaborn. Their licenses are
+included in the licenses directory.

licenses/SEABORN_LICENSE

@@ -0,0 +1,27 @@
+Copyright (c) 2012-2013, Michael L. Waskom
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the {organization} nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

xray/core/dataarray.py

@@ -4,7 +4,7 @@
 
 import pandas as pd
 
-from ..plotting import _PlotMethods
+from ..plotting.plotting import _PlotMethods
 
 from . import indexing
 from . import groupby

xray/plotting/__init__.py

@@ -1,3 +1,2 @@
 from .plotting import (plot, line, contourf, contour,
-                       hist, imshow, pcolormesh,
-                       _infer_interval_breaks, _PlotMethods)
+                       hist, imshow, pcolormesh)

xray/plotting/plotting.py

@@ -202,8 +202,60 @@ def _update_axes_limits(ax, xincrease, yincrease):
         ax.set_ylim(sorted(ax.get_ylim(), reverse=True))
 
 
-# Decorator will fill this list
-_plot2dlist = []
+def _determine_cmap_params(plot_data, vmin=None, vmax=None, cmap=None,
+                           center=None, robust=False, extend=None):
+    """
+    Use some heuristics to set good defaults for colorbar and range.
+
+    Adapted from Seaborn:
+    https://github.com/mwaskom/seaborn/blob/v0.6/seaborn/matrix.py#L158
+    """
+    calc_data = plot_data[~pd.isnull(plot_data)]
+    if vmin is None:
+        vmin = np.percentile(calc_data, 2) if robust else calc_data.min()
+    if vmax is None:
+        vmax = np.percentile(calc_data, 98) if robust else calc_data.max()
+
+    # Simple heuristics for whether these data should  have a divergent map
+    divergent = ((vmin < 0) and (vmax > 0)) or center is not None
+
+    # Now set center to 0 so math below makes sense
+    if center is None:
+        center = 0
+
+    # A divergent map should be symmetric around the center value
+    if divergent:
+        vlim = max(abs(vmin - center), abs(vmax - center))
+        vmin, vmax = -vlim, vlim
+
+    # Now add in the centering value and set the limits
+    vmin += center
+    vmax += center
+
+    # Choose default colormaps if not provided
+    if cmap is None:
+        if divergent:
+            cmap = "RdBu_r"
+        else:
+            cmap = "viridis"
+
+    if cmap == "viridis":
+        cmap = _load_default_cmap()
+
+    if extend is None:
+        extend_min = calc_data.min() < vmin
+        extend_max = calc_data.max() > vmax
+        if extend_min and extend_max:
+            extend = 'both'
+        elif extend_min:
+            extend = 'min'
+        elif extend_max:
+            extend = 'max'
+        else:
+            extend = 'neither'
+
+    return vmin, vmax, cmap, extend
+
 
 def _plot2d(plotfunc):
     """
@@ -224,6 +276,23 @@ def _plot2d(plotfunc):
         if None, use the default for the matplotlib function
     add_colorbar : Boolean, optional
         Adds colorbar to axis
+    vmin, vmax : floats, optional
+        Values to anchor the colormap, otherwise they are inferred from the
+        data and other keyword arguments. When a diverging dataset is inferred,
+        one of these values may be ignored.
+    cmap : matplotlib colormap name or object, optional
+        The mapping from data values to color space. If not provided, this
+        will be either be ``viridis`` (if the function infers a sequential
+        dataset) or ``RdBu_r`` (if the function infers a diverging dataset).
+    center : float, optional
+        The value at which to center the colormap. Passing this value implies
+        use of a diverging colormap.
+    robust : bool, optional
+        If True and ``vmin`` or ``vmax`` are absent, the colormap range is
+        computed with robust quantiles instead of the extreme values.
+    extend : {'neither', 'both', 'min', 'max'}, optional
+        How to draw arrows extending the colorbar beyond its limits. If not
+        provided, extend is inferred from vmin, vmax and the data limits.
     **kwargs : optional
         Additional arguments to wrapped matplotlib function
 
@@ -237,12 +306,10 @@ def _plot2d(plotfunc):
     # Build on the original docstring
     plotfunc.__doc__ = '\n'.join((plotfunc.__doc__, commondoc))
 
-    # global list of 2d plotting functions
-    _plot2dlist.append(plotfunc)
-
     @functools.wraps(plotfunc)
     def wrapper(darray, ax=None, xincrease=None, yincrease=None,
-                add_colorbar=True, **kwargs):
+                add_colorbar=True, vmin=None, vmax=None, cmap=None,
+                center=None, robust=False, extend=None, **kwargs):
         # All 2d plots in xray share this function signature
 
         import matplotlib.pyplot as plt
@@ -264,15 +331,23 @@ def wrapper(darray, ax=None, xincrease=None, yincrease=None,
 
         _ensure_plottable(x, y)
 
-        cmap = kwargs.pop('cmap', _load_default_cmap())
+        vmin, vmax, cmap, extend = _determine_cmap_params(
+            z.data, vmin, vmax, cmap, center, robust, extend)
+
+        if 'contour' in plotfunc.__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'] = extend
 
-        ax, primitive = plotfunc(x, y, z, ax=ax, cmap=cmap, **kwargs)
+        ax, primitive = plotfunc(x, y, z, ax=ax, cmap=cmap, vmin=vmin,
+                                 vmax=vmax, **kwargs)
 
         ax.set_xlabel(xlab)
         ax.set_ylabel(ylab)
 
         if add_colorbar:
-            plt.colorbar(primitive, ax=ax)
+            plt.colorbar(primitive, ax=ax, extend=extend)
 
         _update_axes_limits(ax, xincrease, yincrease)
 

xray/test/test_plotting.py

@@ -2,7 +2,10 @@
 import pandas as pd
 
 from xray import DataArray
+
 import xray.plotting as xplt
+from xray.plotting.plotting import (_infer_interval_breaks,
+                                    _determine_cmap_params)
 
 from . import TestCase, requires_matplotlib
 
@@ -62,11 +65,11 @@ def test_can_pass_in_axis(self):
         self.pass_in_axis(self.darray.plot)
 
     def test__infer_interval_breaks(self):
-        self.assertArrayEqual([-0.5, 0.5, 1.5], xplt._infer_interval_breaks([0, 1]))
+        self.assertArrayEqual([-0.5, 0.5, 1.5], _infer_interval_breaks([0, 1]))
         self.assertArrayEqual([-0.5, 0.5, 5.0, 9.5, 10.5],
-                              xplt._infer_interval_breaks([0, 1, 9, 10]))
+                              _infer_interval_breaks([0, 1, 9, 10]))
         self.assertArrayEqual(pd.date_range('20000101', periods=4) - np.timedelta64(12, 'h'),
-                              xplt._infer_interval_breaks(pd.date_range('20000101', periods=3)))
+                              _infer_interval_breaks(pd.date_range('20000101', periods=3)))
 
 
 class TestPlot1D(PlotTestCase):
@@ -158,6 +161,24 @@ def test_plot_nans(self):
         self.darray.plot.hist()
 
 
+@requires_matplotlib
+class TestDetermineCmapParams(TestCase):
+    def test_robust(self):
+        data = np.random.RandomState(1).rand(100)
+        vmin, vmax, cmap, extend = _determine_cmap_params(data, robust=True)
+        self.assertEqual(vmin, np.percentile(data, 2))
+        self.assertEqual(vmax, np.percentile(data, 98))
+        self.assertEqual(cmap.name, 'viridis')
+        self.assertEqual(extend, 'both')
+
+    def test_center(self):
+        data = np.random.RandomState(2).rand(100)
+        vmin, vmax, cmap, extend = _determine_cmap_params(data, center=0.5)
+        self.assertEqual(vmax - 0.5, 0.5 - vmin)
+        self.assertEqual(cmap, 'RdBu_r')
+        self.assertEqual(extend, 'neither')
+
+
 class Common2dMixin:
     """
     Common tests for 2d plotting go here.
@@ -208,19 +229,34 @@ def test_xyincrease_true_changes_axes(self):
         self.assertTrue(all(abs(x) < 1 for x in diffs))
 
     def test_plot_nans(self):
-        self.darray[0, 0] = np.nan
-        result = self.plotmethod()
-        clim = result.get_clim()
-        self.assertFalse(pd.isnull(np.array(clim)).any())
+        x1 = self.darray[:5]
+        x2 = self.darray.copy()
+        x2[5:] = np.nan
 
-    def test_default_cmap_viridis(self):
+        clim1 = self.plotfunc(x1).get_clim()
+        clim2 = self.plotfunc(x2).get_clim()
+        self.assertEqual(clim1, clim2)
+
+    def test_viridis_cmap(self):
+        cmap_name = self.plotmethod(cmap='viridis').get_cmap().name
+        self.assertEqual('viridis', cmap_name)
+
+    def test_default_cmap(self):
         cmap_name = self.plotmethod().get_cmap().name
+        self.assertEqual('RdBu_r', cmap_name)
+
+        cmap_name = self.plotfunc(abs(self.darray)).get_cmap().name
         self.assertEqual('viridis', cmap_name)
 
     def test_can_change_default_cmap(self):
         cmap_name = self.plotmethod(cmap='jet').get_cmap().name
         self.assertEqual('jet', cmap_name)
 
+    def test_diverging_color_limits(self):
+        artist = self.plotmethod()
+        vmin, vmax = artist.get_clim()
+        self.assertAlmostEqual(-vmin, vmax)
+
 
 class TestContourf(Common2dMixin, PlotTestCase):
 
@@ -235,6 +271,19 @@ def test_primitive_artist_returned(self):
         artist = self.plotmethod()
         self.assertTrue(isinstance(artist, mpl.contour.QuadContourSet))
 
+    def test_extend(self):
+        artist = self.plotmethod()
+        self.assertEqual(artist.extend, 'neither')
+
+        artist = self.plotmethod(robust=True)
+        self.assertEqual(artist.extend, 'both')
+
+        artist = self.plotmethod(vmin=-0, vmax=10)
+        self.assertEqual(artist.extend, 'min')
+
+        artist = self.plotmethod(vmin=-10, vmax=0)
+        self.assertEqual(artist.extend, 'max')
+
 
 class TestContour(Common2dMixin, PlotTestCase):