在另一个数据框中搜索列对的更快方法

Question

我有一个名为df的大数据框，大约有4500万行，如下图所示。下载

gene1   gene2     score
0       PIGA  ATF7IP1 -0.047236
1       PIGB  ATF7IP2 -0.047236
2       PIGC  ATF7IP3 -0.047236
3       PIGD  ATF7IP4 -0.047236
4       PIGE  ATF7IP5 -0.047236

，我有一个小的数据框称为terms，大小约为3k行。

id                                gene_set
1                                 {HDAC4, BCL6}
2                                 {HDAC5, BCL6}
3                                 {HDAC7, BCL6}
4                 {NCOA3, KAT2B, EP300, CREBBP}
5            {NCAPD2, NCAPH, NCAPG, SMC4, SMC2}
...                   
2912                              {FOXO1, ESR1}
2913                               {APP, FOXO3}
2914                               {APP, FOXO1}
2915                               {APP, FOXO4}
2916    {MAP3K20, MAPK14, AKAP13, MAP2K3, PKN1}

，对于每一行，我检查terms数据集中是否存在gene1,gene2对。

我的代码工作得很好，但我想问是否有任何更快的想法?
我已经尝试了几个代码，但运行时间大约是相同的。

def search(g1,g2):
# search gene pair in the go terms
return sum(terms.gene_set.map(set([g1,g2]).issubset))

代码示例1

np.sum(np.vectorize(search)(df.gene1,df.gene2))

代码示例2

[search(g1, g2) for g1, g2 in zip(df.gene1,df.gene2)] 

代码示例3

df[['gene1','gene2']].apply(lambda x: search(x.gene1,x.gene2), axis=1 )

Answer 1

我看到一个明确的方法来加速一点你的search功能通过修改set([g1,g2])到{g1,g2}.issubset。由于不再需要将列表转换为集合，因此可以避免相当多的工作。

下面是一个简单的测试用例

In [4]: %timeit set([1,2]).issubset({1,2,3})
184 ns ± 2.08 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
In [5]: %timeit {1,2}.issubset({1,2,3})
120 ns ± 0.607 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)

你也可以在函数中传递terms数据帧，这样解释器就不需要在函数作用域之外寻找它，也许你可以使用pd.Series.sum()函数而不是调用Python函数(这也应该加快速度)。

def search(g1, g2, terms):
# search gene pair in the go terms
return terms.gene_set.map({g1,g2}.issubset).sum()

除此之外，您可以尝试df.apply()方法，使用polar或Dask来获得更大幅度的加速。

Answer 2

如果你.explode较小的术语数据帧来摆脱集合呢?

long_terms = terms.explode('gene_set')

>>> long_terms
id gene_set
0   1     BCL6
0   1    HDAC4
1   2     BCL6
1   2    HDAC5
2   3    HDAC7
2   3     BCL6
3   4    KAT2B
3   4   CREBBP
3   4    NCOA3
3   4    EP300
4   5     SMC2
4   5    NCAPH
4   5     SMC4
4   5   NCAPD2
4   5    NCAPG
5   6  ATF7IP2
5   6        A
6   7     PIGB
6   7        B
7   8     PIGB
7   8        C
7   8  ATF7IP2

你可以使用.isin

same_row = (
long_terms.gene_set.isin(df.gene1).groupby(level=0).any()
&
long_terms.gene_set.isin(df.gene2).groupby(level=0).any()
)
found = long_terms.loc[same_row]

>>> found
id gene_set
7   8     PIGB
7   8        C
7   8  ATF7IP2

查找对应的匹配项:

>>> df.gene1.isin(found.gene_set) & df.gene2.isin(found.gene_set)
0    False
1     True
2    False
3    False
4    False
dtype: bool

>>> df[df.gene1.isin(found.gene_set) & df.gene2.isin(found.gene_set)]
gene1    gene2     score
1  PIGB  ATF7IP2 -0.047236

示例使用:

>>> df
gene1    gene2     score
0  PIGA  ATF7IP1 -0.047236
1  PIGB  ATF7IP2 -0.047236
2  PIGC  ATF7IP3 -0.047236
3  PIGD  ATF7IP4 -0.047236
4  PIGE  ATF7IP5 -0.047236

>>> terms
id                            gene_set
0   1                       {BCL6, HDAC4}
1   2                       {BCL6, HDAC5}
2   3                       {HDAC7, BCL6}
3   4       {KAT2B, CREBBP, NCOA3, EP300}
4   5  {SMC2, NCAPH, SMC4, NCAPD2, NCAPG}
5   6                        {ATF7IP2, A}
6   7                           {PIGB, B}
7   8                  {PIGB, C, ATF7IP2}

terms = pd.DataFrame({
"id": [1, 2, 3, 4, 5, 6, 7, 8], 
"gene_set": [
{"HDAC4", "BCL6"},
{"HDAC5", "BCL6"},
{"HDAC7", "BCL6"},
{"NCOA3", "KAT2B", "EP300", "CREBBP"},
{"NCAPD2", "NCAPH", "NCAPG", "SMC4", "SMC2"},
{"A", "ATF7IP2"},
{"B", "PIGB"},
{"C", "ATF7IP2", "PIGB"},
]
})