Фильтровать сгруппированные Pandas данные по совокупности столбцов, когда группы относятся к уровню MultiIndex - PullRequest
2 голосов
/ 11 февраля 2020

Как я могу отбросить данные из одного уровня многоуровневого индексированного фрейма данных на основе агрегированной информации, которую я получаю из столбца внутри группы на этом уровне?

Например, с фреймом данных dfmi:

midx = pd.MultiIndex.from_product([['A0','A1','A2'], ['B0','B1','B2']], names=["index_1", "index_2"])
columns = ['foo', 'bar']
dfmi = pd.DataFrame(np.arange(18).reshape((len(midx), len(columns))),
                    index=midx, columns=columns)

dfmi
                 foo  bar
index_1 index_2          
A0      B0         0    1
        B1         2    3
        B2         4    5
A1      B0         6    7
        B1         8    9
        B2        10   11
A2      B0        12   13
        B1        14   15
        B2        16   17

Допустим, я хочу сохранить уровни index_1, только если среднее значение для foo превышает определенный порог.
Например:

thresh = 5

for grp, data in dfmi.groupby("index_1"):
    print(data.foo.mean() > thresh)

False  <-- drop this level
True
True

Желаемый результат:

                 foo  bar
index_1 index_2          
A1      B0         6    7
        B1         8    9
        B2        10   11
A2      B0        12   13
        B1        14   15
        B2        16   17

В этом примере с игрушкой я могу получить то, что хочу, с помощью dfmi.loc[pd.IndexSlice["A1":"A2", :]]. Но я не могу понять, как использовать варианты IndexSlice или loc для агрегирования внутри сгруппированного MultiIndex, а затем нарезать весь кадр данных на основе результатов.

Мое лучшее решение на данный момент - просто отслеживать значения уровней, которые квалифицируются как хранители (с grp), а затем использовать накопленную коллекцию keepers с IndexSlice:

keepers = list()

for grp, data in dfmi.groupby("index_1"):
    if data.foo.mean() > thresh:
        keepers.append(grp)

dfmi.loc[pd.IndexSlice[keepers, :]]

Я ищу более эффективный и / или более элегантный способ достижения sh, чем с нативной Pandas функциональностью.

1 Ответ

2 голосов
/ 11 февраля 2020

Вы можете использовать loc после создания маски, например:

mask = dfmi.groupby(level=0)['foo'].mean()>thresh

dfmi.loc[mask.index[mask]]

Выход:

index_1 index_2          
A1      B0         6    7
        B1         8    9
        B2        10   11
A2      B0        12   13
        B1        14   15
        B2        16   17
...