mean over all columns with the same prefix in python pandas

Viewed 476

I have a pandas df with multiple variables, that I want to compute the mean for. The columns are named with the variable name as prefix and the variable level as suffix. Actually, I got a lot of variables with 30 levels each, but with 2 levels and 2 variables it would look like this:

[index, oneVariable_s1, oneVariable_s2, otherVariable_s1, otherVariable_s2]

data could be generated this way:

data = [{'index':'p001', 'oneVariable_s1': 1, 'oneVariable_s2': 2, 'otherVariable_s1':3, 'otherVariable_s2':6},
        {'index':'p002','oneVariable_s1':10, 'oneVariable_s2': 20, 'otherVariable_s1': 30, 'otherVariable_s2':6}]

df = pd.DataFrame(data).set_index('index')
df

Is there a way to automatically compute the mean for all columns with the same variable name/ the same prefix before "_s1" to "_s30", so that I don't have to write down all variable names for computing the means? So far I've used this to compute the mean over columns:

df['oneVariable_mean'] = df.filter(like='oneVariable').mean(axis=1)
df['otherVariable_mean'] = df.filter(like='otherVariable').mean(axis=1)

This leads to the solution I'm looking for, but I need to type in all variable names.

2 Answers

As you confirmed that the column names are in format <var_name>_s<var_level>, you can use pandas str.split method first to split the column names on "_s" and get the variable names. Use set to pick only the unique variable names.

Then just calculate the mean for all unique variables as you were already doing them individually using filter:

unique_vars = set([var[0] for var in df.columns.str.split('_s')])
for var in unique_vars:
    df[var + '_mean'] = df.filter(like=var).mean(axis=1)

Output:

>>> df
       oneVariable_s1  oneVariable_s2  otherVariable_s1  otherVariable_s2  otherVariable_mean  oneVariable_mean
index                                                                                                          
p001                1               2                 3                 6                 4.5               1.5
p002               10              20                30                 6                18.0              15.0

Always keep data long and not wide for efficient storage, easier computing without loops, and enhanced maintainability and reproducibility. Most analytical steps (aggregation, merging, plotting, testing, modeling) require long format data. Wide format are usually for presentation or reporting purposes.

Specifically, run pandas.wide_to_long. Then, run mean aggregation as needed:

long_df = pd.wide_to_long(
    df.assign(id = df.index),
    stubnames = ["oneVariable", "anotherVariable"],
    i = "id",
    j = "s_num",
    sep = "_",
    suffix = r"\w+"
)

means = long_df.agg("mean")
means

grp_means = long_df.groupby(["group"]).agg("mean")
grp_means
Related