ENH: DataFrame.stack() with 'level' a set or list of sets

Author: seth-p

pandas/core/index.py

@@ -820,6 +820,8 @@ def _validate_index_level(self, level):
                            % (level, self.name))
 
     def _get_level_number(self, level):
+        if isinstance(level, set):
+            return set([self._get_level_number(lev) for lev in level])
         self._validate_index_level(level)
         return 0
 
@@ -3157,6 +3159,8 @@ def _from_elements(values, labels=None, levels=None, names=None,
         return MultiIndex(levels, labels, names, sortorder=sortorder)
 
     def _get_level_number(self, level):
+        if isinstance(level, set):
+            return set([self._get_level_number(lev) for lev in level])
         try:
             count = self.names.count(level)
             if count > 1:

pandas/core/reshape.py

@@ -508,6 +508,14 @@ def get_compressed_ids(labels, sizes):
     return comp_index, obs_ids
 
 
+def _iterate_through_set(x):
+    if isinstance(x, set):
+        for y in x:
+            yield y
+    else:
+        yield x
+
+
 def stack(frame, level=-1, dropna=True):
     """
     Convert DataFrame to Series with multi-level Index. Columns become the
@@ -517,19 +525,18 @@ def stack(frame, level=-1, dropna=True):
     -------
     stacked : Series
     """
-    N, K = frame.shape
     if isinstance(frame.columns, MultiIndex):
-        if frame.columns._reference_duplicate_name(level):
+        if any(frame.columns._reference_duplicate_name(lev)
+               for lev in _iterate_through_set(level)):
             msg = ("Ambiguous reference to {0}. The column "
                    "names are not unique.".format(level))
             raise ValueError(msg)
-
-    # Will also convert negative level numbers and check if out of bounds.
-    level_num = frame.columns._get_level_number(level)
-
-    if isinstance(frame.columns, MultiIndex):
+        # Will also convert negative level numbers and check if out of bounds.
+        level_num = frame.columns._get_level_number(level)
         return _stack_multi_columns(frame, level_num=level_num, dropna=dropna)
-    elif isinstance(frame.index, MultiIndex):
+
+    N, K = frame.shape
+    if isinstance(frame.index, MultiIndex):
         new_levels = list(frame.index.levels)
         new_levels.append(frame.columns)
 
@@ -559,13 +566,13 @@ def stack(frame, level=-1, dropna=True):
 def stack_multiple(frame, level, dropna=True):
     # If all passed levels match up to column names, no
     # ambiguity about what to do
-    if all(lev in frame.columns.names for lev in level):
+    if all(lev in frame.columns.names for levl in level for lev in _iterate_through_set(levl)):
         result = frame
         for lev in level:
             result = stack(result, lev, dropna=dropna)
 
     # Otherwise, level numbers may change as each successive level is stacked
-    elif all(isinstance(lev, int) for lev in level):
+    elif all(isinstance(lev, int) for levl in level for lev in _iterate_through_set(levl)):
         # As each stack is done, the level numbers decrease, so we need
         #  to account for that when level is a sequence of ints
         result = frame
@@ -576,16 +583,19 @@ def stack_multiple(frame, level, dropna=True):
         # Can't iterate directly through level as we might need to change
         # values as we go
         for index in range(len(level)):
-            lev = level[index]
-            result = stack(result, lev, dropna=dropna)
+            levl = level[index]
+            result = stack(result, levl, dropna=dropna)
             # Decrement all level numbers greater than current, as these
-            # have now shifted down by one
+            # have now shifted down
             updated_level = []
             for other in level:
-                if other > lev:
-                    updated_level.append(other - 1)
+                if isinstance(other, set):
+                    updated_level.append({(othr - sum((othr > lev)
+                                                      for lev in _iterate_through_set(levl)))
+                                          for othr in other})
                 else:
-                    updated_level.append(other)
+                    updated_level.append(other - sum((other > lev)
+                                                     for lev in _iterate_through_set(levl)))
             level = updated_level
 
     else:
@@ -616,85 +626,101 @@ def _convert_level_number(level_num, columns):
     this = frame.copy()
 
     # this makes life much simpler
-    if level_num != frame.columns.nlevels - 1:
-        # roll levels to put selected level at end
-        roll_columns = this.columns
-        for i in range(level_num, frame.columns.nlevels - 1):
+    # roll levels to put selected level(s) at end
+    level_nums = level_num if isinstance(level_num, set) else {level_num}
+    roll_columns = this.columns
+    for j, level_num in enumerate(sorted(level_nums, reverse=True)):
+        for i in range(level_num, frame.columns.nlevels - (j + 1)):
             # Need to check if the ints conflict with level names
             lev1 = _convert_level_number(i, roll_columns)
             lev2 = _convert_level_number(i + 1, roll_columns)
             roll_columns = roll_columns.swaplevel(lev1, lev2)
-        this.columns = roll_columns
+    this.columns = roll_columns
 
     if not this.columns.is_lexsorted():
         # Workaround the edge case where 0 is one of the column names,
-        # which interferes with trying to sort based on the first
-        # level
+        # which interferes with trying to sort based on the first level
         level_to_sort = _convert_level_number(0, this.columns)
         this = this.sortlevel(level_to_sort, axis=1)
 
-    # tuple list excluding level for grouping columns
-    if len(frame.columns.levels) > 2:
+    num_levels_to_stack = len(level_nums)
+    level_vals = this.columns.levels[-num_levels_to_stack:]
+    level_labels = sorted(set(zip(*this.columns.labels[-num_levels_to_stack:])))
+    level_vals_used = MultiIndex.from_tuples([tuple(level_vals[i][lab] for i, lab in enumerate(label))
+                                              for label in level_labels],
+                                             names=this.columns.names[-num_levels_to_stack:])
+    levsize = len(level_labels)
+
+    # construct new_index
+    N = len(this)
+    if isinstance(this.index, MultiIndex):
+        new_levels = list(this.index.levels)
+        new_names = list(this.index.names)
+        new_labels = [lab.repeat(levsize) for lab in this.index.labels]
+    else:
+        new_levels = [this.index]
+        new_labels = [np.arange(N).repeat(levsize)]
+        new_names = [this.index.name]  # something better?
+    new_levels += level_vals
+    new_labels += [np.tile(labels, N) for labels in zip(*level_labels)]
+    new_names += level_vals_used.names
+    new_index = MultiIndex(levels=new_levels, labels=new_labels,
+                           names=new_names, verify_integrity=False)
+
+    # if stacking all levels in columns, result will be a Series
+    if len(frame.columns.levels) == num_levels_to_stack:
+        new_data = frame.values.ravel()
+        if dropna:
+            mask = notnull(new_data)
+            new_data = new_data[mask]
+            new_index = new_index[mask]
+        return Series(new_data, index=new_index)
+
+    # result will be a DataFrame
+
+    # construct new_columns
+    if len(frame.columns.levels) > (num_levels_to_stack + 1):
+        # result columns will be a MultiIndex
+        # tuple list excluding level for grouping columns
         tuples = list(zip(*[
             lev.take(lab) for lev, lab in
-            zip(this.columns.levels[:-1], this.columns.labels[:-1])
+            zip(this.columns.levels[:-num_levels_to_stack],
+                this.columns.labels[:-num_levels_to_stack])
         ]))
         unique_groups = [key for key, _ in itertools.groupby(tuples)]
-        new_names = this.columns.names[:-1]
+        new_names = this.columns.names[:-num_levels_to_stack]
         new_columns = MultiIndex.from_tuples(unique_groups, names=new_names)
     else:
+        # result columns will be an Index
         new_columns = unique_groups = this.columns.levels[0]
 
-    # time to ravel the values
+    # construct new_data
     new_data = {}
-    level_vals = this.columns.levels[-1]
-    level_labels = sorted(set(this.columns.labels[-1]))
-    level_vals_used = level_vals[level_labels]
-    levsize = len(level_labels)
     drop_cols = []
     for key in unique_groups:
         loc = this.columns.get_loc(key)
         slice_len = loc.stop - loc.start
         # can make more efficient?
-
         if slice_len == 0:
             drop_cols.append(key)
             continue
         elif slice_len != levsize:
             chunk = this.ix[:, this.columns[loc]]
-            chunk.columns = level_vals.take(chunk.columns.labels[-1])
+            chunk.columns = MultiIndex.from_arrays([vals.take(labels) for (vals, labels)
+                                                    in zip(level_vals, chunk.columns.labels[-num_levels_to_stack:])],
+                                                   names=chunk.columns.names[-num_levels_to_stack:])
             value_slice = chunk.reindex(columns=level_vals_used).values
         else:
             if frame._is_mixed_type:
                 value_slice = this.ix[:, this.columns[loc]].values
             else:
                 value_slice = this.values[:, loc]
-
         new_data[key] = value_slice.ravel()
 
     if len(drop_cols) > 0:
         new_columns = new_columns - drop_cols
 
-    N = len(this)
-
-    if isinstance(this.index, MultiIndex):
-        new_levels = list(this.index.levels)
-        new_names = list(this.index.names)
-        new_labels = [lab.repeat(levsize) for lab in this.index.labels]
-    else:
-        new_levels = [this.index]
-        new_labels = [np.arange(N).repeat(levsize)]
-        new_names = [this.index.name]  # something better?
-
-    new_levels.append(frame.columns.levels[level_num])
-    new_labels.append(np.tile(level_labels, N))
-    new_names.append(frame.columns.names[level_num])
-
-    new_index = MultiIndex(levels=new_levels, labels=new_labels,
-                           names=new_names, verify_integrity=False)
-
     result = DataFrame(new_data, index=new_index, columns=new_columns)
-
     # more efficient way to go about this? can do the whole masking biz but
     # will only save a small amount of time...
     if dropna: