ENH: reindex_multi in Panel per pandas-dev#979

wesm · wesm · commit 5dfb2ee3efd7 · 2012-07-12T22:12:04.000-04:00
diff --git a/RELEASE.rst b/RELEASE.rst
@@ -52,6 +52,7 @@ pandas 0.8.1
   - Improve performance of array-of-Period to PeriodIndex, convert such arrays
     to PeriodIndex inside Index (#1215)
   - More informative string representation for weekly Period objects (#1503)
+  - Accelerate 3-axis multi data selection from homogeneous Panel (#979)
 
 **Bug fixes**
 
diff --git a/pandas/core/common.py b/pandas/core/common.py
@@ -254,7 +254,7 @@ def take_1d(arr, indexer, out=None, fill_value=np.nan):
 
     return out
 
-def take_2d_multi(arr, row_idx, col_idx, fill_value=np.nan):
+def take_2d_multi(arr, row_idx, col_idx, fill_value=np.nan, out=None):
 
     dtype_str = arr.dtype.name
 
@@ -267,21 +267,26 @@ def take_2d_multi(arr, row_idx, col_idx, fill_value=np.nan):
 
         if needs_masking:
             return take_2d_multi(_maybe_upcast(arr), row_idx, col_idx,
-                                 fill_value=fill_value)
+                                 fill_value=fill_value, out=out)
         else:
-            out = np.empty(out_shape, dtype=arr.dtype)
+            if out is None:
+                out = np.empty(out_shape, dtype=arr.dtype)
             take_f = _get_take2d_function(dtype_str, axis='multi')
             take_f(arr, _ensure_int64(row_idx),
                    _ensure_int64(col_idx), out=out,
                    fill_value=fill_value)
             return out
     elif dtype_str in ('float64', 'object', 'datetime64[ns]'):
-        out = np.empty(out_shape, dtype=arr.dtype)
+        if out is None:
+            out = np.empty(out_shape, dtype=arr.dtype)
         take_f = _get_take2d_function(dtype_str, axis='multi')
         take_f(arr, _ensure_int64(row_idx), _ensure_int64(col_idx), out=out,
                fill_value=fill_value)
         return out
     else:
+        if out is not None:
+            raise ValueError('Cannot pass out in this case')
+
         return take_2d(take_2d(arr, row_idx, axis=0, fill_value=fill_value),
                        col_idx, axis=1, fill_value=fill_value)
 
diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py
@@ -141,26 +141,34 @@ def _getitem_tuple(self, tup):
         return retval
 
     def _multi_take_opportunity(self, tup):
-        from pandas.core.frame import DataFrame
+        from pandas.core.generic import NDFrame
 
         # ugly hack for GH #836
-        if not isinstance(self.obj, DataFrame):
+        if not isinstance(self.obj, NDFrame):
             return False
 
         if not all(_is_list_like(x) for x in tup):
             return False
 
         # just too complicated
-        if (isinstance(self.obj.index, MultiIndex) or
-            isinstance(self.obj.columns, MultiIndex)):
-            return False
+        for ax in self.obj._data.axes:
+            if isinstance(ax, MultiIndex):
+                return False
 
         return True
 
     def _multi_take(self, tup):
-        index = self._convert_for_reindex(tup[0], axis=0)
-        columns = self._convert_for_reindex(tup[1], axis=1)
-        return self.obj.reindex(index=index, columns=columns)
+        from pandas.core.frame import DataFrame
+        from pandas.core.panel import Panel
+
+        if isinstance(self.obj, DataFrame):
+            index = self._convert_for_reindex(tup[0], axis=0)
+            columns = self._convert_for_reindex(tup[1], axis=1)
+            return self.obj.reindex(index=index, columns=columns)
+        elif isinstance(self.obj, Panel):
+            conv = [self._convert_for_reindex(x, axis=i)
+                    for i, x in enumerate(tup)]
+            return self.obj.reindex(items=tup[0], major=tup[1], minor=tup[2])
 
     def _convert_for_reindex(self, key, axis=0):
         labels = self.obj._get_axis(axis)
diff --git a/pandas/core/panel.py b/pandas/core/panel.py
@@ -682,6 +682,10 @@ def reindex(self, major=None, items=None, minor=None, method=None,
         major = _mut_exclusive(major, major_axis)
         minor = _mut_exclusive(minor, minor_axis)
 
+        if (method is None and not self._is_mixed_type and
+            com._count_not_none(items, major, minor) == 3):
+            return self._reindex_multi(items, major, minor)
+
         if major is not None:
             result = result._reindex_axis(major, method, 1, copy)
 
@@ -696,6 +700,32 @@ def reindex(self, major=None, items=None, minor=None, method=None,
 
         return result
 
+    def _reindex_multi(self, items, major, minor):
+        a0, a1, a2 = len(items), len(major), len(minor)
+
+        values = self.values
+        new_values = np.empty((a0, a1, a2), dtype=values.dtype)
+
+        new_items, indexer0 = self.items.reindex(items)
+        new_major, indexer1 = self.major_axis.reindex(major)
+        new_minor, indexer2 = self.minor_axis.reindex(minor)
+
+        if indexer0 is None:
+            indexer0 = range(len(new_items))
+
+        if indexer1 is None:
+            indexer1 = range(len(new_major))
+
+        if indexer2 is None:
+            indexer2 = range(len(new_minor))
+
+        for i, ind in enumerate(indexer0):
+            com.take_2d_multi(values[ind], indexer1, indexer2,
+                              out=new_values[i])
+
+        return Panel(new_values, items=new_items, major_axis=new_major,
+                     minor_axis=new_minor)
+
     def reindex_axis(self, labels, axis=0, method=None, level=None, copy=True):
         """Conform Panel to new index with optional filling logic, placing
         NA/NaN in locations having no value in the previous index. A new object
diff --git a/vb_suite/indexing.py b/vb_suite/indexing.py
@@ -103,7 +103,18 @@
 midx = MultiIndex.from_arrays([a, b])
 midx = midx.take(np.random.permutation(np.arange(100000)))
 """
-sort_level_zero = Benchmark("midx.sortlevel(0)", setup, 
+sort_level_zero = Benchmark("midx.sortlevel(0)", setup,
                             start_date=datetime(2012,1,1))
 sort_level_one = Benchmark("midx.sortlevel(1)", setup,
                            start_date=datetime(2012,1,1))
+
+#----------------------------------------------------------------------
+# Panel subset selection
+
+setup = common_setup + """
+p = Panel(np.random.randn(500, 500, 500))
+inds = range(0, 500, 10)
+"""
+
+indexing_panel_subset = Benchmark('p.ix[inds, inds, inds]', setup,
+                                  start_date=datetime(2012, 1, 1))
