ENH: Groupby.plot enhancement

Author: sinhrks

pandas/core/groupby.py

@@ -3297,9 +3297,13 @@ def _apply_to_column_groupbys(self, func):
              in self._iterate_column_groupbys()),
             keys=self._selected_obj.columns, axis=1)
 
-from pandas.tools.plotting import boxplot_frame_groupby
+
+from pandas.tools.plotting import boxplot_frame_groupby, plot_grouped
 DataFrameGroupBy.boxplot = boxplot_frame_groupby
 
+SeriesGroupBy.plot = plot_grouped
+DataFrameGroupBy.plot = plot_grouped
+
 
 class PanelGroupBy(NDFrameGroupBy):
 

pandas/tests/test_graphics.py

@@ -1671,6 +1671,7 @@ def test_line_lim(self):
             self.assertFalse(hasattr(ax, 'right_ax'))
             xmin, xmax = ax.get_xlim()
             lines = ax.get_lines()
+            self.assertTrue(hasattr(ax, 'left_ax'))
             self.assertEqual(xmin, lines[0].get_data()[0][0])
             self.assertEqual(xmax, lines[0].get_data()[0][-1])
 
@@ -1695,10 +1696,8 @@ def test_area_lim(self):
     @slow
     def test_bar_colors(self):
         import matplotlib.pyplot as plt
-
         default_colors = plt.rcParams.get('axes.color_cycle')
 
-
         df = DataFrame(randn(5, 5))
         ax = df.plot(kind='bar')
         self._check_colors(ax.patches[::5], facecolors=default_colors[:5])
@@ -3209,7 +3208,6 @@ def test_hexbin_cmap(self):
     @slow
     def test_no_color_bar(self):
         df = self.hexbin_df
-
         ax = df.plot(kind='hexbin', x='A', y='B', colorbar=None)
         self.assertIs(ax.collections[0].colorbar, None)
 
@@ -3572,9 +3570,8 @@ def test_grouped_plot_fignums(self):
         df = DataFrame({'height': height, 'weight': weight, 'gender': gender})
         gb = df.groupby('gender')
 
-        res = gb.plot()
-        self.assertEqual(len(self.plt.get_fignums()), 2)
-        self.assertEqual(len(res), 2)
+        with tm.assertRaisesRegexp(ValueError, "To plot DataFrameGroupBy, specify 'suplots=True'"):
+            res = gb.plot()
         tm.close()
 
         res = gb.boxplot(return_type='axes')
@@ -3953,6 +3950,102 @@ def test_plotting_with_float_index_works(self):
         df.groupby('def')['val'].apply(lambda x: x.plot())
         tm.close()
 
