pandas-dev#15150 added conditional calculation of crosstable margins based on normalization type, corrected expected margin values in test_margin_dropna

Author: Christoph Möhl

pandas/core/reshape/pivot.py

@@ -199,22 +199,11 @@ def pivot_table(data, values=None, index=None, columns=None, aggfunc='mean',
 
 def _add_margins(table, data, values, rows, cols, aggfunc,
                  margins_name='All'):
-    if not isinstance(margins_name, compat.string_types):
-        raise ValueError('margins_name argument must be a string')
 
-    exception_msg = 'Conflicting name "{0}" in margins'.format(margins_name)
-    for level in table.index.names:
-        if margins_name in table.index.get_level_values(level):
-            raise ValueError(exception_msg)
+    _check_margins_name(margins_name, table)
 
     grand_margin = _compute_grand_margin(data, values, aggfunc, margins_name)
 
-    # could be passed a Series object with no 'columns'
-    if hasattr(table, 'columns'):
-        for level in table.columns.names[1:]:
-            if margins_name in table.columns.get_level_values(level):
-                raise ValueError(exception_msg)
-
     if len(rows) > 1:
         key = (margins_name,) + ('',) * (len(rows) - 1)
     else:
@@ -263,6 +252,21 @@ def _add_margins(table, data, values, rows, cols, aggfunc,
     return result
 
 
+def _check_margins_name(margins_name, table):
+    if not isinstance(margins_name, compat.string_types):
+        raise ValueError('margins_name argument must be a string')
+
+    exception_msg = 'Conflicting name "{0}" in margins'.format(margins_name)
+    for level in table.index.names:
+        if margins_name in table.index.get_level_values(level):
+            raise ValueError(exception_msg)
+    # could be passed a Series object with no 'columns'
+    if hasattr(table, 'columns'):
+        for level in table.columns.names[1:]:
+            if margins_name in table.columns.get_level_values(level):
+                raise ValueError(exception_msg)
+
+
 def _compute_grand_margin(data, values, aggfunc,
                           margins_name='All'):
 
@@ -507,17 +511,32 @@ def crosstab(index, columns, values=None, rownames=None, colnames=None,
         df = DataFrame(data)
         df['__dummy__'] = 0
         table = df.pivot_table('__dummy__', index=rownames, columns=colnames,
-                               aggfunc=len, margins=margins,
+                               aggfunc=len, margins=False,
                                margins_name=margins_name, dropna=dropna)
         table = table.fillna(0).astype(np.int64)
 
     else:
         data['__dummy__'] = values
         df = DataFrame(data)
         table = df.pivot_table('__dummy__', index=rownames, columns=colnames,
-                               aggfunc=aggfunc, margins=margins,
+                               aggfunc=aggfunc, margins=False,
                                margins_name=margins_name, dropna=dropna)
 
+    if margins:
+        _check_margins_name(margins_name, table)
+
+        if normalize != 'index':
+            # add margin column
+            table[margins_name] = table.sum(axis=1)
+
+        if normalize != 'columns':
+            # add margin row
+            if type(table.index) is MultiIndex:
+                table = table.transpose()
+                table[margins_name] = table.sum(axis=1)
+                table = table.transpose()
+            else:
+                table.loc[margins_name] = table.sum(axis=0)
     # Post-process
     if normalize is not False:
         table = _normalize(table, normalize=normalize, margins=margins,
@@ -534,49 +553,37 @@ def _normalize(table, normalize, margins, margins_name='All'):
         try:
             normalize = axis_subs[normalize]
         except KeyError:
-            raise ValueError("Not a valid normalize argument")
+            raise ValueError(
+                "Not a valid normalize argument: {!r}".format(normalize))
+
+    # Actual Normalizations
+    normalizers = {
+        'columns': lambda x: x / x.sum(),
+        'index': lambda x: x.div(x.sum(axis=1), axis=0)
+    }
 
     if margins is False:
-        # Actual Normalizations
-        normalizers = {
-            'all': lambda x: x / x.sum(axis=1).sum(axis=0),
-            'columns': lambda x: x / x.sum(),
-            'index': lambda x: x.div(x.sum(axis=1), axis=0)
-        }
+        normalizers['all'] = lambda x: x / x.sum(axis=1).sum(axis=0)
 
     elif margins is True:
-        # skip margin rows and/or cols for normalization
-        normalizers = {
-            'all': lambda x: x / x.iloc[:-1, :-1].sum(axis=1).sum(axis=0),
-            'columns': lambda x: x.div(x.iloc[:-1, :].sum()).iloc[:-1, :],
-            'index': lambda x: (x.div(x.iloc[:, :-1].sum(axis=1),
-                                axis=0)).iloc[:, :-1]
-        }
+        # skip margin rows and cols for normalization
+        normalizers['all'] = lambda x: x / x.iloc[:-1, :-1].sum(axis=1)\
+                                                           .sum(axis=0)
 
     else:
-        raise ValueError("Not a valid margins argument")
+        raise ValueError("Not a valid margins argument: {!r}".format(margins))
 
     normalizers[True] = normalizers['all']
 
     try:
         f = normalizers[normalize]
     except KeyError:
-        raise ValueError("Not a valid normalize argument")
+        raise ValueError(
+            "Not a valid normalize argument: {!r}".format(normalize))
 
     table = f(table)
     table = table.fillna(0)
 
-    if margins is True:
-        # reset index to ensure default index dtype
-        if normalize == 'index':
-            colnames = table.columns.names
-            table.columns = Index(table.columns.tolist())
-            table.columns.names = colnames
-        if normalize == 'columns':
-            rownames = table.index.names
-            table.index = Index(table.index.tolist())
-            table.index.names = rownames
-
     return table
 
 

pandas/tests/reshape/test_pivot.py

@@ -1184,15 +1184,15 @@ def test_margin_dropna(self):
         df = pd.DataFrame({'a': [1, 2, 2, 2, 2, np.nan],
                            'b': [3, 3, 4, 4, 4, 4]})
         actual = pd.crosstab(df.a, df.b, margins=True, dropna=False)
-        expected = pd.DataFrame([[1, 0, 1], [1, 3, 4], [2, 4, 6]])
+        expected = pd.DataFrame([[1, 0, 1], [1, 3, 4], [2, 3, 5]])
         expected.index = Index([1.0, 2.0, 'All'], name='a')
         expected.columns = Index([3, 4, 'All'], name='b')
         tm.assert_frame_equal(actual, expected)
 
         df = DataFrame({'a': [1, np.nan, np.nan, np.nan, 2, np.nan],
                         'b': [3, np.nan, 4, 4, 4, 4]})
         actual = pd.crosstab(df.a, df.b, margins=True, dropna=False)
-        expected = pd.DataFrame([[1, 0, 1], [0, 1, 1], [1, 4, 6]])
+        expected = pd.DataFrame([[1, 0, 1], [0, 1, 1], [1, 1, 2]])
         expected.index = Index([1.0, 2.0, 'All'], name='a')
         expected.columns = Index([3.0, 4.0, 'All'], name='b')
         tm.assert_frame_equal(actual, expected)
@@ -1209,8 +1209,8 @@ def test_margin_dropna(self):
         m = MultiIndex.from_arrays([['one', 'one', 'two', 'two', 'All'],
                                     ['dull', 'shiny', 'dull', 'shiny', '']],
                                    names=['b', 'c'])
-        expected = DataFrame([[1, 0, 1, 0, 2], [2, 0, 1, 1, 5],
-                              [3, 0, 2, 1, 7]], columns=m)
+        expected = DataFrame([[1, 0, 1, 0, 2], [2, 0, 1, 1, 4],
+                              [3, 0, 2, 1, 6]], columns=m)
         expected.index = Index(['bar', 'foo', 'All'], name='a')
         tm.assert_frame_equal(actual, expected)
 
@@ -1220,7 +1220,7 @@ def test_margin_dropna(self):
                                     ['one', 'two', 'one', 'two', '']],
                                    names=['a', 'b'])
         expected = DataFrame([[1, 0, 1], [1, 0, 1], [2, 0, 2], [1, 1, 2],
-                              [5, 2, 7]], index=m)
+                              [5, 1, 6]], index=m)
         expected.columns = Index(['dull', 'shiny', 'All'], name='c')
         tm.assert_frame_equal(actual, expected)
 
