Group-by with multiple variables

Question

I have a df that looks a bit like this:

Date   Animal   Height   Weight
Jan-00   Cat      102      56
Jan-00   Cat      98       75
Jan-00   Cat      50       100
Jan-00   Cat      46       46
Jan-00   Cat      100      50

I am trying to find the average weight of cats in Jan-00 that are more than the median height of Cats in Jan-00. So in this example, the median height is 98; and the average weight of cats more than the median height is 53 (avg of 50, 56). I have a variable number of different animal types so I don't want to have to manually specify the animal types. In time there will be other things I want to measure over and above the weight of the cats (so I'm trying to future proof my code somewhat).

From scouring StackOverflow, my approach has been: (1) Write a function which tells the rest of the code which things I want to measure:

def column_index(df, query_cols):
    cols = df.columns.values
    sidx = np.argsort(cols)
    return sidx[np.searchsorted(cols,query_cols,sorter=sidx)]

cols = (column_index(df, ["Weight"]))

Then (2), write some code that splits my original df between dates and animal type, split by the median height, and then return the things I want to measure:

x = (df["Height"]
     .gt(df.groupby(["Date","Animal"])["Height"]
     .transform('median')))

df_Tall = df[x].mean(level = 0)[df.columns[cols]]

However, when I do this, it is only returning a single series of data for the weight; I'm trying to get multiple columns, each of which represent a different animal - ie my intended output should look like:

       Weight
Date   Cat   Animal_x   Animal_y   Animal_z
Jan-00  53     xx          xx         xx

I think my error likes in the line

.gt(df.groupby(["Date","Animal"])["Height"]

But I can't figure out how to fix it. Any thoughts would be really appreciated!

Thanks

user_6345518 · Answer

Your approach was almost correct. I just added another groupby of the median-filtered height data to group the mean weight (and height) by the Animal and Date columns:

df.loc[
    df['Height'].gt(df.groupby(['Date', 'Animal'])['Height'].transform('median')), :
].groupby(['Date', 'Animal']).mean().unstack()

Of course you can get the mean of the Weight by selecting the Weight column after (or also before) the second groupby, f.i. by changing the last line to: ].groupby(['Date', 'Animal'])[['Weight']].mean().unstack(), with the double brackets around 'Weight' retaining the df dimensions/structure.

step-by-step:

groupby Date and Animal, get median and select where Height is greater (same as in your question)
.loc explicit first axis indexing over your df[x] is just my personal preference. both work well to filter the data by the median height.
second groupby on height-filtered data, selecting Weight before or after groupby is optional.
get the mean of each animal and date group
unstack transforms the dataframe into multiindex-columns with Animal as the second level

The second groupby could afaik only be avoided, if groupby.median() supported a level-argument. But since level is not supported by (most or all?) groupby methods, a second groupby is required.

Group-by with multiple variables

1 Answers