如何用Smith-Waterman算法计算序列相似度并生成相似分组排名DataFrame

首先要确认你已经定义了sub_cost函数（你的smith_waterman依赖这个但未给出实现），下面是一个通用的替换成本示例，你可以根据需求调整匹配/不匹配的分值：

import numpy as np
import pandas as pd
from itertools import combinations

# 定义替换成本函数：匹配得2分，不匹配扣1分
def sub_cost(a, b):
    return 2 if a == b else -1

# 保留你提供的Smith-Waterman函数
def smith_waterman(seq2, seq1, d=-8):
    m = len(seq1)
    n = len(seq2)
    mat = np.zeros((m+1, n+1)) # 创建空矩阵
    # 填充矩阵
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            diag = mat[i-1][j-1] + sub_cost(seq1[i-1], seq2[j-1])
            up = mat[i-1][j] + d
            left = mat[i][j-1] + d
            mat[i][j] = max(0, diag, up, left)
    # 找到矩阵中的最高分位置
    highest_value = np.where(mat == np.amax(mat))
    highest_value_location = list(zip(highest_value[0], highest_value[1]))[0]
    traceback_seq1, traceback_seq2 = '', ''
    i, j = highest_value_location[0], highest_value_location[1]
    # 回溯生成比对序列（此处我们只需要最高分，所以回溯部分不影响最终返回值）
    while i > 0 or j > 0:
        current_score = mat[i][j]
        diag_score = mat[i-1][j-1]
        left_score = mat[i][j-1]
        up_score = mat[i-1][j]
        if (current_score==0):
            break
        if (current_score == diag_score + sub_cost(seq1[i-1], seq2[j-1])):
            t1, t2 = seq2[j-1], seq1[i-1]
            i,j = i-1,j-1
        elif (current_score == up_score + d):
            t1, t2 = '-', seq1[i-1]
            i -= 1
        elif (current_score == left_score + d):
            t1, t2 = seq2[j-1], '-'
            j -= 1
        traceback_seq1 += t1
        traceback_seq2 += t2
    traceback_seq1 = (traceback_seq1[::-1])
    traceback_seq2 = (traceback_seq2[::-1])
    return np.amax(mat)

步骤1：生成所有唯一的变体对 #

我们用itertools.combinations生成不重复的变体组合（避免重复计算A vs B和B vs A，因为Smith-Waterman得分是对称的）：

# 基于你的DataFrame生成所有两两变体对
variant_pairs = list(combinations(finalDF['Variant ID'], 2))

步骤2：批量计算相似度得分 #

遍历每个变体对，提取对应序列并计算Smith-Waterman得分：

# 初始化结果存储列表
results = []

for var1, var2 in variant_pairs:
    # 根据Variant ID获取对应的序列
    seq1 = finalDF.loc[finalDF['Variant ID'] == var1, 'original_sequence'].iloc[0]
    seq2 = finalDF.loc[finalDF['Variant ID'] == var2, 'original_sequence'].iloc[0]
    # 计算相似度得分
    sim_score = smith_waterman(seq1, seq2)
    # 将结果存入列表
    results.append({
        'Variants': f"{var1} & {var2}",
        'Similarity': sim_score
    })

步骤3：生成带排名的结果DataFrame #

将结果转换为DataFrame，按相似度降序排序并添加排名列：

# 转换为DataFrame
result_df = pd.DataFrame(results)

# 按相似度降序排序，添加排名（相同得分将获得相同排名）
result_df = result_df.sort_values(by='Similarity', ascending=False)
result_df['RANK'] = result_df['Similarity'].rank(method='min', ascending=False).astype(int)

# 查看前5条结果
print(result_df.head())

优化提示（针对大数量序列） #

如果你的DataFrame包含成百上千个变体，上述方法会因Smith-Waterman的O(mn)复杂度变慢，可以尝试：

• 用joblib库实现并行计算，加速遍历过程
• 先过滤长度差异极大的序列对，减少不必要的计算

内容的提问来源于stack exchange，提问作者Varsha Gupta