Surface plots

doc/api-hidden.rst

@@ -597,6 +597,7 @@
    plot.imshow
    plot.pcolormesh
    plot.scatter
+   plot.surface
 
    plot.FacetGrid.map_dataarray
    plot.FacetGrid.set_titles

doc/api.rst

@@ -587,6 +587,7 @@ Plotting
    DataArray.plot.line
    DataArray.plot.pcolormesh
    DataArray.plot.step
+   DataArray.plot.surface
 
 .. _api.ufuncs:
 

doc/whats-new.rst

@@ -23,6 +23,9 @@ v0.17.1 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Add :py:meth:`DataArray.plot.surface` which wraps matplotlib's `plot_surface` to make
+  surface plots (:issue:`#2235` :issue:`#5084` :pull:`5101`).
+  By `John Omotani <https://github.com/johnomotani>`_.
 - Add :py:meth:`Dataset.query` and :py:meth:`DataArray.query` which enable indexing
   of datasets and data arrays by evaluating query expressions against the values of the
   data variables (:pull:`4984`). By `Alistair Miles <https://github.com/alimanfoo>`_.
@@ -78,6 +81,9 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+- Fix 2d plot failure for certain combinations of dimensions when `x` is 1d and `y` is
+  2d (:issue:`5079`, :pull:`5099`).
+  By `John Omotani <https://github.com/johnomotani>`_
 - Ensure standard calendar times encoded with large values (i.e. greater than approximately 292 years), can be decoded correctly without silently overflowing (:pull:`5050`).  This was a regression in xarray 0.17.0.  By `Zeb Nicholls <https://github.com/znicholls>`_.
 - Added support for `numpy.bool_` attributes in roundtrips using `h5netcdf` engine with `invalid_netcdf=True` [which casts `bool`s to `numpy.bool_`] (:issue:`4981`, :pull:`4986`).
   By `Victor Negîrneac <https://github.com/caenrigen>`_.

xarray/plot/__init__.py

@@ -1,6 +1,6 @@
 from .dataset_plot import scatter
 from .facetgrid import FacetGrid
-from .plot import contour, contourf, hist, imshow, line, pcolormesh, plot, step
+from .plot import contour, contourf, hist, imshow, line, pcolormesh, plot, step, surface
 
 __all__ = [
     "plot",
@@ -13,4 +13,5 @@
     "pcolormesh",
     "FacetGrid",
     "scatter",
+    "surface",
 ]

xarray/plot/facetgrid.py

@@ -263,7 +263,9 @@ def map_dataarray(self, func, x, y, **kwargs):
             if k not in {"cmap", "colors", "cbar_kwargs", "levels"}
         }
         func_kwargs.update(cmap_params)
