使用 groupby+apply 对每个组进行聚类 - 性能问题答案

【问题标题】：Clustering each group with groupby+apply - performance issue使用 groupby+apply 对每个组进行聚类 - 性能问题
【发布时间】：2020-04-29 14:20:01
【问题描述】：

我有一个如下的数据框：

import pandas as pd
import numpy as np

df = pd.DataFrame.from_dict(
    {'id': {(1, 0, 'obj11'): '3',
  (1, 0, 'obj12'): '9',
  (1, 0, 'obj13'): '5',
  (1, 0, 'obj14'): '4',
  (1, 0, 'obj15'): '23',
  (1, 0, 'obj16'): '52',
  (1, 0, 'obj17'): '22',
  (1, 0, 'obj18'): '13',
  (1, 0, 'obj19'): '8',
  (1, 0, 'obj20'): '26',
  (1, 1000, 'obj11'): '3',
  (1, 1000, 'obj12'): '9',
  (1, 1000, 'obj13'): '5',
  (1, 1000, 'obj14'): '4',
  (1, 1000, 'obj15'): '23',
  (1, 1000, 'obj16'): '52',
  (1, 1000, 'obj17'): '22',
  (1, 1000, 'obj18'): '13',
  (1, 1000, 'obj19'): '8',
  (1, 1000, 'obj20'): '26',
  (1, 2000, 'obj11'): '3',
  (1, 2000, 'obj12'): '9',
  (1, 2000, 'obj13'): '5',
  (1, 2000, 'obj14'): '4',
  (1, 2000, 'obj15'): '23',
  (1, 2000, 'obj16'): '52',
  (1, 2000, 'obj17'): '22',
  (1, 2000, 'obj18'): '13',
  (1, 2000, 'obj19'): '8',
  (1, 2000, 'obj20'): '26',
  (1, 3000, 'obj11'): '3',
  (1, 3000, 'obj12'): '9',
  (1, 3000, 'obj13'): '5',
  (1, 3000, 'obj14'): '4',
  (1, 3000, 'obj15'): '23',
  (1, 3000, 'obj16'): '52',
  (1, 3000, 'obj17'): '22',
  (1, 3000, 'obj18'): '13',
  (1, 3000, 'obj19'): '8',
  (1, 3000, 'obj20'): '26',
  (1, 4000, 'obj11'): '3',
  (1, 4000, 'obj12'): '9',
  (1, 4000, 'obj13'): '5',
  (1, 4000, 'obj14'): '4',
  (1, 4000, 'obj15'): '23',
  (1, 4000, 'obj16'): '52',
  (1, 4000, 'obj17'): '22',
  (1, 4000, 'obj18'): '13',
  (1, 4000, 'obj19'): '8',
  (1, 4000, 'obj20'): '26'},
 'var': {(1, 0, 'obj11'): 61.05099868774414,
  (1, 0, 'obj12'): 52.6510009765625,
  (1, 0, 'obj13'): 61.422000885009766,
  (1, 0, 'obj14'): 75.99199676513672,
  (1, 0, 'obj15'): 72.30999755859375,
  (1, 0, 'obj16'): 63.79999923706055,
  (1, 0, 'obj17'): 52.604000091552734,
  (1, 0, 'obj18'): 61.02899932861328,
  (1, 0, 'obj19'): 65.16999816894531,
  (1, 0, 'obj20'): 71.26699829101562,
  (1, 1000, 'obj11'): 59.92499923706055,
  (1, 1000, 'obj12'): 49.4630012512207,
  (1, 1000, 'obj13'): 60.25299835205078,
  (1, 1000, 'obj14'): 77.15299987792969,
  (1, 1000, 'obj15'): 73.43199920654297,
  (1, 1000, 'obj16'): 62.207000732421875,
  (1, 1000, 'obj17'): 49.805999755859375,
  (1, 1000, 'obj18'): 60.459999084472656,
  (1, 1000, 'obj19'): 65.0199966430664,
  (1, 1000, 'obj20'): 71.9520034790039,
  (1, 2000, 'obj11'): 58.72600173950195,
  (1, 2000, 'obj12'): 45.98500061035156,
  (1, 2000, 'obj13'): 58.21099853515625,
  (1, 2000, 'obj14'): 78.35800170898438,
  (1, 2000, 'obj15'): 75.06199645996094,
  (1, 2000, 'obj16'): 59.23500061035156,
  (1, 2000, 'obj17'): 46.32699966430664,
  (1, 2000, 'obj18'): 57.902000427246094,
  (1, 2000, 'obj19'): 65.1510009765625,
  (1, 2000, 'obj20'): 72.99099731445312,
  (1, 3000, 'obj11'): 57.47800064086914,
  (1, 3000, 'obj12'): 42.904998779296875,
  (1, 3000, 'obj13'): 55.89699935913086,
  (1, 3000, 'obj14'): 79.41999816894531,
  (1, 3000, 'obj15'): 76.78800201416016,
  (1, 3000, 'obj16'): 55.53099822998047,
  (1, 3000, 'obj17'): 42.67900085449219,
  (1, 3000, 'obj18'): 55.277000427246094,
  (1, 3000, 'obj19'): 65.21199798583984,
  (1, 3000, 'obj20'): 74.27400207519531,
  (1, 4000, 'obj11'): 56.189998626708984,
  (1, 4000, 'obj12'): 41.14099884033203,
  (1, 4000, 'obj13'): 54.09000015258789,
  (1, 4000, 'obj14'): 80.78099822998047,
  (1, 4000, 'obj15'): 78.38999938964844,
  (1, 4000, 'obj16'): 57.492000579833984,
  (1, 4000, 'obj17'): 40.400001525878906,
  (1, 4000, 'obj18'): 53.159000396728516,
  (1, 4000, 'obj19'): 63.72200012207031,
  (1, 4000, 'obj20'): 75.40399932861328}}
)
df.index.names = ['k', 'group', 'object']

