Pandas групповая и взвешенная сумма для нескольких столбцов

Question

Я вижу дюжину Pandas групповых вопросов из нескольких столбцов, но я не знаю, как заставить это работать в разумные сроки. Моя цель состоит в том, чтобы сгруппировать несколько столбцов, и с полученным подмножеством применить np.dot к каждому оставшемуся столбцу относительно моих весов:

# Example data:
weights = np.array([.20, .60, .20])
data = pd.DataFrame([[0, "TX", 10, 55], [0, "TX", 5, 30], [0, "TX", 2, 75], [1, "TX", 4, 30], [1, "TX", 8, 100], [1, "TX", 2, 30]], columns=["sim", "state", "x1", "x2"])

print(data)
   sim state  x1   x2
0    0    TX  10   55
1    0    TX   5   30
2    0    TX   2   75
3    1    TX   4   30
4    1    TX   8  100
5    1    TX   2   30

Я не смог заставить np.dot работать из коробки, поэтому мне пришлось разбить умножение и суммирование на отдельные этапы. Вот что я попробовал, но в моем наборе данных из нескольких миллионов строк это занимает ~ 2 минуты, не говоря уже о том, что он довольно нечитабелен:

results = data.groupby(["sim", "state"]).apply(lambda sdf: (sdf[["x1", "x2"]] * weights.reshape((3,1))).sum())

print(results.reset_index())
   sim state   x1    x2
0    0    TX  5.4  44.0
1    1    TX  6.0  72.0

Answer 1

Как насчет ...

(df.set_index(['sim', 'state'])
   .mul(np.tile(weights, len(df) // len(weights)), axis=0)
   .sum(level=[0, 1]))

            x1    x2
sim state           
0   TX     5.4  44.0
1   TX     6.0  72.0

Как это работает,

• установить индекс на то, что не следует умножать (первичные ключи df, по сути)
• используйте mul для выполнения широковещательного умножения с весовыми коэффициентами
• на индексы и суммирования взвешенных значений.

Это работает в предположении, что len(df) % len(weights) == 0.