[Pandas] MultiIndexな列を最上位のindex毎に処理する方法で悩んだ

先週、groupby.aggに複数の集約関数を与えてできた、列がMultiIndexのDataFrameを、元の列毎に処理する方法がわからず悩んだが、良い方法が見つかったのでメモしておく。

In [1]: 
import numpy as np
import pandas as pd

def get_data(date):
    n = np.random.randint(5, 10)
    return pd.DataFrame({
        'person': list("ABCDEFGHIJ"[:n]), 
        'date': date,
        'hour': map(lambda x: ['AM', 'PM'][x], sorted(np.random.randint(2, size=n))),
        'Left': np.random.randint(5, size=n),
        'Right': np.random.randint(10, size=n)
    }).set_index(['person', 'date', 'hour'])

# test
np.random.seed(3)
get_data('2020-03-01')
Out[1]:
Left Right
person date hour
A 2020-03-01 AM 3 4
B 2020-03-01 AM 2 7
C 2020-03-01 AM 3 8
D 2020-03-01 AM 1 1
E 2020-03-01 PM 1 6
F 2020-03-01 PM 2 2
G 2020-03-01 PM 0 2

このような形式で、来た人が午前か午後のどちらに左の通路と右の通路をそれぞれ何回通ったかというデータがあり、時間帯毎の１人当たりのそれぞれの通路を通った平均回数を計算したいとする。

実際のデータはサイズが大きく、１日分のデータはDRAMに載るが全期間のデータは載らずMemory Errorになったので、平均は合計÷人数ということで、次のように１日分ずつデータを読み込んで時間帯毎、通路毎に通った回数の合計と人数を加算していくようにした。

In [2]: 
np.random.seed(3)

sumdf = None
for date in pd.date_range('2020-03-01', '2020-03-07'):
    df = get_data(date.date())

    df = df.groupby('hour').agg(['sum', 'size'])
    if sumdf is None:
        sumdf = df
    else:
        sumdf = sumdf.add(df, fill_value=0)

sumdf
Out [2]:
Left Right
sum size sum size
hour
AM 40 27 114 27
PM 49 23 72 23

groupby.aggに'sum', 'size'という複数の集約関数を渡しているので、結果の列がMultiIndexになっている。
後はこれの 'Left' と 'Right' をそれぞれの 'sum' / 'size' に置換すれば良いのだが、その方法がわからなかった。
結局、調べながら試行錯誤して筆者が最もシンプルだと思ったコードは次のようになった。

In [3]: 
meandf = sumdf.groupby(level=0, axis=1).apply(lambda x: x[x.name]['sum'] / x[x.name]['size'])
meandf
Out [3]:
Left Right
hour
AM 1.481481 4.222222
PM 2.130435 3.130435
In [4]: 
meandf.plot(kind='bar', title='The average per person')
Out [4]:

この例だと 'Left', 'Right' の両方の列の 'size' が同じなので、1日分ずつ集計する時の結果を'Left'の合計, 'Right'の合計, 'size'の3列にすれば、扱いがややこしいMultiIndexにする必要も無いので、それもやってみたが、コードが余計に煩雑になってしまった。

np.random.seed(3)

sumdf = None
for date in pd.date_range('2020-03-01', '2020-03-07'):
    df = get_data(date.date())

    grouped = df.groupby('hour')
    sizesr = grouped.agg('size').rename('size')
    df = grouped.agg('sum').join(sizesr)
    
    if sumdf is None:
        sumdf = df
    else:
        sumdf = sumdf.add(df, fill_value=0)

sizesr = sumdf['size']
meandf = sumdf.drop(columns='size').div(sizesr, axis=0)

今回の例含め、MultiIndexにする方が処理がスマートになるケースが結構ありそうに思うので、やはり嫌がらずにMultiIndexの扱いに慣れるようにしようと思った。