我想对组内的对象进行聚类，所以我定义了一个对组进行聚类的函数。

def cluster_group(x, index):
    clustering = KMeans(n_clusters = 3, random_state = 42).fit(x.values)
    return pd.Series(clustering.labels_, index = index)

并将其应用于我的df，如下所示：

df \ 
    .groupby(['k', 'group']) \ 
    .filter(lambda x: x.shape[0] > 3)['var'] \ 
    .reset_index('object') \ 
    .groupby(['k', 'group']) \ 
    .apply(lambda x: cluster_group(x['var'], x['object']))

但是，由于原始 DataFrame 非常庞大，因此此解决方案的运行速度非常慢。因此，我想问一下是否有办法以某种方式优化性能？

我的机器是 Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz，有 6 个内核和 12 个线程。df.shape 是 (1286135, 9)，但我需要为许多数据帧计算这个，每个数据帧大约是这个尺寸。所以我需要确保代码尽可能的优化。

【问题讨论】：

标签： python pandas scikit-learn

【解决方案1】：

您最初的解决方案还不错。我尝试使用pd.DataFrame.unstack() 方法来加快速度。

可能的解决方案：

df = df.unstack()['var'].apply(lambda x: cluster_group(x, x.index), axis=1)

上面例子的速度测试：

%timeit -n 10 test = df \
    .groupby(['k', 'group']) \
    .filter(lambda x: x.shape[0] > 3)['var'] \
    .reset_index('object') \
    .groupby(['k', 'group']) \
    .apply(lambda x: cluster_group(x['var'], x['object']))

每个循环 234 毫秒 ± 20.9 毫秒（7 次运行的平均值 ± 标准偏差，每次 10 次循环）

%timeit -n 10 test = df.unstack()['var'].apply(lambda x: cluster_group(x, x.index), axis=1)

每个循环 205 毫秒 ± 26.8 毫秒（7 次运行的平均值 ± 标准偏差，每次 10 次循环）

【讨论】：