Search code examples
pythonpandasdataframeindexingmulti-index

How to set values based on a condition on a subset of MultiIndex pandas dataframe


I want to take a subset of a MultiIndex pandas dataframe, test for values less than zero and set them to zero.

For example:

df = pd.DataFrame({('A','a'): [-1,-1,0,10,12],
                   ('A','b'): [0,1,2,3,-1],
                   ('B','a'): [-20,-10,0,10,20],
                   ('B','b'): [-200,-100,0,100,200]})

df[df['A']<0] = 0.0

gives

    A        B
    a   b    a     b
0  -1   0  -20  -200
1  -1   1  -10  -100
2   0   2    0     0
3  10   3   10   100
4  12  -1   20   200

Which shows that it was not able to set based on the condition. Alternatively if I do chained assignment:

df.loc[:,'A'][df['A']<0] = 0.0

this gives the same result (and setting with copy warning).

I could loop through each column based on the condition that the first level is the one that I want:

for one,two in df.columns.values:
    if one == 'A':
        df.loc[df[(one,two)]<0, (one,two)] = 0.0

which gives the desired result:

    A       B
    a  b    a     b
0   0  0  -20  -200
1   0  1  -10  -100
2   0  2    0     0
3  10  3   10   100
4  12  0   20   200

What is the best way to do this in pandas?


Solution

  • This is an application of (and one of the main motivations for using MultiIndex slicers), see docs here

    In [20]: df = pd.DataFrame({('A','a'): [-1,-1,0,10,12],
                       ('A','b'): [0,1,2,3,-1],
                       ('B','a'): [-20,-10,0,10,20],
                       ('B','b'): [-200,-100,0,100,200]})
    
    In [21]: df
    Out[21]: 
        A      B     
        a  b   a    b
    0  -1  0 -20 -200
    1  -1  1 -10 -100
    2   0  2   0    0
    3  10  3  10  100
    4  12 -1  20  200
    
    In [22]: idx = pd.IndexSlice
    
    In [23]: mask = df.loc[:,idx['A',:]]<0
    
    In [24]: mask
    Out[24]: 
           A       
           a      b
    0   True  False
    1   True  False
    2  False  False
    3  False  False
    4  False   True
    
    In [25]: df[mask] = 0
    
    In [26]: df
    Out[26]: 
        A      B     
        a  b   a    b
    0   0  0 -20 -200
    1   0  1 -10 -100
    2   0  2   0    0
    3  10  3  10  100
    4  12  0  20  200
    

    Since you are working with the 1st level of the columns index, the following will work as well. The above example is more general, say you wanted to do this for 'a'.

    In [30]: df[df[['A']]<0] = 0
    
    In [31]: df
    Out[31]: 
        A      B     
        a  b   a    b
    0   0  0 -20 -200
    1   0  1 -10 -100
    2   0  2   0    0
    3  10  3  10  100
    4  12  0  20  200