python - Pandas : Proper way to set values based on condition for subset of multiindex dataframe

Question

I'm not sure of how to do this without chained assignments (which probably wouldn't work anyways because I'd be setting a copy).

I wan't 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

In [37]:

df

Out[37]:
    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 did a 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:

In [64]:

df

Out[64]:
    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

But somehow I feel there is a better way to do this than looping through the columns. What is the best way to do this in pandas?

Answer 1

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