Добавить промежуточные столбцы в pandas с несколькими индексами

У меня есть фреймворк с 3-уровневым глубоким мультииндексиром на столбцах. Я хотел бы вычислить промежуточные итоги по строкам (sum(axis=1)), где я суммируюсь на одном из уровней, сохраняя остальные. Я думаю, что знаю, как это сделать, используя аргумент ключевого слова level pd.DataFrame.sum. Тем не менее, у меня возникают проблемы с тем, как включить результат этой суммы обратно в исходную таблицу.

Настройка:

import numpy as np
import pandas as pd
from itertools import product

np.random.seed(0)

colors = ['red', 'green']
shapes = ['square', 'circle']
obsnum = range(5)

rows = list(product(colors, shapes, obsnum))
idx = pd.MultiIndex.from_tuples(rows)
idx.names = ['color', 'shape', 'obsnum']

df = pd.DataFrame({'attr1': np.random.randn(len(rows)), 
                   'attr2': 100 * np.random.randn(len(rows))},
                  index=idx)

df.columns.names = ['attribute']

df = df.unstack(['color', 'shape'])

Дает хороший кадр:

Original frame

Скажем, я хотел уменьшить уровень shape. Я мог бы запустить:

tots = df.sum(axis=1, level=['attribute', 'color'])

чтобы получить мои итоговые значения:

totals

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

tots = df.sum(axis=1, level=['attribute', 'color'])
newcols = pd.MultiIndex.from_tuples(list((i[0], i[1], 'sum(shape)') for i in tots.columns))
tots.columns = newcols
bigframe = pd.concat([df, tots], axis=1).sort_index(axis=1)

aggregated

Есть ли более естественный способ сделать это?

Ответы

Ответ 1

Вот путь без циклов:

s = df.sum(axis=1, level=[0,1]).T
s["shape"] = "sum(shape)"
s.set_index("shape", append=True, inplace=True)
df.combine_first(s.T)

Хитрость заключается в использовании транспонированной суммы. Поэтому мы можем вставить еще один столбец (например, строку) с именем дополнительного уровня, который мы называем точно так же, как тот, который мы суммировали. Этот столбец можно преобразовать в уровень в индексе с помощью set_index. Затем мы объединяем df с транспонированной суммой. Если суммарный уровень не является последним, вам может потребоваться некоторое переупорядочение уровня.

Ответ 2

Вот мой грубоватый способ сделать это.

После выполнения вашего хорошо написанного (спасибо) образца кода я сделал следующее:

attributes = pd.unique(df.columns.get_level_values('attribute'))
colors = pd.unique(df.columns.get_level_values('color'))

for attr in attributes:
    for clr in colors:
        df[(attr, clr, 'sum')] = df.xs([attr, clr], level=['attribute', 'color'], axis=1).sum(axis=1)

df

Что дает мне:

big table