Allow plotting categorical data

xarray/plot/plot.py

@@ -919,24 +919,31 @@ def imshow(x, y, z, ax, **kwargs):
     The pixels are centered on the coordinates. For example, if the coordinate
     value is 3.2, then the pixels for those coordinates will be centered on 3.2.
     """
+    if any(np.issubdtype(v.dtype, str) for v in [x, y]):
+        raise TypeError("ax.imshow does not support categorical data.")
 
     if x.ndim != 1 or y.ndim != 1:
         raise ValueError(
             "imshow requires 1D coordinates, try using pcolormesh or contour(f)"
         )
 
-    # Centering the pixels- Assumes uniform spacing
-    try:
-        xstep = (x[1] - x[0]) / 2.0
-    except IndexError:
-        # Arbitrary default value, similar to matplotlib behaviour
-        xstep = 0.1
-    try:
-        ystep = (y[1] - y[0]) / 2.0
-    except IndexError:
-        ystep = 0.1
-    left, right = x[0] - xstep, x[-1] + xstep
-    bottom, top = y[-1] + ystep, y[0] - ystep
+    def _maybe_center_pixels(x):
+        """Center the pixels on the coordinates."""
+        if np.issubdtype(x.dtype, str):
+            # Categorical data cannot be centered:
+            return None, None
+
+        try:
+            # Center the pixels assuming uniform spacing:
+            xstep = 0.5 * (x[1] - x[0])
+        except IndexError:
+            # Arbitrary default value, similar to matplotlib behaviour:
+            xstep = 0.1
+        return x[0] - xstep, x[-1] + xstep
+
+    # Try to center the pixels:
+    left, right = _maybe_center_pixels(x)
+    top, bottom = _maybe_center_pixels(y)
 
     defaults = {"origin": "upper", "interpolation": "nearest"}
 
@@ -1011,9 +1018,13 @@ def pcolormesh(x, y, z, ax, infer_intervals=None, **kwargs):
         else:
             infer_intervals = True
 
-    if infer_intervals and (
-        (np.shape(x)[0] == np.shape(z)[1])
-        or ((x.ndim > 1) and (np.shape(x)[1] == np.shape(z)[1]))
+    if (
+        infer_intervals
+        and not np.issubdtype(x.dtype, str)
+        and (
+            (np.shape(x)[0] == np.shape(z)[1])
+            or ((x.ndim > 1) and (np.shape(x)[1] == np.shape(z)[1]))
+        )
     ):
         if len(x.shape) == 1:
             x = _infer_interval_breaks(x, check_monotonic=True)
@@ -1022,7 +1033,11 @@ def pcolormesh(x, y, z, ax, infer_intervals=None, **kwargs):
             x = _infer_interval_breaks(x, axis=1)
             x = _infer_interval_breaks(x, axis=0)
 
-    if infer_intervals and (np.shape(y)[0] == np.shape(z)[0]):
+    if (
+        infer_intervals
+        and not np.issubdtype(y.dtype, str)
+        and (np.shape(y)[0] == np.shape(z)[0])
+    ):
         if len(y.shape) == 1:
             y = _infer_interval_breaks(y, check_monotonic=True)
         else:
@@ -1049,5 +1064,8 @@ def surface(x, y, z, ax, **kwargs):
 
     Wraps :py:meth:`matplotlib:mpl_toolkits.mplot3d.axes3d.Axes3D.plot_surface`.
     """
+    if any(np.issubdtype(v.dtype, str) for v in [x, y]):
+        raise TypeError("ax.plot_surface does not support categorical data.")
+
     primitive = ax.plot_surface(x, y, z, **kwargs)
     return primitive

xarray/plot/utils.py

@@ -604,7 +604,14 @@ def _ensure_plottable(*args):
     Raise exception if there is anything in args that can't be plotted on an
     axis by matplotlib.
     """
-    numpy_types = [np.floating, np.integer, np.timedelta64, np.datetime64, np.bool_]
+    numpy_types = [
+        np.floating,
+        np.integer,
+        np.timedelta64,
+        np.datetime64,
+        np.bool_,
+        np.str_,
+    ]
     other_types = [datetime]
     try:
         import cftime

xarray/tests/test_plot.py

@@ -684,10 +684,9 @@ def test_format_string(self):
     def test_can_pass_in_axis(self):
         self.pass_in_axis(self.darray.plot.line)
 
-    def test_nonnumeric_index_raises_typeerror(self):
+    def test_nonnumeric_index(self):
         a = DataArray([1, 2, 3], {"letter": ["a", "b", "c"]}, dims="letter")
-        with pytest.raises(TypeError, match=r"[Pp]lot"):
-            a.plot.line()
+        a.plot.line()
 
     def test_primitive_returned(self):
         p = self.darray.plot.line()
@@ -1162,9 +1161,13 @@ def test_3d_raises_valueerror(self):
         with pytest.raises(ValueError, match=r"DataArray must be 2d"):
             self.plotfunc(a)
 
-    def test_nonnumeric_index_raises_typeerror(self):
+    def test_nonnumeric_index(self):
         a = DataArray(easy_array((3, 2)), coords=[["a", "b", "c"], ["d", "e"]])
-        with pytest.raises(TypeError, match=r"[Pp]lot"):
+        if self.plotfunc.__name__ in ["imshow", "surface"]:
+            # ax.imshow and ax.plot_surface errors with nonnumerics:
+            with pytest.raises(TypeError, match=r"does not support categorical data"):
+                self.plotfunc(a)
+        else:
             self.plotfunc(a)
 
     def test_multiindex_raises_typeerror(self):