Center-Star-Approximation: Center-String in Python identifizieren

#!/usr/bin/env python3
"""
Einfache Funktionen zur Berechnung des Center-Strings einer Liste von Sequenzen.
Berechnet auch verwandte Matrizen.

Version 1.1: Python3-ready
"""
import numpy
import pandas

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

def substitutionCost(a, b):
    """Gibt die Kosten für die Substitution von Zeichen a durch Zeichen b zurück"""
    return 0 if a == b else 3

def calculateCostMatrix(a, b, gapCost=2):
    """Berechnet für zwei Sequenzen a, b deren Edit-Distanz-Matrix"""
    #Initialisiert n+1 x m+1 Matrix. x-Achse <--> a, y-Achse <--> b
    xdim = len(b) + 1
    ydim = len(a) + 1
    matrix = numpy.zeros((xdim, ydim))
    #Erste Zeile und erste Spalte mit n*gap-Kosten initialisieren
    matrix[0,:] = numpy.arange(0, gapCost*ydim, step=gapCost)
    matrix[:,0] = numpy.arange(0, gapCost*xdim, step=gapCost)
    #Array füllen
    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):
    """Gibt die Edit-Distanz zweier Sequenzen a und b zurück"""
    return calculateCostMatrix(a, b, gapCost)[-1,-1]

def calculateSequenceDistanceMatrix(sequences, gapCost=2):
    """
    Berechnet eine Matrix der Edit-Distanzen zwischen einer Liste von Sequenzen.
    Gibt die Matrix zurück.
    """
    #Hinweis: Der Einfachheit halber wird JEDES Paar von Sequenzen berechnet.
    # Wenn Performance wichtig ist, würde man nur die Hälfte davon berechnen
    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):
    """
    Berechnet den Center-String (repräsentiert durch seinen Index im Sequenz-Array)
    für eine gegebene Sequenz-Edit-Distanz-Matrix.
    Gibt den 1-basierten Index des Center-Strings zurück
    """
    #Summen für jede Zeile berechnen
    rowSums = [sum(distanceMatrix[:,i]) for i in range(distanceMatrix.shape[1])]
    #Wenn dieselbe Zeilensumme mehrfach auftritt, das erste Vorkommen verwenden
    return rowSums.index(min(rowSums)) + 1 #+ 1: Eind-basierte Indizes

def getPandasCostMatrix(a, b, gapCost=2):
    """
    Berechnet eine Kostenmatrix und wandelt sie in einen pandas-DataFrame um.
    Vorteil gegenüber numpy ndarray: Beschriftete Zeilen und Spalten
    """
    return pandas.DataFrame(calculateCostMatrix(a, b, gapCost), index=list(" " + b), columns=list(" " + a))

def getPandasDistanceMatrix(sequences, gapCost=2):
    """
    Berechnet die Sequenzabstandsmatrix einer Liste von Sequenzen
    und wandelt das Ergebnis in einen beschrifteten pandas-DataFrame um
    """
    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__":
    #Beispielaufruf
    sequences = ["ACGTGC","ACCTG", "AGGCTT", "AGCC"]
    #EINE (beliebige) Edit-Distanz-Matrix ausgeben
    print("s2 vs s3: \n", getPandasCostMatrix(sequences[1], sequences[3]),"\n")
    #Sequenzabstandsmatrix ausgeben
    print("Sequence distances: \n", getPandasDistanceMatrix(sequences),"\n")
    #Center-String ausgeben
    print("Center string: s%d" % calculateCenterString(calculateSequenceDistanceMatrix(sequences, gapCost=2)))