Python: inconsistent handling of IF statement in loop

Viewed 315

I have a dataframe df containing conditions and values.

import pandas as pd
df=pd.DataFrame({'COND':['X','X','X','Y','Y','Y'], 'VALUE':[1,2,3,1,2,3]})

Therefore df looks like:

  COND  VALUE
     X      1
     X      2
     X      3
     Y      1
     Y      2
     Y      3

I'm using a loop to subset df according to COND, and write separate text files containing values for each condition

conditions = {'X','Y'}
for condition in conditions:
    df2 = df[df['COND'].isin([condition])][['VALUE']]
    df2.to_csv(condition + '_values.txt', header=False, index=False)

The end results is two text files: X_vals.txt and Y_vals.txt, both of which contain 1 2 3. Up until this point everything is working as expected.

I would like to further subset df for one condition only. For example, perhaps I want all values from condition Y, but ONLY values < 3 from condition X. In this scenario, X_vals.txt should contain 1 2 and Y_vals.txt should contain 1 2 3. I tried implementing this with an IF statement:

conditions = {'X','Y'}
for condition in conditions:
    if condition == 'X':
        df = df[df['VALUE'] < 3]

    df2 = df[df['COND'].isin([condition])][['VALUE']]
    df2.to_csv(condition + '_values.txt', header=False, index=False)

Here is where the inconsistency occurs. The above code works fine (i.e. X_vals.txt contains 1 2, and Y_vals.txt 1 2 3, as intended), but when I use if condition=='Y' instead of if condition=='X', it breaks, and both text files only contain 1 2.

In other words, if I specify the first element of conditions in the IF statement then it works as intended, however if I specify the second element then it breaks and applies the < 3 subset to values from both conditions.

What is going on here and how can I resolve it?

Thanks!

3 Answers

The problem you are encountering arises because you are overwriting df inside the loop.

conditions = {'X','Y'}
for condition in conditions:
    if condition == 'X':
        df = df[df['VALUE'] < 3]  # <-- HERE'S YOUR ISSUE

    df2 = df[df['COND'].isin([condition])][['VALUE']]
    df2.to_csv(condition + '_values.txt', header=False, index=False)

What slightly surprised me is that when you are looping over the set conditions you get condition = 'Y' first, then condition = 'X'. But as a set is an unordered collection (i.e. it doesn't claim to have an inherent order of its elements), this ought not to be too disturbing: python is just reading out the elements in the most internally convenient way.

You could use conditions = ['X', 'Y'] to loop over a list (an ordered collection) instead. Then it will do X first, then Y. However, if you do that you will get the same bug but in reverse (i.e. it works for if condition == 'Y' but not if condition == 'X').

This is because after the loop runs once, df has been reassigned to the subset of the original df that only contains values less than three. That's why you get only the values 1 and 2 in both files if the if condition statement triggers on the first pass through the loop.

Now for the fix:


conditions = ['X', 'Y']

for condition in conditions:

    csv_name = f"{condition}_values.txt"

    if condition == 'X':
        df_filter = f"VALUE < 3 & COND == '{condition}'"
    else:
        df_filter = f"COND == '{condition}'"

    df.query(df_filter).VALUE.to_csv(csv_name, header=False, index=False)

Here I've introduced the DataFrame.query method, which is typically more concise than trying to create a Boolean series to use as a mask as you were doing.

The f-string syntax only works on python 3.6+, if you're on a lower version then modify as appropriate (e.g. df_filter = "COND == '{}'".format(condition))

We can write the condition to dict then use map filter the df before groupby

cond = {'X' : 2, 'Y' : 3}
subdf = df[df['VALUE']<df.COND.map(cond)]
for x, y in subdf.groupby('COND'):
    y.to_csv(x + '_values.txt')
df=pd.DataFrame({'COND':['X','X','X','Y','Y','Y'], 'VALUE':[1,2,3,1,2,3]})

conditions = df.COND

for condition in conditions:
 print(condition)
 df2=df[df['COND'].isin([condition])][['VALUE']]
 df2.to_csv(condition + '_values.txt',header=False, index=False)


for condition in conditions:
 if condition=='X':
     df=df[df['VALUE'] < 3]

 df2=df[df['COND'].isin([condition])][['VALUE']]
 df2.to_csv(condition + '_values.txt',header=False, index=False)

You didn't specify the variable "Condition", so it gave you an error. try doing :

conditions = df.COND

before the for loop

Related