如何加速1亿行基因对相关性数据与5k行金标准基因集的匹配及PR-AUC计算

我不确定标题能不能准确描述我的问题，还是一步步解释吧：

我有一个10k×10k的基因相关性矩阵，把它转换成了仅包含上三角对的DataFrame，大概有1亿行数据，格式如下：

另外我还有一个金标准TERMS表，大概5k行，包含ID、名称和该条目对应的基因列表：

我的处理流程是这样的：

• 遍历金标准表中的每一个复合物条目
• 将条目中的基因列表转换成所有可能的基因对（比如geneA-geneB、geneA-geneC等）
• 检查这些基因对是否存在于之前生成的相关性基因对DataFrame中
• 存在的标记为TP=1，不存在的标记为TP=0
• 根据TP的计数计算精确率、召回率以及PR-AUC分数

最终每个金标准条目都会得到对应的PR-AUC分数，示例结果如下：

下面是我实现的代码，处理1亿行的相关性基因对和5k条金标准数据大概需要25分钟，我想找到优化方法来缩短运行时间。

PS：PR-AUC计算部分我已经用C++编译代码实现了，只需要把排序后的TP序列传入就能返回分数，但运行时间还是没变化，我猜测问题出在遍历匹配的环节。

from sklearn import metrics
def compute_per_complex_pr(corr_df, terms_df):

    pairwise_df = binary(corr_df)
    pairwise_df = quick_sort(pairwise_df).reset_index(drop=True)
    
    # Precompute a mapping from each gene to the row indices in the pairwise DataFrame where it appears.
    gene_to_pair_indices = {}
    for i, (gene_a, gene_b) in enumerate(zip(pairwise_df["gene1"], pairwise_df["gene2"])):
        gene_to_pair_indices.setdefault(gene_a, []).append(i)
        gene_to_pair_indices.setdefault(gene_b, []).append(i)
    
    # Initialize AUC scores (one for each complex) with NaNs.
    auc_scores = np.full(len(terms_df), np.nan)
    
    # Loop over each gene complex
    for idx, row in terms_df.iterrows():
        gene_set = set(row.used_genes)
   
        
        # Collect all row indices in the pairwise data where either gene belongs to the complex.
        candidate_indices = set()
        for gene in gene_set:
            candidate_indices.update(gene_to_pair_indices.get(gene, []))
        candidate_indices = sorted(candidate_indices)
        
        if not candidate_indices:
            continue
        
        # Select only the relevant pairwise comparisons.
        sub_df = pairwise_df.loc[candidate_indices]
        # A prediction is 1 if both genes in the pair are in the complex; otherwise 0.
        predictions = (sub_df["gene1"].isin(gene_set) & sub_df["gene2"].isin(gene_set)).astype(int)
        
        if predictions.sum() == 0:
            continue
        
        # Compute cumulative true positives and derive precision and recall.
        true_positive_cumsum = predictions.cumsum()
        precision = true_positive_cumsum / (np.arange(len(predictions)) + 1)
        recall = true_positive_cumsum / true_positive_cumsum.iloc[-1]
        
        if len(recall) < 2 or recall.iloc[-1] == 0:
            continue
        
        auc_scores[idx] = metrics.auc(recall, precision)
    
    # Add the computed AUC scores to the terms DataFrame.
    terms_df["auc_score"] = auc_scores
    return terms_df
    
    
 
def binary(corr):
    stack = corr.stack().rename_axis(index=['gene1', 'gene2']).reset_index(name='score')
    stack = drop_mirror_pairs(stack)
    return stack
    
    
def quick_sort(df, ascending=False):
    order = 1 if ascending else -1
    sorted_df = df.iloc[np.argsort(order * df["score"].values)].reset_index(drop=True)
    return sorted_df


def drop_mirror_pairs(df):
    gene_pairs = np.sort(df[["gene1", "gene2"]].to_numpy(), axis=1)
    df.loc[:, ["gene1", "gene2"]] = gene_pairs
    df = df.loc[~df.duplicated(subset=["gene1", "gene2"], keep="first")]
    return df

以下是用于测试的模拟数据生成代码：

import numpy as np
import pandas as pd

# Set a random seed for reproducibility
np.random.seed(0)

# -------------------------------
# Create the 10,000 x 10,000 correlation matrix
# -------------------------------
num_genes = 10000
genes = [f"Gene{i}" for i in range(num_genes)]


rand_matrix = np.random.uniform(-1, 1, (num_genes, num_genes))
corr_matrix = (rand_matrix + rand_matrix.T) / 2
np.fill_diagonal(corr_matrix, 1.0)

corr_df = pd.DataFrame(corr_matrix, index=genes, columns=genes)


num_terms = 5000
terms_list = []

for i in range(1, num_terms + 1):
    # Randomly choose a number of genes between 10 and 40 for this term
    n_genes = np.random.randint(10, 41)
    used_genes = np.random.choice(genes, size=n_genes, replace=False).tolist()
    term = {
        "id": i,
        "name": f"Complex {i}",
        "used_genes": used_genes
    }
    terms_list.append(term)

terms_df = pd.DataFrame(terms_list)

# Display sample outputs (for verification, you might want to show the first few rows)
print("Correlation Matrix Sample:")
print(corr_df.iloc[:5, :5])  # print a 5x5 sample

print("\nTerms DataFrame Sample:")
print(terms_df.head())

运行函数的方式：

compute_per_complex_pr(corr_df, terms_df)

备注：内容来源于stack exchange，提问作者Yasir