Pandas rename multilevel looking column names

Question

I have a dataframe which its column names have changed to

MultiIndex([(     'ID',       ''),
            ('Probability',   'mean'),
            ('Probability', 'median'),
            ('Uncertainty',   'mean'),
            ('Uncertainty', 'median')],
          )

because I did

data[data["ID"].notnull()].groupby(["ID"]).agg({"Probability":["mean", "median"], "Uncertainty":["mean", "median"]}).reset_index()

I want to rename the column names to something like this:

["ID", "Probability_mean", "Probability_median", "Uncertainty_mean", "Uncertainty_median"]

I can rename each original column name individually but I can't rename them together. I also tried flattening the dataframe as I thought it was a multi index or a multilevel. Although it looks so, it is not one. Most of the multi index functionalities don't work for it. Is there a way to rename these columns? Am I missing something?

Answer 1

If using pd.__version__ > 0.25.0 you can create a dictionary with NamedAggs.

Before reset_index you can use str.join to collapse the MultiIndex after the groupby. Then reset_index. This avoids issues with 'ID'. Also, there is no need to drop the null group keys as groupby will ignore these by default.

Sample Data

import pandas as pd
import numpy as np
N = 6
df = pd.DataFrame({'ID': np.arange(N)//2, 
                   'Probability': np.random.normal(0,1,N),
                   'Uncertainty': np.random.normal(0,1,N)})
agg_d = {'Probability': ['mean', 'median'], 'Uncertainty': ['mean', 'median']}

Code:

>= 0.25.0

d = {f'{k}_{x}': pd.NamedAgg(column=k, aggfunc=x) for k,v in agg_d.items() for x in v}
df.groupby('ID').agg(**d).reset_index()

< 0.25.0

res = df.groupby('ID').agg(agg_d)
res.columns = ['_'.join(tup) for tup in res.columns]
res = res.reset_index()

Output:

   ID  Probability_mean  Probability_median  Uncertainty_mean  Uncertainty_median
0   0          0.795119            0.795119          0.466417            0.466417
1   1          0.150184            0.150184         -0.132942           -0.132942
2   2          1.250202            1.250202         -0.102760           -0.102760