如何加快（并行化）分组的逐行滚动平均值计算？

Question

我正在计算一个大型数据集的分组逐行移动平均值。 但是，该过程在单个线程上花费的时间太长。 我怎样才能有效地加快这个过程？

请在下面找到一个可重现的示例：

dataframe = pd.DataFrame({'id': range(2),
                   'group_id': range(2),
                   'Date_1_F1': [1,2], 
                   'Date_2_F1': [2,4], 
                   'Date_3_F1': [3, 6],
                   'Date_4_F1': [4,8], 
                   'Date_1_F2': [2,11], 
                   'Date_2_F2': [6, 13], 
                   'Date_3-F2': [10, 15],
                   'Date_4_F2': [14, 17]})
dataframe

   id  group_id  Date_1_F1  ...  Date_2_F2  Date_3-F2  Date_4_F2
0   0         0          1  ...          6         10         14
1   1         1          2  ...         13         15         17

我有一个函数可以返回数据集的（逐行）平滑版本。

def smooth_ts(dataframe, ma_parameter = 2):

    dataframe = (dataframe
                 .set_index(["id", "group_id"])
                 .groupby(lambda x: x.split("_")[-1], axis = 1, group_keys=False)
                 .apply(lambda x: x.rolling(ma_parameter, axis = 1)
                        .mean()
                        .dropna(axis=1, how='all')))
    
    dataframe.reset_index(inplace = True)
    
    return dataframe

smoothed_df = smooth_ts(dataframe)

非常感谢

Answer 1

您可以 (1) 使用pd.melt融化您的数据框，(2) 创建您的分组变量，(3) 通过rolling.mean(2)对其进行排序和分组。 然后您可以使用df.pivot显示所需的数据。 在这种方法中，有一种可以使用swifter并行化的 apply 方法。 这是一个例子：

import pandas as pd
import numpy as np

import swifter

dataframe = pd.DataFrame({'id': range(2),
                   'group_id': range(2),
                   'Date_1_F1': [1,2], 
                   'Date_2_F1': [2,4], 
                   'Date_3_F1': [3, 6],
                   'Date_4_F1': [4,8], 
                   'Date_1_F2': [2,11], 
                   'Date_2_F2': [6, 13], 
                   'Date_3-F2': [10, 15],
                   'Date_4_F2': [14, 17]})

df_melted = pd.melt(dataframe, id_vars=['id', 'group_id'])
# Use next line if you want to parallelize the apply method 
# df_melted['groups'] = df_melted['variable'].str.split('_').swifter.apply(lambda v: v[-1])
df_melted['groups'] = df_melted['variable'].str.split('_').apply(lambda v: v[-1])

df_melted = df_melted.sort_values(['id', 'group_id', 'groups'])

df_tmp = df_melted.copy()
df_tmp['rolling_val'] = df_tmp.groupby(['id', 'group_id', 'groups'])['value'].rolling(2).mean().values
df_tmp.pivot(index=['id', 'group_id'], columns='variable', values='rolling_val').dropna(axis=1).reset_index().rename_axis(None, axis=1)

如果你想坚持你的方法，你可以使用multiprocessing库中的Pool对象来加速它，它将函数映射到迭代器并行化。

import pandas as pd
import numpy as np
from multiprocessing import Pool


dataframe = pd.DataFrame({'id': range(2),
                   'group_id': range(2),
                   'Date_1_F1': [1,2], 
                   'Date_2_F1': [2,4], 
                   'Date_3_F1': [3, 6],
                   'Date_4_F1': [4,8], 
                   'Date_1_F2': [2,11], 
                   'Date_2_F2': [6, 13], 
                   'Date_3-F2': [10, 15],
                   'Date_4_F2': [14, 17]})
dataframe

def smooth_ts(dataframe, ma_parameter = 2):

    dataframe = (dataframe
                 .set_index(["id", "group_id"])
                 .groupby(lambda x: x.split("_")[-1], axis = 1, group_keys=False)
                 .apply(lambda x: x.rolling(ma_parameter, axis = 1)
                        .mean()
                        .dropna(axis=1, how='all')))
    
    dataframe.reset_index(inplace = True)
    
    return dataframe

id_chunks = np.array_split(dataframe.id.unique(), 2) # 2 : number of splits => corresponds to number of chunks
df_chunks = [dataframe[dataframe['id'].isin(i)] for i in id_chunks] # list containing chunked data frames
with Pool(2) as p: dfs_chunks = p.map(smooth_ts, df_chunks) # applies function smooth_ts to list of data frames, use two processors as dfs_chunks only contain two data frames. For more chunks, number of processors can be increased
pd.concat(dfs_chunks).reset_index(drop=True)