中心星近似：在 Python 中识别中心字符串

#!/usr/bin/env python3
"""
计算序列列表中心字符串的简单函数。
还计算相关矩阵。

版本 1.1：Python3 就绪
"""
import numpy
import pandas

__author__  = "Uli Köhler"
__license__ = "Apache license v2.0"
__version__ = "1.1"

def substitutionCost(a, b):
    """返回用字符 b 替换字符 a 的成本"""
    return 0 if a == b else 3

def calculateCostMatrix(a, b, gapCost=2):
    """对于两个序列 a、b，计算它们的编辑距离矩阵"""
    #初始化 n+1 x m+1 矩阵。x 轴 <--> a，y 轴 <--> b
    xdim = len(b) + 1
    ydim = len(a) + 1
    matrix = numpy.zeros((xdim, ydim))
    #用 n*gap 成本初始化第一行和第一列
    matrix[0,:] = numpy.arange(0, gapCost*ydim, step=gapCost)
    matrix[:,0] = numpy.arange(0, gapCost*xdim, step=gapCost)
    #填充数组
    for y in range(1, ydim):
        for x in range(1, xdim):
            substCost = matrix[x-1, y-1] + substitutionCost(b[x-1], a[y-1])
            matrix[x,y] = min(substCost, matrix[x-1, y] + gapCost, matrix[x, y-1] + gapCost)
    return matrix

def getEditDistance(a, b, gapCost=2):
    """获取两个序列 a 和 b 的编辑距离"""
    return calculateCostMatrix(a, b, gapCost)[-1,-1]

def calculateSequenceDistanceMatrix(sequences, gapCost=2):
    """
    计算序列列表之间的编辑距离矩阵。
    返回矩阵。
    """
    #注意：为简单起见，这计算任意序列对。
    # 如果性能重要，只需计算其中一半
    dim = len(sequences)
    matrix = numpy.zeros((dim, dim))
    for i in range(len(sequences)):
        for j in range(len(sequences)):
            matrix[i,j] = getEditDistance(sequences[i], sequences[j], gapCost)
    return matrix

def calculateCenterString(distanceMatrix):
    """
    为给定序列编辑距离矩阵计算中心字符串（由其在序列数组中的索引表示）。
    返回中心字符串的 1 基索引
    """
    #计算任何行的总和
    rowSums = [sum(distanceMatrix[:,i]) for i in range(distanceMatrix.shape[1])]
    #如果相同的行分数出现多次，使用第一次出现
    return rowSums.index(min(rowSums)) + 1 #+ 1: One-based indices

def getPandasCostMatrix(a, b, gapCost=2):
    """
    计算成本矩阵并将其转换为 pandas data frame。
    相比 numpy ndarray 的优势：带标签的行和列
    """
    return pandas.DataFrame(calculateCostMatrix(a, b, gapCost), index=list(" " + b), columns=list(" " + a))

def getPandasDistanceMatrix(sequences, gapCost=2):
    """
    计算序列列表的序列距离矩阵
    并将结果转换为带标签的 pandas DataFrame
    """
    labels = ["s%d" % i for i in range(1, len(sequences) + 1)]
    return pandas.DataFrame(calculateSequenceDistanceMatrix(sequences, gapCost), index=labels, columns=labels)

if __name__ == "__main__":
    #示例调用
    sequences = ["ACGTGC","ACCTG", "AGGCTT", "AGCC"]
    #打印一个（任意）编辑距离矩阵
    print("s2 vs s3: \n", getPandasCostMatrix(sequences[1], sequences[3]),"\n")
    #打印序列距离矩阵
    print("Sequence distances: \n", getPandasDistanceMatrix(sequences),"\n")
    #打印中心字符串
    print("Center string: s%d" % calculateCenterString(calculateSequenceDistanceMatrix(sequences, gapCost=2)))