如何使这段代码更高效KS测试

问题描述

我的数据看起来像这样

id1,id2,similarity
CHEMBL1,CHEMBL1,1
CHEMBL2,0.18
CHEMBL3,0.56
CHEMBL4,0.64
CHEMBL5,0.12
CHEMBL1,CHEMBL2,0.18
CHEMBL2,1
CHEMBL3,0.26
CHEMBL4,0.78
CHEMBL5,0.33
CHEMBL1,CHEMBL3,0.56
CHEMBL2,0.26
CHEMBL3,1
CHEMBL4,0.04
CHEMBL5,0.85
CHEMBL1,CHEMBL4,0.64
CHEMBL2,0.78
CHEMBL3,0.04
CHEMBL4,1
CHEMBL5,0.49
CHEMBL1,CHEMBL5,12
CHEMBL2,0.33
CHEMBL3,0.85
CHEMBL4,0.49
CHEMBL5,1

整个文件约为1.97亿行（10 GB）。我的目标是比较第1列中每种化合物的第3列的分布。通过大量重构，我设法拥有了这段代码

import pandas as pd
from scipy.stats import ks_2samp
import re

with open('example.csv','r') as f,open('Metrics.tsv','a') as f_out:
    f_out.write('compound_1' + '\t' + 'compound_2' + '\t' + 'Similarity' + '\t' + 'KS Distance' + '\n')
    df = pd.read_csv(f,delimiter = ',',lineterminator = '\n',header = None)
    d = {}
    l_id1 = []
    l_id2 = []
    l_sim = []
    uniq_comps = df.iloc[:,0].unique().tolist()
    for i in uniq_comps:
        d[i] = []
    for j in range(df.shape[0]):
        d[df.iloc[j,0]].append(df.iloc[j,2])
        l_id1.append(df.iloc[j,0])
        l_id2.append(df.iloc[j,1])
        l_sim.append(df.iloc[j,2])
    for k in range(len(l_id1)):
        sim = round(l_sim[k]*100,0)/100
        ks = re.findall(r"statistic=(.*)\,.*$",str(ks_2samp(d[l_id1[k]],d[l_id2[k]])))
        f_out.write(l_id1[k] + '\t' + l_id2[k] + '\t' + str(sim) + '\t' + str(''.join(ks)) + '\n')

运行，但是按预期运行非常慢。是否有人对如何使其更快化有任何想法？我想要的输出看起来像这样

 compound_1,compound_2,Similarity,KS Distance
CHEMBL1,1.0,0.0
CHEMBL2,0.18,0.4
CHEMBL3,0.56,0.2
CHEMBL4,0.64,0.2
CHEMBL5,0.12,0.4
CHEMBL1,0.4
CHEMBL2,0.0
CHEMBL3,0.26,0.78,0.4
CHEMBL5,0.33,0.2
CHEMBL1,0.2
CHEMBL2,0.2
CHEMBL3,0.0
CHEMBL4,0.04,0.85,0.0
CHEMBL5,0.49,12.0,0.0

由于数据量大，在Pyspark中运行它会更明智吗？如果是这样，怎么能达到类似的效果？

解决方法

代码检查

有些重点应该突出显示，并且可以提高性能：

• 在熊猫中打开CSV时，您已经将所有数据加载到RAM中，因此不必将该数据复制到列表中（例如l_id1，l_id2等）。 ）。尽可能避免拥有多个数据副本，这会降低性能，并使代码难以调试。
• 在处理Pandas DataFrame时，请尝试避免尽可能频繁地编写显式循环，应该有一个适合您的方法，例如groupby。
• Scipy统计信息包会返回一个Result对象，该对象始终公开statistic和pvalue成员，使用它而不是将其转换为字符串，然后使用正则表达式提取值。
• 避免不必要地调用昂贵的函数，在上一个循环中，您为“两次样本KS测试”计算的时间量相同，而不是每次计算一次，然后将结果与数据集合并。

重构

