分组 dataframe 并使用 Python/Dask/multiprocessing 以并行方式对结果矩阵执行操作？

Question

我正在做一个项目，我需要根据分子的 ID 对数据库中的分子进行分组，并对生成的矩阵执行操作。 我正在使用 Python，我想通过并行处理来提高性能。

我目前正在从 SDF 文件加载分子并将它们存储在 Pandas dataframe 中。每个分子都有一个 ID、一个唯一的 Pose ID 和一个唯一的结构。 我的目标是按 ID 对 dataframe 进行分组，并为每个 ID 组创建一个矩阵。 矩阵的行和列将对应于该 ID 组中分子的唯一姿势 ID。 然后，我可以计算矩阵中每个单元格的值，例如定义该单元格的分子之间的相似性。 然而，对分子的具体操作对于这个问题并不重要。 我主要是就如何使用 Dask 或 Multiprocessing 为并行计算设置这样的系统寻求建议，或者是否有其他更好的选择。

这是没有任何并行化的版本的要点（请注意我已经进行了大量修改以使我的问题更清楚，下面的代码输出所需的东西，但我正在寻找计算分子上的细胞而不是姿势 ID）： https：/ /gist.github.com/Tonylac77/abfd54b1ceef39f0d161fb6b21950edb

#Generate sample dataframe

import pandas as pd
df = pd.DataFrame(columns=['ID', 'Pose ID'])
ids = ['ID' + str(i) for i in range(1, 6)]
pose_ids = ['Pose ' + str(i) for i in range(1, 11)]
# For each ID, add 10 rows to the dataframe with the corresponding Pose IDs
df_list = []
for i in ids:
    temp_df = pd.DataFrame({'ID': [i] * 10, 'Pose ID': pose_ids})
    df_list.append(temp_df)
df= pd.concat(df_list)
print(df)

################

from tqdm import tqdm
import itertools
import functools
import numpy as np
from IPython.display import display
def full_matrix_calculation(df):
    #Here I am using just string concatenation as an example calculation, in reality i am calling external functions
    def matrix_calculation(df, id_list):
        matrices = {}
        calc_dataframes = []
        for id in tqdm(id_list):
            df_name = df[df['ID']==id]
            df_name.index = range(len(df_name['Pose ID']))
            matrix = pd.DataFrame(0.0, index=[df_name['Pose ID']], columns=df_name['Pose ID'])
            for subset in itertools.combinations(df_name['Pose ID'], 2):
                result = subset[0]+subset[1]
                matrix.iloc[df_name[df_name['Pose ID']==subset[0]].index.values, df_name[df_name['Pose ID']==subset[1]].index.values] = result
                matrix.iloc[df_name[df_name['Pose ID']==subset[1]].index.values, df_name[df_name['Pose ID']==subset[0]].index.values] = result
            matrices[id] = matrix
        return matrices
    id_list = np.unique(np.array(df['ID']))
    calculated_dfs = matrix_calculation(df, id_list)
    return calculated_dfs

calculated_dfs = full_matrix_calculation(df)
display(calculated_dfs)

我尝试过使用多处理，但是，我的实现似乎比非并行版本慢： https://gist.github.com/Tonylac77/b4bbada97ee2bab7c37d4a29079af574

def function(tuple):
    return tuple[0]+tuple[1]
def full_matrix_calculation(df):
    #Here I am using just string concatenation as an example calculation, in reality i am calling external functions
    def matrix_calculation(df, id_list):
        matrices = {}
        calc_dataframes = []
        for id in tqdm(id_list):
            df_name = df[df['ID']==id]
            df_name.index = range(len(df_name['Pose ID']))
            matrix = pd.DataFrame(0.0, index=[df_name['Pose ID']], columns=df_name['Pose ID'])
            with multiprocessing.Pool() as p:
                try:
                    results = p.map(function, itertools.combinations(df_name['Pose ID'], 2))
                except KeyError:
                    print('Incorrect clustering method selected')
                    return
                results_list = list(zip(itertools.combinations(df_name['Pose ID'], 2), results))
                for subset, result in results_list:
                    matrix.iloc[df_name[df_name['Pose ID']==subset[0]].index.values, df_name[df_name['Pose ID']==subset[1]].index.values] = result
                    matrix.iloc[df_name[df_name['Pose ID']==subset[1]].index.values, df_name[df_name['Pose ID']==subset[0]].index.values] = result
                matrices[id] = matrix
            for subset in itertools.combinations(df_name['Pose ID'], 2):
                result = subset[0]+subset[1]
                matrix.iloc[df_name[df_name['Pose ID']==subset[0]].index.values, df_name[df_name['Pose ID']==subset[1]].index.values] = result
                matrix.iloc[df_name[df_name['Pose ID']==subset[1]].index.values, df_name[df_name['Pose ID']==subset[0]].index.values] = result
            matrices[id] = matrix
        return matrices
    id_list = np.unique(np.array(df['ID']))
    calculated_dfs = matrix_calculation(df, id_list)
    return calculated_dfs

calculated_dfs = full_matrix_calculation(df)
display(calculated_dfs)

我也开始玩 Dask，但是我面临的主要问题是我需要一个 ID 的所有值都在同一个 dask 分区中，否则我将得到不完整的矩阵（如果我至少理解正确的话） . 我试图找到解决这个问题的方法（比如在 x 分区中分块等），但到目前为止无济于事。 如果有变化，将更新此线程。

欢迎任何建议来加快这些计算。 作为参考，我正在处理的实际数据集包含约 10000 个唯一 ID 和约 300000 个姿势 ID。 通过我在分子上运行的计算，其中一些需要 40 小时才能完成。

Answer 1

使用 Dask Dataframe 和 groupBy 应该非常简单：

ddf = your_dataframe_as_dask

def matrix_calculation(df, id):
    matrix = pd.DataFrame(0.0, index=[df['Pose ID']], columns=df_name['Pose ID'])
    for subset in itertools.combinations(df['Pose ID'], 2):
        result = subset[0]+subset[1]
        matrix.iloc[df[df['Pose ID']==subset[0]].index.values, df_name[df_name['Pose ID']==subset[1]].index.values] = result
        matrix.iloc[df[df['Pose ID']==subset[1]].index.values, df_name[df_name['Pose ID']==subset[0]].index.values] = result
    return matrix

ddf.groupby('ID').apply(matrix_calculation).compute()

请参阅https://examples.dask.org/dataframes/02-groupby.html#Groupby-Apply 。

这将使每个 ID 的工作并行化。

然后您可能想查看https://docs.dask.org/en/stable/scheduling.html以选择适合您需要的调度程序（Dataframe 默认是线程，这可能效率不高，具体取决于您的代码）。