Pandas:将掩码应用于多索引数据帧

Pandas: Apply mask to multiindex dataframe

我有一个包含 MultiIndex 列的 pandas 数据框,有 3 个级别:

import itertools
import numpy as np

def mklbl(prefix, n):
    return ["%s%s" % (prefix, i) for i in range(n)]


miindex = pd.MultiIndex.from_product([mklbl('A', 4)])

micolumns = pd.MultiIndex.from_tuples(list(itertools.product(['A', 'B'], ['a', 'b', 'c'], ['foo', 'bar'])),
                                      names=['lvl0', 'lvl1', 'lvl2'])

dfmi = pd.DataFrame(np.arange(len(miindex) * len(micolumns)).reshape((len(miindex), len(micolumns))),
                    index=miindex,
                    columns=micolumns).sort_index().sort_index(axis=1)

lvl0   A                       B                    
lvl1   a       b       c       a       b       c    
lvl2 bar foo bar foo bar foo bar foo bar foo bar foo
A0     1   0   3   2   5   4   7   6   9   8  11  10
A1    13  12  15  14  17  16  19  18  21  20  23  22
A2    25  24  27  26  29  28  31  30  33  32  35  34
A3    37  36  39  38  41  40  43  42  45  44  47  46

我想根据另一个数据框屏蔽这个数据框,它有最后两层索引:

cols = micolumns.droplevel(0).unique()
a_mask = pd.DataFrame(np.random.randn(len(dfmi.index), len(cols)), index=dfmi.index, columns=cols)
a_mask = (np.sign(a_mask) > 0).astype(bool)

        a             b             c       
      foo    bar    foo    bar    foo    bar
A0  False  False  False   True   True  False
A1   True  False   True  False   True   True
A2   True   True   True   True  False  False
A3   True  False  False   True   True  False

我想做的是根据a_mask屏蔽原始数据帧。 假设我想在 a_mask 为真时将原始条目设置为零。

我尝试使用 pd.IndexSlice,但它默默地失败了(即我可以 运行 以下代码,但没有效果:

dfmi.loc[:, pd.IndexSlice[:, a_mask]] = 0  #dfmi is unchanged

有什么实现方法的建议吗?

编辑 在我的用例中,标签是用笛卡尔积构造的,因此会有 (lev0、lev1、lev2) 的所有组合。 但情况是 lev0 可以取 2 个值 {A, B},而 lev1 可以取 3 个值 {a, b, c}

使用底层数组数据进行原位编辑以提高内存效率(不创建任何其他数据帧)-

d = len(dfmi.columns.levels[0])
n = dfmi.shape[1]//d
for i in range(0,d*n,n):
    dfmi.values[:,i:i+n][a_mask] = 0

样本运行-

In [833]: dfmi
Out[833]: 
lvl0   A                       B                    
lvl1   a       b       c       a       b       c    
lvl2 bar foo bar foo bar foo bar foo bar foo bar foo
A0     1   0   3   2   5   4   7   6   9   8  11  10
A1    13  12  15  14  17  16  19  18  21  20  23  22
A2    25  24  27  26  29  28  31  30  33  32  35  34
A3    37  36  39  38  41  40  43  42  45  44  47  46

In [834]: a_mask
Out[834]: 
        a             b             c       
      foo    bar    foo    bar    foo    bar
A0   True   True   True  False  False  False
A1  False   True  False  False   True  False
A2  False   True   True   True  False  False
A3  False  False  False  False  False   True

In [835]: d = len(dfmi.columns.levels[0])
     ...: n = dfmi.shape[1]//d
     ...: for i in range(0,d*n,n):
     ...:     dfmi.values[:,i:i+n][a_mask] = 0

In [836]: dfmi
Out[836]: 
lvl0   A                       B                    
lvl1   a       b       c       a       b       c    
lvl2 bar foo bar foo bar foo bar foo bar foo bar foo
A0     0   0   0   2   5   4   0   0   0   8  11  10
A1    13   0  15  14   0  16  19   0  21  20   0  22
A2    25   0   0   0  29  28  31   0   0   0  35  34
A3    37  36  39  38  41   0  43  42  45  44  47   0

更新后的解决方案更加稳健而不是对级别值进行硬编码:

lvl0_values = dfmi.columns.get_level_values(0).unique()
pd.concat([dfmi[i].mask(a_mask.rename_axis(['lvl1','lvl2'],axis=1),0) for i in lvl0_values],
          keys=lvl0_values, axis=1)

输出:

lvl0   A               B            
lvl1   a       b       a       b    
lvl2 bar foo bar foo bar foo bar foo
A0     1   0   0   0   5   0   0   0
A1     9   0  11   0  13   0  15   0
A2    17  16  19   0  21  20  23   0
A3     0  24   0  26   0  28   0  30

您可以这样做的一种方法:

pd.concat([dfmi['A'].mask(a_mask.rename_axis(['lvl1','lvl2'],axis=1),0),
           dfmi['B'].mask(a_mask.rename_axis(['lvl1','lvl2'],axis=1),0)],
           keys=['A','B'], axis=1)

print(a_mask)

lvl1      a             b       
lvl2    foo    bar    foo    bar
A0     True  False   True   True
A1     True  False   True  False
A2    False  False   True  False
A3    False   True  False   True

输出:

       A               B            
lvl1   a       b       a       b    
lvl2 bar foo bar foo bar foo bar foo
A0     1   0   0   0   5   0   0   0
A1     9   0  11   0  13   0  15   0
A2    17  16  19   0  21  20  23   0
A3     0  24   0  26   0  28   0  30

我会这样做:

mask = pd.concat({k: a_mask for k in dfmi.columns.levels[0]}, axis=1)
dfmi.where(~mask, 0)

我觉得用这种方式比较安全。

dfmi.where(a_mask.loc[:,dfmi.columns.droplevel(0)].values,0)
Out[191]: 
lvl0   A               B            
lvl1   a       b       a       b    
lvl2 bar foo bar foo bar foo bar foo
A0     0   0   0   2   0   0   0   6
A1     9   8  11   0  13  12  15   0
A2     0  16  19  18   0  20  23  22
A3    25   0   0   0  29   0   0   0