由于您似乎可以使用熊猫打开CSV，因此我认为完整的文件适合您的内存。检查两次，数字数据应适合2Gb RAM。

8 bytes*197e6 rows/1024**3 ~ 1.47 Gb

不清楚要计算什么。我将假设您要基于id1列收集数据，然后要使用基于每个可能的标识符对的“两次样本Kolmogorov-Smirnow”检验来检查分布是否等效。如果这不是您想要执行的操作，请更新您的帖子以详细说明您打算计算的内容。

让我们创建一个试用数据框：

import itertools
import numpy as np
import pandas as pd
from scipy import stats

N = 10**6
df = pd.DataFrame({
    "id1": np.random.choice([f"CHEMBL{i:d}" for i in np.arange(1,6)],N),"id2": np.random.choice([f"CHEMBL{i:d}" for i in np.arange(1,"value": np.random.uniform(0,12,N)
})

试验数据集如下：

       id1      id2      value
0  CHEMBL4  CHEMBL3  10.719870
1  CHEMBL2  CHEMBL5   2.911339
2  CHEMBL4  CHEMBL4   0.001595
3  CHEMBL2  CHEMBL3   0.148120
4  CHEMBL5  CHEMBL2   4.683689

创建DataFrame后，很容易使用groupby方法按标识符对数据进行分组。然后，我们可以对所有可能的标识符对进行统计检验。如果我们将所有组件组装在一个生成器中，那就是：

def apply_test(df,idkey="id",valuekey="value",test=stats.ks_2samp):
    """
    Apply statistical test to each possible pair of identifier
    """
    # Group by identifier:
    g = df.groupby(idkey)
    # Generate all 2-combination of identifier:
    for k1,k2 in itertools.combinations(g.groups.keys(),2):
        # Apply Statistical Test to grouped data:
        t = test(df.loc[g.groups[k1],valuekey],df.loc[g.groups[k2],valuekey])
        # Store Identifier pair:
        res = {"id1": k1,"id2": k2}
        # Store statistics and p-value:
        res.update({k: getattr(t,k) for k in t._fields})
        # Yield result:
        yield res

这时，只需将功能应用于数据框：

r = pd.DataFrame([x for x in apply_test(df)])

返回试验数据集：

       id1      id2  statistic    pvalue
0  CHEMBL1  CHEMBL2   0.002312  0.657859
1  CHEMBL1  CHEMBL3   0.002125  0.756018
2  CHEMBL1  CHEMBL4   0.001701  0.934290
3  CHEMBL1  CHEMBL5   0.002560  0.527594
4  CHEMBL2  CHEMBL3   0.002155  0.741524
5  CHEMBL2  CHEMBL4   0.001766  0.914602
6  CHEMBL2  CHEMBL5   0.003035  0.315677
7  CHEMBL3  CHEMBL4   0.001668  0.944053
8  CHEMBL3  CHEMBL5   0.002603  0.507482
9  CHEMBL4  CHEMBL5   0.002661  0.479805

然后我们可以将这些结果与原始数据框合并：

df.merge(r)

            id1      id2      value  statistic    pvalue
0       CHEMBL2  CHEMBL5   2.911339   0.003035  0.315677
1       CHEMBL2  CHEMBL5   6.583948   0.003035  0.315677
2       CHEMBL2  CHEMBL5  10.237092   0.003035  0.315677
3       CHEMBL2  CHEMBL5   8.049175   0.003035  0.315677
4       CHEMBL2  CHEMBL5   3.977925   0.003035  0.315677
...         ...      ...        ...        ...       ...
400776  CHEMBL4  CHEMBL5   4.339528   0.002661  0.479805
400777  CHEMBL4  CHEMBL5   5.353133   0.002661  0.479805
400778  CHEMBL4  CHEMBL5  10.599985   0.002661  0.479805
400779  CHEMBL4  CHEMBL5   9.701375   0.002661  0.479805
400780  CHEMBL4  CHEMBL5   7.951454   0.002661  0.479805