ENH: categorical scatter plot

pandas/plotting/_matplotlib/core.py

@@ -9,6 +9,7 @@
 from pandas.util._decorators import cache_readonly
 
 from pandas.core.dtypes.common import (
+    is_categorical_dtype,
     is_extension_array_dtype,
     is_float,
     is_float_dtype,
@@ -440,7 +441,9 @@ def _compute_plot_data(self):
         if is_empty:
             raise TypeError("no numeric data to plot")
 
-        self.data = numeric_data.apply(self._convert_to_ndarray)
+        self.data = numeric_data.apply(
+            lambda x: self._convert_to_ndarray(x) if not is_categorical_dtype(x) else x
+        )
 
     def _make_plot(self):
         raise AbstractMethodError(self)
@@ -973,7 +976,7 @@ def _plot_colorbar(self, ax: "Axes", **kwds):
 
         if mpl_ge_3_0_0():
             # The workaround below is no longer necessary.
-            return
+            return cbar
 
         points = ax.get_position().get_points()
         cbar_points = cbar.ax.get_position().get_points()
@@ -991,6 +994,8 @@ def _plot_colorbar(self, ax: "Axes", **kwds):
         # print(points[1, 1] - points[0, 1])
         # print(cbar_points[1, 1] - cbar_points[0, 1])
 
+        return cbar
+
 
 class ScatterPlot(PlanePlot):
     _kind = "scatter"
@@ -1013,6 +1018,8 @@ def _make_plot(self):
 
         c_is_column = is_hashable(c) and c in self.data.columns
 
+        color_by_categorical = c_is_column and is_categorical_dtype(self.data[c])
+
         # pandas uses colormap, matplotlib uses cmap.
         cmap = self.colormap or "Greys"
         cmap = self.plt.cm.get_cmap(cmap)
@@ -1024,10 +1031,22 @@ def _make_plot(self):
         elif color is not None:
             c_values = color
         elif c_is_column:
-            c_values = self.data[c].values
+            if color_by_categorical:
+                c_values = self.data[c].cat.codes
+            else:
+                c_values = self.data[c].values
         else:
             c_values = c
 
+        if color_by_categorical:
+            from matplotlib import colors
+
+            n_cats = len(self.data[c].cat.categories)
+            cmap = colors.ListedColormap([cmap(i) for i in range(cmap.N)])
+            bounds = np.linspace(0, n_cats, n_cats + 1)
+            norm = colors.BoundaryNorm(bounds, cmap.N)
+        else:
+            norm = None
         # plot colorbar if
         # 1. colormap is assigned, and
         # 2.`c` is a column containing only numeric values
@@ -1044,11 +1063,15 @@ def _make_plot(self):
             c=c_values,
             label=label,
             cmap=cmap,
+            norm=norm,
             **self.kwds,
         )
         if cb:
             cbar_label = c if c_is_column else ""
-            self._plot_colorbar(ax, label=cbar_label)
+            cbar = self._plot_colorbar(ax, label=cbar_label)
+            if color_by_categorical:
+                cbar.set_ticks(np.linspace(0.5, n_cats - 0.5, n_cats))
+                cbar.ax.set_yticklabels(self.data[c].cat.categories)
 
         if label is not None:
             self._add_legend_handle(scatter, label)

pandas/tests/plotting/frame/test_frame.py

@@ -696,6 +696,36 @@ def test_scatterplot_object_data(self):
         _check_plot_works(df.plot.scatter, x="a", y="b")
         _check_plot_works(df.plot.scatter, x=0, y=1)
 
+    @pytest.mark.parametrize("ordered", [True, False])
+    @pytest.mark.parametrize(
+        "categories",
+        (["setosa", "versicolor", "virginica"], ["versicolor", "virginica", "setosa"]),
+    )
+    def test_scatterplot_color_by_categorical(self, ordered, categories):
+        df = DataFrame(
+            [[5.1, 3.5], [4.9, 3.0], [7.0, 3.2], [6.4, 3.2], [5.9, 3.0]],
+            columns=["length", "width"],
+        )
+        df["species"] = pd.Categorical(
+            ["setosa", "setosa", "virginica", "virginica", "versicolor"],
+            ordered=ordered,
+            categories=categories,
+        )
+        ax = df.plot.scatter(x=0, y=1, c="species")
+        colorbar = ax.collections[0].colorbar
+
+        expected_ticks = np.array([0.5, 1.5, 2.5])
+        result_ticks = colorbar.get_ticks()
+        tm.assert_numpy_array_equal(result_ticks, expected_ticks)
+
+        expected_boundaries = np.array([0.0, 1.0, 2.0, 3.0])
+        result_boundaries = colorbar._boundaries
+        tm.assert_numpy_array_equal(result_boundaries, expected_boundaries)
+
+        expected_yticklabels = categories
+        result_yticklabels = [i.get_text() for i in colorbar.ax.get_ymajorticklabels()]
+        assert all(i == j for i, j in zip(result_yticklabels, expected_yticklabels))
+
     @pytest.mark.parametrize("x, y", [("x", "y"), ("y", "x"), ("y", "y")])
     def test_plot_scatter_with_categorical_data(self, x, y):
         # after fixing GH 18755, should be able to plot categorical data