ENH: Allow column grouping in DataFrame.plot

doc/source/whatsnew/v1.1.0.rst

@@ -87,6 +87,7 @@ Other enhancements
 - Positional slicing on a :class:`IntervalIndex` now supports slices with ``step > 1`` (:issue:`31658`)
 - :class:`Series.str` now has a `fullmatch` method that matches a regular expression against the entire string in each row of the series, similar to `re.fullmatch` (:issue:`32806`).
 - :meth:`DataFrame.sample` will now also allow array-like and BitGenerator objects to be passed to ``random_state`` as seeds (:issue:`32503`)
+- :meth:`DataFrame.plot` will now allow for the ``subplots`` parameter to be a list of tuples so that columns may be grouped together in the same plot. (:issue:`29944`)
 - :meth:`MultiIndex.union` will now raise `RuntimeWarning` if the object inside are unsortable, pass `sort=False` to suppress this warning (:issue:`33015`)
 -
 

pandas/plotting/_matplotlib/core.py

@@ -1,3 +1,4 @@
+from collections import Iterable
 import re
 from typing import Optional
 import warnings
@@ -115,6 +116,67 @@ def __init__(
         self.sort_columns = sort_columns
 
         self.subplots = subplots
+        if not isinstance(self.subplots, (bool, Iterable)):
+            raise ValueError("subplots should be a bool or an iterable")
+        if isinstance(self.subplots, Iterable):
+
+            supported_kinds = (
+                "line",
+                "bar",
+                "barh",
+                "hist",
+                "kde",
+                "density",
+                "area",
+                "pie",
+            )
+            if self._kind not in supported_kinds:
+                raise ValueError(
+                    "When subplots is an iterable, kind must be "
+                    f"one of {', '.join(supported_kinds)}. Got {self._kind}"
+                )
+
+            if any(
+                not isinstance(group, Iterable) or isinstance(group, str)
+                for group in subplots
+            ):
+                raise ValueError(
+                    "When subplots is an iterable, each entry "
+                    "should be a list/tuple of column names "
+                    "or column indices."
+                )
+
+            cols_in_groups = {col for group in self.subplots for col in group}
+            cols_remaining = set(data.columns) - cols_in_groups
+            bad_columns = cols_in_groups - set(data.columns)
+
+            if bad_columns:
+                raise ValueError(
+                    "Subplots contains the following column(s) "
+                    f"which are invalid names: {bad_columns}"
+                )
+
+            # subplots is a list of tuples where each tuple is a group of
+            # columns to be grouped together (one ax per group).
+            # we consolidate the subplots list such that:
+            # - the tuples contain indexes instead of column names
+            # - the columns that aren't yet in the list are added in a group
+            #   of their own.
+            # For example with columns from a to g, and
+            # subplots = [(a, c), (b, f, e)],
+            # we end up with [(ai, ci), (bi, fi, ei), (di,), (gi,)]
+            # This way, we can handle self.subplots in a homogeneous manner
+            # later.
+            # TODO: also accept indexes instead of just names?
+
+            subplots = []
+            index = list(data.columns).index
+            for group in self.subplots:
+                subplots.append(tuple(index(col) for col in group))
+            for col in cols_remaining:
+                subplots.append((index(col),))
+
+            self.subplots = subplots
 
         if sharex is None:
             if ax is None:
@@ -306,8 +368,11 @@ def _maybe_right_yaxis(self, ax, axes_num):
 
     def _setup_subplots(self):
         if self.subplots:
+            naxes = (
+                self.nseries if isinstance(self.subplots, bool) else len(self.subplots)
+            )
             fig, axes = _subplots(
-                naxes=self.nseries,
+                naxes=naxes,
                 sharex=self.sharex,
                 sharey=self.sharey,
                 figsize=self.figsize,
@@ -676,9 +741,21 @@ def _get_ax_layer(cls, ax, primary=True):
         else:
             return getattr(ax, "right_ax", ax)
 
+    def _col_idx_to_axis_idx(self, col_idx):
+        """Return the index of the axis where the column at col_idx should be plotted"""
+        if isinstance(self.subplots, list):
+            # Subplots is a list: some columns are be grouped together in the same ax
+            for group_idx, group in enumerate(self.subplots):
+                if col_idx in group:
+                    return group_idx
+        else:
+            # subplots is True: one ax per column
+            return col_idx
+
     def _get_ax(self, i):
         # get the twinx ax if appropriate
         if self.subplots:
+            i = self._col_idx_to_axis_idx(i)
             ax = self.axes[i]
             ax = self._maybe_right_yaxis(ax, i)
             self.axes[i] = ax

pandas/tests/plotting/test_frame.py

@@ -408,6 +408,9 @@ def test_subplots(self):
             for ax in axes:
                 assert ax.get_legend() is None
 
+        with pytest.raises(ValueError, match="should be a bool or an iterable"):
+            axes = df.plot(subplots=123)
+
     def test_groupby_boxplot_sharey(self):
         # https://github.com/pandas-dev/pandas/issues/20968
         # sharey can now be switched check whether the right
@@ -3279,6 +3282,64 @@ def test_plot_no_numeric_data(self):
         with pytest.raises(TypeError):
             df.plot()
 
+    @td.skip_if_no_scipy
+    @pytest.mark.parametrize(
+        "kind", ("line", "bar", "barh", "hist", "kde", "density", "area", "pie")
+    )
+    def test_group_subplot(self, kind):
+        d = {
+            "a": np.arange(10),
+            "b": np.arange(10) + 1,
+            "c": np.arange(10) + 1,
+            "d": np.arange(10),
+            "e": np.arange(10),
+        }
+        df = pd.DataFrame(d)
+
+        axes = df.plot(subplots=[("b", "e"), ("c", "d")], kind=kind)
+        assert len(axes) == 3  # 2 groups + single column a
+
+        expected_labels = (["b", "e"], ["c", "d"], ["a"])
+        for ax, labels in zip(axes, expected_labels):
+            if kind != "pie":
+                self._check_legend_labels(ax, labels=labels)
+            if kind == "line":
+                assert len(ax.lines) == len(labels)
+
+    @pytest.mark.parametrize(
+        "subplots",
+        [
+            "a",  # iterable of non-iterable
+            (1,),  # iterable of non-iterable
+            ("a",),  # iterable of strings
+        ],
+    )
+    def test_group_subplot_bad_input(self, subplots):
+        # Make sure error is raised when subplots is not a properly
+        # formatted iterable. Only iterables of iterables are permitted, and
+        # entries should not be strings.
+        d = {"a": np.arange(10), "b": np.arange(10)}
+        df = pd.DataFrame(d)
+
+        with pytest.raises(ValueError, match="each entry should be a list/tuple"):
+            df.plot(subplots=subplots)
+
+    def test_group_subplot_invalid_column_name(self):
+        d = {"a": np.arange(10), "b": np.arange(10)}
+        df = pd.DataFrame(d)
+
+        with pytest.raises(ValueError, match="invalid names: {'bad_name'}"):
+            df.plot(subplots=[("a", "bad_name")])
+
+    @pytest.mark.parametrize("kind", ("box", "scatter", "hexbin"))
+    def test_group_subplot_invalid_kind(self, kind):
+        d = {"a": np.arange(10), "b": np.arange(10)}
+        df = pd.DataFrame(d)
+        with pytest.raises(
+            ValueError, match="When subplots is an iterable, kind must be one of"
+        ):
+            df.plot(subplots=[("a", "b")], kind=kind)
+
     def test_missing_markers_legend(self):
         # 14958
         df = pd.DataFrame(np.random.randn(8, 3), columns=["A", "B", "C"])