如何让这段代码更高效（KS测试）

Question

我的数据看起来像这样

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

整个文件大约有 1.97 亿行 (10GB)。我的目标是比较第 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,CHEMBL1,1.0,0.0
CHEMBL2,CHEMBL1,0.18,0.4
CHEMBL3,CHEMBL1,0.56,0.2
CHEMBL4,CHEMBL1,0.64,0.2
CHEMBL5,CHEMBL1,0.12,0.4
CHEMBL1,CHEMBL2,0.18,0.4
CHEMBL2,CHEMBL2,1.0,0.0
CHEMBL3,CHEMBL2,0.26,0.2
CHEMBL4,CHEMBL2,0.78,0.4
CHEMBL5,CHEMBL2,0.33,0.2
CHEMBL1,CHEMBL3,0.56,0.2
CHEMBL2,CHEMBL3,0.26,0.2
CHEMBL3,CHEMBL3,1.0,0.0
CHEMBL4,CHEMBL3,0.04,0.2
CHEMBL5,CHEMBL3,0.85,0.2
CHEMBL1,CHEMBL4,0.64,0.2
CHEMBL2,CHEMBL4,0.78,0.4
CHEMBL3,CHEMBL4,0.04,0.2
CHEMBL4,CHEMBL4,1.0,0.0
CHEMBL5,CHEMBL4,0.49,0.2
CHEMBL1,CHEMBL5,12.0,0.4
CHEMBL2,CHEMBL5,0.33,0.2
CHEMBL3,CHEMBL5,0.85,0.2
CHEMBL4,CHEMBL5,0.49,0.2
CHEMBL5,CHEMBL5,1.0,0.0

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

Answer 1

代码检查

有一些关键点应该突出显示并可能会提高性能：

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

重构

由于您似乎可以使用 pandas 打开 CSV，因此我假设完整的文件适合您的记忆。检查两次，数值数据应该适合 2Gb 的 RAM。

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

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

让我们创建一个试用 DataFrame：

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, 6)], N),
    "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