I have a DataFrame with a multi-index as a column. I want to group by level 1 and apply a function that generates a new column. I want this calculated column to be added to each group, and thus my dataframe to have new columns for each group.
I made a little dummy script and function to replicate what I want to do.
import pandas as pd
import numpy as np
columns = [('A','julian'),('A','geoffrey'),
('B','julian'),('B','geoffrey'),
('C','julian'),('C','geoffrey')]
columns = pd.MultiIndex.from_tuples(columns)
dataframe = pd.DataFrame(data=np.random.rand(10,6),columns=columns)
def addColumn(inputDF):
group = inputDF.columns[0][1]
inputDF['sum', group] = inputDF.sum(axis=1)
return inputDF
newColumnsDataframe = dataframe.groupby(level=1, axis=1).apply(addColumn)
The original dataframe looks as follows:
A B C
julian geoffrey julian geoffrey julian geoffrey
0 0.204082 0.073676 0.795725 0.279702 0.258185 0.258112
1 0.263235 0.096733 0.507324 0.541198 0.525919 0.757652
2 0.196243 0.028613 0.653408 0.364365 0.174911 0.924733
3 0.528785 0.831569 0.654160 0.738029 0.940831 0.294473
4 0.853517 0.263250 0.803087 0.855270 0.701937 0.264698
5 0.239797 0.069519 0.943544 0.374411 0.189361 0.846647
6 0.980734 0.290414 0.850097 0.873785 0.903645 0.118713
7 0.591942 0.088387 0.566298 0.062140 0.568482 0.872064
8 0.818167 0.061483 0.282050 0.008404 0.449198 0.658370
9 0.217424 0.427602 0.471933 0.171458 0.390549 0.234426
The resulting dataframe should look as follows (I constructed the sum DataFrame separately and concatenated the two dataframes to achieve this results):
A B C sum A B C \
geoffrey geoffrey geoffrey geoffrey julian julian julian
0 0.073676 0.279702 0.258112 0.611491 0.204082 0.795725 0.258185
1 0.096733 0.541198 0.757652 1.395584 0.263235 0.507324 0.525919
2 0.028613 0.364365 0.924733 1.317710 0.196243 0.653408 0.174911
3 0.831569 0.738029 0.294473 1.864071 0.528785 0.654160 0.940831
4 0.263250 0.855270 0.264698 1.383219 0.853517 0.803087 0.701937
5 0.069519 0.374411 0.846647 1.290578 0.239797 0.943544 0.189361
6 0.290414 0.873785 0.118713 1.282912 0.980734 0.850097 0.903645
7 0.088387 0.062140 0.872064 1.022590 0.591942 0.566298 0.568482
8 0.061483 0.008404 0.658370 0.728257 0.818167 0.282050 0.449198
9 0.427602 0.171458 0.234426 0.833486 0.217424 0.471933 0.390549
sum
julian
0 1.257992
1 1.296478
2 1.024561
3 2.123776
4 2.358542
5 1.372703
6 2.734476
7 1.726721
8 1.549415
9 1.079906
The approach in the script above is based on what makes sense to me and what others have written online about doing this sort of thing. However, newColumnsDataframe still only has 6 columns, not 8 (one added for each name).
I did notice that when I group by level=0 (so by A, B, or C) and use transform(but NOT when I use apply on this level), newColumnsDataframe does have 9 columns, one sum column added for each group. See the code below:
import pandas as pd
import numpy as np
columns = [('A','julian'),('A','geoffrey'),
('B','julian'),('B','geoffrey'),
('C','julian'),('C','geoffrey')]
columns = pd.MultiIndex.from_tuples(columns)
dataframe = pd.DataFrame(data=np.random.rand(10,6),columns=columns)
def addColumn(inputDF):
group = inputDF.columns[0][1]
inputDF[group, 'sum'] = inputDF.sum(axis=1)
return inputDF
newColumnsDataframe = dataframe.groupby(level=0, axis=1).transform(addColumn)
It was always my understanding that transform worked on each column within the group, whereas apply operated on the group as a whole dataframe. This seems to contradict that. I also noticed, that when I group by level=1 and use transform in stead of apply, it throws the following error:
ValueError: Length mismatch: Expected axis has 10 elements, new values have 6 elements
I am very confused about what is going on. Does anyone know why this DOES work when I use transform and group on level=0. Why there is an error when I do the same but group on level=1. And why grouping on EITHER level and APPLYING the function does not add columns to my final dataframe? Thanks in advance!
(PS: This is not the actual DataFrame or function I am using to add a column, just an easier illustration)