+    def test_line_groupby(self):
+        df = DataFrame(np.random.rand(30, 5), columns=['A', 'B', 'C', 'D', 'E'])
+        df['by'] = ['Group {0}'.format(i) for i in [0]*10 + [1]*10 + [2]*10]
+        grouped = df.groupby(by='by')
+
+        # SeriesGroupBy
+        sgb = grouped['A']
+        ax = _check_plot_works(sgb.plot, colors=['r', 'g', 'b'])
+        self._check_legend_labels(ax, labels=['Group 0', 'Group 1', 'Group 2'])
+        self._check_colors(ax.get_lines(), linecolors=['r', 'g', 'b'])
+
+        axes = _check_plot_works(sgb.plot, subplots=True)
+        self._check_axes_shape(axes, axes_num=3, layout=(3, 1))
+        self._check_legend_labels(axes[0], labels=['Group 0'])
+        self._check_legend_labels(axes[1], labels=['Group 1'])
+        self._check_legend_labels(axes[2], labels=['Group 2'])
+
+        # DataFrameGroupBy
+        import matplotlib.pyplot as plt
+        default_colors = plt.rcParams.get('axes.color_cycle')
+
+        axes = _check_plot_works(grouped.plot, subplots=True)
+        self._check_axes_shape(axes, axes_num=3, layout=(3, 1))
+        for ax in axes:
+            self._check_legend_labels(ax, labels=['A', 'B', 'C', 'D', 'E'])
+            self._check_colors(ax.get_lines(), linecolors=default_colors[:5])
+
+        axes = _check_plot_works(grouped.plot, subplots=True, axis=1)
+        self._check_axes_shape(axes, axes_num=5, layout=(5, 1))
+        for ax in axes:
+            self._check_legend_labels(ax, labels=['Group 0', 'Group 1', 'Group 2'])
+            self._check_colors(ax.get_lines(), linecolors=default_colors[:3])
+
+    def test_hist_groupby(self):
+        df = DataFrame(np.random.rand(30, 5), columns=['A', 'B', 'C', 'D', 'E'])
+        df['by'] = ['Group {0}'.format(i) for i in [0]*10 + [1]*10 + [2]*10]
+        grouped = df.groupby(by='by')
+
+        # SeriesGroupBy
+        sgb = grouped['A']
+        ax = sgb.plot(kind='hist', color=['r', 'g', 'b'])
+        self._check_legend_labels(ax, labels=['Group 0', 'Group 1', 'Group 2'])
+        self._check_colors(ax.patches[::10], facecolors=['r', 'g', 'b'])
+
+        import matplotlib.pyplot as plt
+        default_colors = plt.rcParams.get('axes.color_cycle')
+        axes = sgb.plot(kind='hist', subplots=True)
+        self._check_axes_shape(axes, axes_num=3, layout=(3, 1))
+        self._check_legend_labels(axes[0], labels=['Group 0'])
+        self._check_legend_labels(axes[1], labels=['Group 1'])
+        self._check_legend_labels(axes[2], labels=['Group 2'])
+        self._check_colors([axes[0].patches[0]], facecolors=default_colors[0])
+        self._check_colors([axes[1].patches[0]], facecolors=default_colors[1])
+        self._check_colors([axes[2].patches[0]], facecolors=default_colors[2])
+
+        # DataFrameGroupBy
+        axes = grouped.plot(kind='hist', subplots=True)
+        self._check_axes_shape(axes, axes_num=3, layout=(3, 1))
+        for ax in axes:
+            self._check_legend_labels(ax, labels=['A', 'B', 'C', 'D', 'E'])
+        self._check_colors(axes[0].patches[::10], facecolors=default_colors[:5])
+        self._check_colors(axes[1].patches[::10], facecolors=default_colors[:5])
+        self._check_colors(axes[2].patches[::10], facecolors=default_colors[:5])
+
+        axes = grouped.plot(kind='hist', subplots=True, axis=1)
+        self._check_axes_shape(axes, axes_num=5, layout=(5, 1))
+        for ax in axes:
+            self._check_legend_labels(ax, labels=['Group 0', 'Group 1', 'Group 2'])
+            self._check_colors(ax.patches[::10], facecolors=['b', 'g', 'r'])
+
+    def test_scatter_groupby(self):
+        df = DataFrame(np.random.rand(30, 5), columns=['A', 'B', 'C', 'D', 'E'])
+        df['by'] = ['Group {0}'.format(i) for i in [0]*10 + [1]*10 + [2]*10]
+        grouped = df.groupby(by='by')
+
+        ax = _check_plot_works(grouped.plot, kind='scatter', x='A', y='B', subplots=False)
+        self._check_legend_labels(ax, labels=['Group 0', 'Group 1', 'Group 2'])
+
+        axes = _check_plot_works(grouped.plot, kind='scatter', x='A', y='B', subplots=True)
+        self._check_axes_shape(axes, axes_num=3, layout=(1, 3))
+        self._check_legend_labels(axes[0], labels=['Group 0'])
+        self._check_legend_labels(axes[1], labels=['Group 1'])
+        self._check_legend_labels(axes[2], labels=['Group 2'])
+
+    def test_hexbin_groupby(self):
+        df = DataFrame(np.random.rand(30, 5), columns=['A', 'B', 'C', 'D', 'E'])
+        df['by'] = ['Group {0}'.format(i) for i in [0]*10 + [1]*10 + [2]*10]
+        grouped = df.groupby(by='by')
+
+        msg = "To plot DataFrameGroupBy, specify 'suplots=True'"
+        with tm.assertRaisesRegexp(ValueError, msg):
+            grouped.plot(kind='hexbin', x='A', y='B', subplots=False)
+
+        axes = _check_plot_works(grouped.plot, kind='hexbin', x='A', y='B', subplots=True)
+        self._check_axes_shape(axes, axes_num=3, layout=(1, 3))
+
 
 def assert_is_valid_plot_return_object(objs):
     import matplotlib.pyplot as plt