@@ -1421,22 +1421,23 @@ def test_crosstab_errors(self):
         df = pd.DataFrame({'a': [1, 2, 2, 2, 2], 'b': [3, 3, 4, 4, 4],
                            'c': [1, 1, np.nan, 1, 1]})
 
-        error = 'values cannot be used without an aggfunc.'
+        error = "values cannot be used without an aggfunc."
         with tm.assert_raises_regex(ValueError, error):
             pd.crosstab(df.a, df.b, values=df.c)
 
-        error = 'aggfunc cannot be used without values'
+        error = "aggfunc cannot be used without values"
         with tm.assert_raises_regex(ValueError, error):
             pd.crosstab(df.a, df.b, aggfunc=np.mean)
 
-        error = 'Not a valid normalize argument'
+        error = "Not a valid normalize argument: '42'"
         with tm.assert_raises_regex(ValueError, error):
             pd.crosstab(df.a, df.b, normalize='42')
 
+        error = "Not a valid normalize argument: 42"
         with tm.assert_raises_regex(ValueError, error):
             pd.crosstab(df.a, df.b, normalize=42)
 
-        error = 'Not a valid margins argument'
+        error = "Not a valid margins argument: 42"
         with tm.assert_raises_regex(ValueError, error):
             pd.crosstab(df.a, df.b, normalize='all', margins=42)
 
@@ -1495,6 +1496,7 @@ def test_crosstab_with_numpy_size(self):
         expected = pd.DataFrame(expected_data,
                                 index=expected_index,
                                 columns=expected_column)
+
         tm.assert_frame_equal(result, expected)