Pandas: imputing descriptive stats using a groupby with a variable-CodePudding

I have a data frame like this:

input_df = pd.DataFrame({"sex": ["M", "F", "F", "M", "M"], "Class": [1, 2, 2, 1, 1], "Age":[40, 30, 30, 50, NaN]})

What I want to do is to impute the missing value for the age based on the sex and class columns. I have tried doing it with a function, conditional_impute. What the function does is take a data frame and a condition and then use it to impute the age based on the sex and class grouping. Butthe caveat is that the condition can either be a mean or median and if not either of these two, the function has to raise an error. So I did this:

### START FUNCTION
def conditional_impute(input_df, choice='median'):
    my_df = input_df.copy()
    # if choice is not median or mean, raise valueerror
    if choice == "mean" or choice == "median":
        my_df['Age'] = my_df['Age'].fillna(my_df.groupby(["Sex","Pclass"])['Age'].transform(choice))
    else:
        raise ValueError()    
    # round the values in Age colum
    my_df['Age'] = round(my_df['Age'], 1)
    return my_df
### END FUNCTION

But I am getting an error when I call it.

conditional_impute(train_df, choice='mean')

What could I possibly be doing wrong? I really cannot get a handle on this.

CodePudding user response：

If you give the right inputs, it outputs just fine...

# Fixed input to match function:
df = pd.DataFrame({"Sex": ["M", "F", "F", "M", "M"], "Pclass": [1, 2, 2, 1, 1], "Age":[40, 30, 30, 50, np.nan]})

def conditional_impute(input_df, choice='median'):
    my_df = input_df.copy()
    # if choice is not median or mean, raise valueerror
    if choice == "mean" or choice == "median":
        my_df['Age'] = my_df['Age'].fillna(my_df.groupby(["Sex","Pclass"])['Age'].transform(choice))
    else:
        raise ValueError()    
    # round the values in Age colum
    my_df['Age'] = round(my_df['Age'], 1)
    return my_df

conditional_impute(df, choice='mean')

Output:

  Sex  Pclass   Age
0   M       1  40.0
1   F       2  30.0
2   F       2  30.0
3   M       1  50.0
4   M       1  45.0