ENH: DataFrame plot improvements, make pivot_table more flexible per #941

Author: wesm

pandas/core/frame.py

@@ -3845,8 +3845,9 @@ def boxplot(self, column=None, by=None, ax=None, fontsize=None,
         return ax
 
     def plot(self, subplots=False, sharex=True, sharey=False, use_index=True,
-             figsize=None, grid=True, legend=True, rot=30, ax=None,
-             kind='line', sort_columns=True, **kwds):
+             figsize=None, grid=True, legend=True, rot=30, ax=None, title=None,
+             xlim=None, ylim=None, xticks=None, yticks=None, kind='line',
+             sort_columns=True, fontsize=None, **kwds):
         """
         Make line plot of DataFrame's series with the index on the x-axis using
         matplotlib / pylab.
@@ -3934,22 +3935,37 @@ def plot(self, subplots=False, sharex=True, sharey=False, use_index=True,
                     ax_.set_xticklabels(xticklabels, rotation=rot)
         elif kind == 'bar':
             self._bar_plot(axes, subplots=subplots, grid=grid, rot=rot,
-                           legend=legend)
+                           legend=legend, ax=ax, fontsize=fontsize)
 
-        if not subplots or (subplots and sharex):
+        if self.index.is_all_dates and not subplots or (subplots and sharex):
             try:
                 fig.autofmt_xdate()
             except Exception:  # pragma: no cover
                 pass
 
+        if yticks is not None:
+            ax.set_yticks(yticks)
+
+        if xticks is not None:
+            ax.set_xticks(xticks)
+
+        if ylim is not None:
+            ax.set_ylim(ylim)
+
+        if xlim is not None:
+            ax.set_xlim(xlim)
+
+        if title and not subplots:
+            ax.set_title(title)
+
         plt.draw_if_interactive()
         if subplots:
             return axes
         else:
             return ax
 
     def _bar_plot(self, axes, subplots=False, use_index=True, grid=True,
-                  rot=30, legend=True, **kwds):
+                  rot=30, legend=True, ax=None, fontsize=None, **kwds):
         import pandas.tools.plotting as gfx
 
         N, K = self.shape
@@ -3958,7 +3974,7 @@ def _bar_plot(self, axes, subplots=False, use_index=True, grid=True,
         rects = []
         labels = []
 
-        if not subplots:
+        if not subplots and ax is None:
             ax = axes[0]
 
         for i, col in enumerate(self.columns):
@@ -3970,17 +3986,20 @@ def _bar_plot(self, axes, subplots=False, use_index=True, grid=True,
                        bottom=np.zeros(N), linewidth=1, **kwds)
                 ax.set_title(col)
             else:
-                rects.append(ax.bar(xinds + i * 0.5 / K, y, 0.5 / K,
+                rects.append(ax.bar(xinds + i * 0.75 / K, y, 0.75 / K,
                                     bottom=np.zeros(N), label=col,
                                     color=colors[i % len(colors)], **kwds))
                 labels.append(col)
 
-        if N < 10:
-            fontsize = 12
-        else:
-            fontsize = 10
+        if fontsize is None:
+            if N < 10:
+                fontsize = 12
+            else:
+                fontsize = 10
+
+        ax.set_xlim([xinds[0] - 1, xinds[-1] + 1])
 
-        ax.set_xticks(xinds + 0.25)
+        ax.set_xticks(xinds + 0.375)
         ax.set_xticklabels([gfx._stringify(key) for key in self.index],
                            rotation=rot,
                            fontsize=fontsize)

pandas/core/groupby.py

@@ -342,7 +342,7 @@ def _python_apply_general(self, func, *args, **kwargs):
 
         not_indexed_same = False
         for key, group in self:
-            group.name = key
+            object.__setattr__(group, 'name', key)
 
             # group might be modified
             group_axes = _get_axes(group)
@@ -1003,7 +1003,7 @@ def transform(self, func, *args, **kwargs):
         result = self.obj.copy()
 
         for name, group in self:
-            group.name = name
+            object.__setattr__(group, 'name', name)
             res = func(group, *args, **kwargs)
             indexer = self.obj.index.get_indexer(group.index)
             np.put(result, indexer, res)
@@ -1363,7 +1363,7 @@ def transform(self, func, *args, **kwargs):
 
         obj = self._obj_with_exclusions
         for name, group in self:
-            group.name = name
+            object.__setattr__(group, 'name', name)
 
             try:
                 wrapper = lambda x: func(x, *args, **kwargs)

pandas/tools/pivot.py

@@ -4,6 +4,7 @@
 from pandas.tools.merge import concat
 import pandas.core.common as com
 import numpy as np
+import types
 
 def pivot_table(data, values=None, rows=None, cols=None, aggfunc='mean',
                 fill_value=None, margins=False):
@@ -16,10 +17,10 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc='mean',
     ----------
     data : DataFrame
     values : column to aggregate, optional
-    rows : list
-        Columns to group on the x-axis of the pivot table
-    cols : list
-        Columns to group on the x-axis of the pivot table
+    rows : list of column names or arrays to group on
+        Keys to group on the x-axis of the pivot table
+    cols : list of column names or arrays to group on
+        Keys to group on the x-axis of the pivot table
     aggfunc : function, default numpy.mean, or list of functions
         If list of functions passed, the resulting pivot table will have
         hierarchical columns whose top level are the function names (inferred
@@ -83,14 +84,23 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc='mean',
         values = list(data.columns.drop(keys))
 
     if values_passed:
-        data = data[keys + values]
+        to_filter = []
+        for x in keys + values:
+            try:
+                if x in data:
+                    to_filter.append(x)
+            except TypeError:
+                pass
+        if len(to_filter) < len(data.columns):
+            data = data[to_filter]
 
     grouped = data.groupby(keys)
     agged = grouped.agg(aggfunc)
 
     table = agged
-    for k in cols:
-        table = table.unstack(level=k)
+    for i in range(len(cols)):
+        name = table.index.names[len(rows)]
+        table = table.unstack(name)
 
     if fill_value is not None:
         table = table.fillna(value=fill_value)
@@ -183,7 +193,8 @@ def _all_key(key):
 def _convert_by(by):
     if by is None:
         by = []
-    elif np.isscalar(by):
+    elif (np.isscalar(by) or isinstance(by, np.ndarray)
+          or hasattr(by, '__call__')):
         by = [by]
     else:
         by = list(by)

pandas/tools/tests/test_pivot.py

@@ -47,6 +47,18 @@ def test_pivot_table(self):
         expected = self.data.groupby(rows + [cols])['D'].agg(np.mean).unstack()
         tm.assert_frame_equal(table, expected)
 
+    def test_pass_array(self):
+        result = self.data.pivot_table('D', rows=self.data.A, cols=self.data.C)
+        expected = self.data.pivot_table('D', rows='A', cols='C')
+        tm.assert_frame_equal(result, expected)
+
+    def test_pass_function(self):
+        result = self.data.pivot_table('D', rows=lambda x: x // 5,
+                                       cols=self.data.C)
+        expected = self.data.pivot_table('D', rows=self.data.index // 5,
+                                         cols='C')
+        tm.assert_frame_equal(result, expected)
+
     def test_pivot_table_multiple(self):
         rows = ['A', 'B']
         cols=  'C'

setup.py

@@ -166,8 +166,8 @@
 
 MAJOR = 0
 MINOR = 7
-MICRO = 2
-ISRELEASED = True
+MICRO = 3
+ISRELEASED = False
 VERSION = '%d.%d.%d' % (MAJOR, MINOR, MICRO)
 QUALIFIER = ''