-        func_kwargs.update({"add_colorbar": False, "add_labels": False})
+        func_kwargs["add_colorbar"] = False
+        if func.__name__ != "surface":
+            func_kwargs["add_labels"] = False
 
         # Get x, y labels for the first subplot
         x, y = _infer_xy_labels(

xarray/plot/plot.py

@@ -633,7 +633,9 @@ def newplotfunc(
 
         # Decide on a default for the colorbar before facetgrids
         if add_colorbar is None:
-            add_colorbar = plotfunc.__name__ != "contour"
+            add_colorbar = plotfunc.__name__ != "contour" and not (
+                plotfunc.__name__ == "surface" and cmap is None
+            )
         imshow_rgb = plotfunc.__name__ == "imshow" and darray.ndim == (
             3 + (row is not None) + (col is not None)
         )
@@ -646,6 +648,24 @@ def newplotfunc(
                 darray = _rescale_imshow_rgb(darray, vmin, vmax, robust)
                 vmin, vmax, robust = None, None, False
 
+        if subplot_kws is None:
+            subplot_kws = dict()
+
+        if plotfunc.__name__ == "surface" and not kwargs.get("_is_facetgrid", False):
+            if ax is None:
+                # TODO: Importing Axes3D is not necessary in matplotlib >= 3.2.
+                # Remove when minimum requirement of matplotlib is 3.2:
+                from mpl_toolkits.mplot3d import Axes3D  # type: ignore  # noqa: F401
+
+                del Axes3D
+
+                # Need to create a "3d" Axes instance for surface plots
+                subplot_kws["projection"] = "3d"
+
+            # In facet grids, shared axis labels don't make sense for surface plots
+            sharex = False
+            sharey = False
+
         # Handle facetgrids first
         if row or col:
             allargs = locals().copy()
@@ -658,6 +678,19 @@ def newplotfunc(
 
         plt = import_matplotlib_pyplot()
 
+        if (
+            plotfunc.__name__ == "surface"
+            and not kwargs.get("_is_facetgrid", False)
+            and ax is not None
+        ):
+            import mpl_toolkits  # type: ignore
+
+            if not isinstance(ax, mpl_toolkits.mplot3d.Axes3D):
+                raise ValueError(
+                    "If ax is passed to surface(), it must be created with "
+                    'projection="3d"'
+                )
+
         rgb = kwargs.pop("rgb", None)
         if rgb is not None and plotfunc.__name__ != "imshow":
             raise ValueError('The "rgb" keyword is only valid for imshow()')
@@ -671,28 +704,22 @@ def newplotfunc(
             darray=darray, x=x, y=y, imshow=imshow_rgb, rgb=rgb
         )
 
-        # better to pass the ndarrays directly to plotting functions
-        xval = darray[xlab].values
-        yval = darray[ylab].values
-
-        # check if we need to broadcast one dimension
-        if xval.ndim < yval.ndim:
-            dims = darray[ylab].dims
-            if xval.shape[0] == yval.shape[0]:
-                xval = np.broadcast_to(xval[:, np.newaxis], yval.shape)
-            else:
-                xval = np.broadcast_to(xval[np.newaxis, :], yval.shape)
+        xval = darray[xlab]
+        yval = darray[ylab]
 
-        elif yval.ndim < xval.ndim:
-            dims = darray[xlab].dims
-            if yval.shape[0] == xval.shape[0]:
-                yval = np.broadcast_to(yval[:, np.newaxis], xval.shape)
-            else:
-                yval = np.broadcast_to(yval[np.newaxis, :], xval.shape)
-        elif xval.ndim == 2:
-            dims = darray[xlab].dims
+        if xval.ndim > 1 or yval.ndim > 1 or plotfunc.__name__ == "surface":
+            # Passing 2d coordinate values, need to ensure they are transposed the same
+            # way as darray.
+            # Also surface plots always need 2d coordinates
+            xval = xval.broadcast_like(darray)
+            yval = yval.broadcast_like(darray)
+            dims = darray.dims
         else:
-            dims = (darray[ylab].dims[0], darray[xlab].dims[0])
+            dims = (yval.dims[0], xval.dims[0])
+
+        # better to pass the ndarrays directly to plotting functions
+        xval = xval.values
+        yval = yval.values
 
         # May need to transpose for correct x, y labels
         # xlab may be the name of a coord, we have to check for dim names
@@ -741,8 +768,6 @@ def newplotfunc(
             # forbid usage of mpl strings
             raise ValueError("plt.imshow's `aspect` kwarg is not available in xarray")
 
-        if subplot_kws is None:
-            subplot_kws = dict()
         ax = get_axis(figsize, size, aspect, ax, **subplot_kws)
 
         primitive = plotfunc(
@@ -762,6 +787,8 @@ def newplotfunc(
             ax.set_xlabel(label_from_attrs(darray[xlab], xlab_extra))
             ax.set_ylabel(label_from_attrs(darray[ylab], ylab_extra))
             ax.set_title(darray._title_for_slice())
+            if plotfunc.__name__ == "surface":
+                ax.set_zlabel(label_from_attrs(darray))
 
         if add_colorbar:
             if add_labels and "label" not in cbar_kwargs:
@@ -994,3 +1021,14 @@ def pcolormesh(x, y, z, ax, infer_intervals=None, **kwargs):
         ax.set_ylim(y[0], y[-1])
 
     return primitive
+
+
+@_plot2d
+def surface(x, y, z, ax, **kwargs):
+    """
+    Surface plot of 2d DataArray
+
+    Wraps :func:`matplotlib:mpl_toolkits.mplot3d.axes3d.plot_surface`
+    """
+    primitive = ax.plot_surface(x, y, z, **kwargs)
+    return primitive

xarray/plot/utils.py

@@ -804,6 +804,14 @@ def _process_cmap_cbar_kwargs(
     cmap_params
     cbar_kwargs
     """
+    if func.__name__ == "surface":
+        # Leave user to specify cmap settings for surface plots
+        kwargs["cmap"] = cmap
+        return {
+            k: kwargs.get(k, None)
+            for k in ["vmin", "vmax", "cmap", "extend", "levels", "norm"]
+        }, {}
+
     cbar_kwargs = {} if cbar_kwargs is None else dict(cbar_kwargs)
 
     if "contour" in func.__name__ and levels is None:

xarray/tests/__init__.py

@@ -60,6 +60,9 @@ def LooseVersion(vstring):
 
 
 has_matplotlib, requires_matplotlib = _importorskip("matplotlib")
+has_matplotlib_3_3_0, requires_matplotlib_3_3_0 = _importorskip(
+    "matplotlib", minversion="3.3.0"
+)
 has_scipy, requires_scipy = _importorskip("scipy")
 has_pydap, requires_pydap = _importorskip("pydap.client")
 has_netCDF4, requires_netCDF4 = _importorskip("netCDF4")

xarray/tests/test_plot.py

@@ -2,6 +2,7 @@
 import inspect
 from copy import copy
 from datetime import datetime
+from typing import Any, Dict
 
 import numpy as np
 import pandas as pd
@@ -28,6 +29,7 @@
     requires_cartopy,
     requires_cftime,
     requires_matplotlib,
+    requires_matplotlib_3_3_0,
     requires_nc_time_axis,
     requires_seaborn,
 )
@@ -36,6 +38,7 @@
 try:
     import matplotlib as mpl
     import matplotlib.pyplot as plt
+    import mpl_toolkits  # type: ignore
 except ImportError:
     pass
 
@@ -132,8 +135,9 @@ def setup(self):
         # Remove all matplotlib figures
         plt.close("all")
 
-    def pass_in_axis(self, plotmethod):
-        fig, axes = plt.subplots(ncols=2)
+    def pass_in_axis(self, plotmethod, subplot_kw=None):
+        subplot_kw = {} if subplot_kw is None else subplot_kw
+        fig, axes = plt.subplots(ncols=2, subplot_kw=subplot_kw)
         plotmethod(ax=axes[0])
         assert axes[0].has_data()
 
@@ -1079,6 +1083,9 @@ class Common2dMixin:
     Should have the same name as the method.
     """
 
+    # Needs to be overridden in TestSurface for facet grid plots
+    subplot_kws: Dict[Any, Any] = {}
+
     @pytest.fixture(autouse=True)
     def setUp(self):
         da = DataArray(
@@ -1394,7 +1401,7 @@ def test_colorbar_kwargs(self):
     def test_verbose_facetgrid(self):
         a = easy_array((10, 15, 3))
         d = DataArray(a, dims=["y", "x", "z"])
-        g = xplt.FacetGrid(d, col="z")
+        g = xplt.FacetGrid(d, col="z", subplot_kws=self.subplot_kws)
         g.map_dataarray(self.plotfunc, "x", "y")
         for ax in g.axes.flat:
             assert ax.has_data()
@@ -1794,6 +1801,105 @@ def test_origin_overrides_xyincrease(self):
             assert plt.ylim()[0] < 0
 
 
+# The try/except/else is needed for the matplotlib version check, to handle the case
+# when matplotlib is not installed. It should be possible to remove it once we require
+# matplotlib>=3.2.0.
+# Note, importing as mpl2 to avoid redefining mpl, which is a flake8 error.
+try:
+    import matplotlib as mpl2  # noqa: F401
+except ImportError:
+    pass
+else:
+
+    class TestSurface(Common2dMixin, PlotTestCase):
+
+        plotfunc = staticmethod(xplt.surface)
+        subplot_kws = {"projection": "3d"}
+
+        def test_primitive_artist_returned(self):
+            artist = self.plotmethod()
+            assert isinstance(artist, mpl_toolkits.mplot3d.art3d.Poly3DCollection)
+
+        @pytest.mark.slow
+        def test_2d_coord_names(self):
+            self.plotmethod(x="x2d", y="y2d")
+            # make sure labels came out ok
+            ax = plt.gca()
+            assert "x2d" == ax.get_xlabel()
+            assert "y2d" == ax.get_ylabel()
+            assert f"{self.darray.long_name} [{self.darray.units}]" == ax.get_zlabel()
+
+        def test_xyincrease_false_changes_axes(self):
+            # Does not make sense for surface plots
+            pytest.skip("does not make sense for surface plots")
+
+        def test_xyincrease_true_changes_axes(self):
+            # Does not make sense for surface plots
+            pytest.skip("does not make sense for surface plots")
+
+        def test_can_pass_in_axis(self):
+            self.pass_in_axis(self.plotmethod, subplot_kw={"projection": "3d"})
+
+        def test_default_cmap(self):
+            # Does not make sense for surface plots with default arguments
+            pytest.skip("does not make sense for surface plots")
+
+        def test_diverging_color_limits(self):
+            # Does not make sense for surface plots with default arguments
+            pytest.skip("does not make sense for surface plots")
+
+        def test_colorbar_kwargs(self):
+            # Does not make sense for surface plots with default arguments
+            pytest.skip("does not make sense for surface plots")
+
+        def test_cmap_and_color_both(self):
+            # Does not make sense for surface plots with default arguments
+            pytest.skip("does not make sense for surface plots")
+
+        def test_seaborn_palette_as_cmap(self):
+            # seaborn does not work with mpl_toolkits.mplot3d
+            with pytest.raises(ValueError):
+                super().test_seaborn_palette_as_cmap()
+
+        # Need to modify this test for surface(), because all subplots should have labels,
+        # not just left and bottom
+        @pytest.mark.filterwarnings("ignore:tight_layout cannot")
+        def test_convenient_facetgrid(self):
+            a = easy_array((10, 15, 4))
+            d = DataArray(a, dims=["y", "x", "z"])
+            g = self.plotfunc(d, x="x", y="y", col="z", col_wrap=2)
+
+            assert_array_equal(g.axes.shape, [2, 2])
+            for (y, x), ax in np.ndenumerate(g.axes):
+                assert ax.has_data()
+                assert "y" == ax.get_ylabel()
+                assert "x" == ax.get_xlabel()
+
+            # Infering labels
+            g = self.plotfunc(d, col="z", col_wrap=2)
+            assert_array_equal(g.axes.shape, [2, 2])
+            for (y, x), ax in np.ndenumerate(g.axes):
+                assert ax.has_data()
+                assert "y" == ax.get_ylabel()
+                assert "x" == ax.get_xlabel()
+
+        @requires_matplotlib_3_3_0
+        def test_viridis_cmap(self):
+            return super().test_viridis_cmap()
+
+        @requires_matplotlib_3_3_0
+        def test_can_change_default_cmap(self):
+            return super().test_can_change_default_cmap()
+
+        @requires_matplotlib_3_3_0
+        def test_colorbar_default_label(self):
+            return super().test_colorbar_default_label()
+
+        @requires_matplotlib_3_3_0
+        def test_facetgrid_map_only_appends_mappables(self):
+            return super().test_facetgrid_map_only_appends_mappables()
+
+
 class TestFacetGrid(PlotTestCase):
     @pytest.fixture(autouse=True)
     def setUp(self):