Repository: pandas-dev/pandas
Pull request: ENH: Allow column grouping in DataFrame.plot
- [X] closes #29688
- [x] tests added / passed
- [x] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry

This is an attempt at #29688 i.e. group some columns on the same ax when calling `df.plot`

See [this gist](https://gist.github.com/NicolasHug/4d1f60212ba69e6000982951ae47f203) for illustration of different supported plots

Thanks!

Problem statement:
ENH: Group some columns in subplots with DataFrame.plot
`df.plot(subplots=True)`  will create one subplot per column. Is there a way to group multiple columns on the same subplot (and leave the rest of the column separated)?

I'd be happy to submit a PR if that's something you'd consider? In terms of API  `subplot` could accept a list of tuples where each tuple indicates which columns should be grouped together.

Hints from issue discussion:
Just to make sure I understand, `df.plot(subplots=[('a', 'b'), ('d', 'e', 'f')])` would have two axes?

I'm not sure, but I'm guessing it'd be somewhat complicated to work that into our current `subplots` logic. At glance, it looks like we assume `subplots=True` -> one axes per column in several places.
> `df.plot(subplots=[('a', 'b'), ('d', 'e', 'f')])` would have two axes?

Yes that's the idea. My motivation is that sometimes some columns are comparable, in which case it would make sense to plot them together.

Commit message: linting again

Patch to review:
diff --git a/pandas/plotting/_matplotlib/core.py b/pandas/plotting/_matplotlib/core.py
index 2d68bb46a8ada..210ecdd4e3590 100644
--- a/pandas/plotting/_matplotlib/core.py
+++ b/pandas/plotting/_matplotlib/core.py
@@ -1,3 +1,4 @@
+from collections import Iterable
 import re
 from typing import Optional
 import warnings
@@ -115,6 +116,31 @@ def __init__(
         self.sort_columns = sort_columns
 
         self.subplots = subplots
+        if isinstance(self.subplots, Iterable):
+            # 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?
+            # TODO: we're potentially messing with the order of the axes here
+            cols_in_groups = {col for group in self.subplots for col in group}
+            cols_remaining = set(data.columns) - cols_in_groups
+
+            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 +332,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,
@@ -677,9 +706,20 @@ def _get_ax_layer(cls, ax, primary=True):
         else:
             return getattr(ax, "right_ax", ax)
 
+    def _col_idx_to_axis_idx(self, i):
+        if isinstance(self.subplots, list):
+            # Some columns are be grouped together in the same ax
+            for group_idx, group in enumerate(self.subplots):
+                if i in group:
+                    return group_idx
+        else:
+            # One ax per column
+            return i
+
     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
diff --git a/pandas/tests/plotting/test_frame.py b/pandas/tests/plotting/test_frame.py
index 1c429bafa9a19..cdc9dd7e53d2b 100644
--- a/pandas/tests/plotting/test_frame.py
+++ b/pandas/tests/plotting/test_frame.py
@@ -3271,6 +3271,19 @@ def test_plot_no_numeric_data(self):
         with pytest.raises(TypeError):
             df.plot()
 
+    def test_group_subplot(self):
+        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)
+
+        ax = df.plot(subplots=[("b", "e"), ("c", "d")])
+        assert len(ax) == 3  # 2 groups + single column a
+
     def test_missing_markers_legend(self):
         # 14958
         df = pd.DataFrame(np.random.randn(8, 3), columns=["A", "B", "C"])


Code review comments:
diff_hunk: @@ -115,6 +116,31 @@ def __init__(
         self.sort_columns = sort_columns
 
         self.subplots = subplots
+        if isinstance(self.subplots, Iterable):
+            # 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?
+            # TODO: we're potentially messing with the order of the axes here
+            cols_in_groups = {col for group in self.subplots for col in group}
path: pandas/plotting/_matplotlib/core.py
text: @user1:
maybe better to raise other than silence if wrong col names are assigned?

diff_hunk: @@ -3271,6 +3271,19 @@ def test_plot_no_numeric_data(self):
         with pytest.raises(TypeError):
             df.plot()
 
+    def test_group_subplot(self):
+        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)
+
+        ax = df.plot(subplots=[("b", "e"), ("c", "d")])
path: pandas/tests/plotting/test_frame.py
text: @user1:
minor: maybe use `axes` is more appropriate here since it has three ax in this.

diff_hunk: @@ -115,6 +116,31 @@ def __init__(
         self.sort_columns = sort_columns
 
         self.subplots = subplots
+        if isinstance(self.subplots, Iterable):
+            # 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?
+            # TODO: we're potentially messing with the order of the axes here
+            cols_in_groups = {col for group in self.subplots for col in group}
+            cols_remaining = set(data.columns) - cols_in_groups
+
+            subplots = []
+            index = list(data.columns).index
+            for group in self.subplots:
+                subplots.append(tuple(index(col) for col in group))
path: pandas/plotting/_matplotlib/core.py
text: @user1:
like said above, this will fail if `subplots` is not a nested iterable.

diff_hunk: @@ -115,6 +116,31 @@ def __init__(
         self.sort_columns = sort_columns
 
         self.subplots = subplots
+        if isinstance(self.subplots, Iterable):
path: pandas/plotting/_matplotlib/core.py
text: @user1:
i think this would also need to check the element in `subplots`? because we probably will take nested `Iterable` here, e.g. `[('a', 'c'), ('b', 'd')]`?  or we need to convert to this for example: `[('a', 'c'), ('b', 'd'), 'e']`, if users specify like this (although they shouldn't), we should decide what to do with such cases? should we convert it to `[('a', 'c'), ('b', 'd'), ('e', )]` or not?