version bump: 0.1.3 -- fixed some bugs related to stacking and unstac…

…king DataFrames (and added new tests to check for those bugs)

HISTORY.rst

@@ -2,6 +2,11 @@
 History
 =======
 
+0.1.3 (2021-08-04)
+------------------
+
+* Fixed some bugs related to stacking and unstacking DataFrames
+
 0.1.2 (2021-08-04)
 ------------------
 

datawrangler/__init__.py

@@ -2,8 +2,7 @@
 
 __author__ = """Contextual Dynamics Lab"""
 __email__ = 'contextualdynamics@gmail.com'
-__version__ = '0.1.2'
-
 
 from .zoo import wrangle
 from .decorate.decorate import funnel, pandas_stack as stack, pandas_unstack as unstack
+from .core.configurator import __version__

datawrangler/decorate/decorate.py

@@ -311,7 +311,7 @@ def pandas_unstack(x):
 
     Parameters
     ----------
-    :param x: a single DataFrame or MultiIndex DataFrame
+    :param x: a single DataFrame, a MultiIndex DataFrame, or a list of DataFrames
 
     Returns
     -------
@@ -320,6 +320,8 @@ def pandas_unstack(x):
     if not is_multiindex_dataframe(x):
         if is_dataframe(x):
             return x
+        elif type(x) is list and all([is_dataframe(d) for d in x]):
+            return x
         else:
             raise Exception(f'Unsupported datatype: {type(x)}')
 
@@ -338,7 +340,17 @@ def pandas_unstack(x):
     assert names[0] == grouper, 'Unstacking error'
 
     x.index.rename(names, inplace=True)
-    return [d[1].set_index(d[1].index.get_level_values(1)) for d in list(x.groupby(grouper))]
+
+    groups = list(x.groupby(grouper))
+    n_levels = len(groups[0][1].index.levels)
+    if n_levels > 2:
+        g = groups[0][1]
+        index = pd.MultiIndex.from_arrays([g.index.get_level_values(len(g.index.levels) - n)
+                                           for n in range(1, len(g.index.levels))][::-1])
+        return [d[1].set_index(index) for d in groups]
+    else:
+        return [d[1].set_index(d[1].index.get_level_values(len(d[1].index.levels) - 1))
+                for d in list(x.groupby(grouper))]
 
 
 def apply_stacked(f):

setup.py

@@ -41,6 +41,6 @@
     test_suite='tests',
     tests_require=test_requirements,
     url='https://github.com/ContextLab/data-wrangler',
-    version='0.1.2',
+    version='0.1.3',
     zip_safe=False,
 )

tests/wrangler/test_decorate.py

@@ -118,6 +118,31 @@ def f(x):
     assert np.allclose(means.iloc[0], data1.mean(axis=0))
     assert np.allclose(means.iloc[1], data2.mean(axis=0))
 
+    xs = [np.cumsum(np.random.randn(100, 5), axis=0) for _ in range(10)]
+
+    @dw.decorate.apply_unstacked
+    def g(x):
+        return x
+
+    assert all(np.allclose(x, y) for x, y in zip(xs, g(xs)))
+    assert all(np.allclose(x, y) for x, y in zip(xs, dw.unstack(g(dw.stack(xs)))))
+
+
+def test_unstack(data):
+    xs = dw.wrangle([np.cumsum(np.random.randn(100, 5), axis=0) for _ in range(10)])
+    ys = dw.wrangle([np.cumsum(np.random.randn(100, 5), axis=0) for _ in range(10)])
+
+    stacked_xs = dw.stack(xs)
+    stacked_ys = dw.stack(ys)
+    stacked_xy = dw.stack([stacked_xs, stacked_ys])
+
+    assert np.allclose(dw.unstack(stacked_xs)[0], xs[0])
+    assert np.allclose(dw.unstack(stacked_xs)[0].index.values, xs[0].index.values)
+
+    assert np.allclose(dw.unstack(stacked_xy)[0], stacked_xs)
+    assert dw.zoo.is_multiindex_dataframe(dw.unstack(stacked_xy)[0])
+    assert np.allclose(dw.unstack(stacked_xy)[0].index.to_frame(), stacked_xs.index.to_frame())
+
 
 def test_apply_stacked(data):
     i = 4