GlobalAlignment.py

from sys import maxsize


class Alignmenter(object):
    def __init__(self, array1, array2):
        self.array1 = array1
        self.array2 = array2
        self.rows = len(self.array1) + 1
        self.cols = len(self.array2) + 1
        self.table = []
        self._penalty = 1

    def get_rows(self):
        return self.rows

    def get_cols(self):
        return self.cols

    def set_table(self, table):
        self.table = table

    def set_penalty_score(self, penalty: int):
        self._penalty = penalty

    def score_func(self, i, j):
        if self.array1[i-1] == self.array2[j-1]:
            return 1
        else:
            return -self._penalty

    def align_make_table(self):
        table = []
        for row in range(self.rows):
            table.append([0 for col in range(self.cols)])
        for row in range(self.rows):
            table[row][0] = -row * self._penalty
        for col in range(self.cols):
            table[0][col] = -col * self._penalty
        for row in range(1, self.rows):
            for col in range(1, self.cols):
                cell_above = table[row-1][col] - self._penalty
                cell_left = table[row][col-1] - self._penalty
                cell_left_above = \
                    table[row-1][col-1] + self.score_func(row, col)
                table[row][col] = max(
                    cell_above, cell_left, cell_left_above)
        return table

    def _traceback(self, cur_pos, path):
        row = cur_pos[0]
        col = cur_pos[1]
        if(row == 0 and col == 0):
            return [(0, 0)] + path
        if row > 0:
            cell_above_score = self.table[row-1][col]
            cell_above = (cell_above_score, (row-1, col))
        else:
            cell_above_score = -maxsize
            cell_above = (cell_above_score, (row-1, col))
        if col > 0:
            cell_left_score = self.table[row][col-1]
            cell_left = (cell_left_score, (row, col-1))
        else:
            cell_left_score = -maxsize
            cell_left = (cell_left_score, (row, col-1))
        if row > 0 and col > 0:
            cell_left_above_score = \
                self.table[row-1][col-1] + self.score_func(row, col)
            cell_left_above = (cell_left_above_score, (row-1, col-1))
        else:
            cell_left_above_score = -maxsize
            cell_left_above = (cell_left_above_score, (row-1, col-1))

        data_list = [cell_above, cell_left, cell_left_above]
        next_pos = max(data_list, key=lambda x: x[0])[1]
        return self._traceback(next_pos, [(row, col)] + path)

    def traceback(self):
        return self._traceback((self.rows-1, self.cols-1), [])

    def get_opt_align(self, path):
        align1 = []
        align2 = []
        prev_pos = path[0]
        path = path[1:]
        for cur_pos in path:
            move_below = cur_pos[0] - prev_pos[0]
            move_right = cur_pos[1] - prev_pos[1]
            prev_pos = cur_pos
            if move_below and move_right:
                align1.append(self.array1[cur_pos[0]-1])
                align2.append(self.array2[cur_pos[1]-1])
            elif move_below:
                align1.append(self.array1[cur_pos[0]-1])
                align2.append("-")
            elif move_right:
                align1.append("-")
                align2.append(self.array2[cur_pos[1]-1])
        return (align1, align2)


if __name__ == '__main__':
    arr1 = "MVHLTPEEKSAVTALWGKVNVDEVGGEALGRLLVVYPWTQRFFESFGDLSTPDAVMGNPKVKAHGKKVLGAFSDGLAHLDNLKGTFATLSELHCDKLHVDPENFRLLGNVLVCVLAHHFGKEFTPPVQAAYQKVVAGVANALAHKYH"
    arr2 = "KVKAHGKKVLHSFGEGVHHLDNLKGTFAALSELHCDKLHVDPENFRLLGNVLVVVLARHFGVQLSGEEKAAVLALWDKVNEEEVGGEALGRLLVVYPWTQRFFDSFGDLSNPGAVMGNPKDFTPELQASYQKVVAGVANALAHKYH"
    a = Alignmenter(list(arr1), list(arr2))
    table = a.align_make_table()
    a.set_table(table)
    path = a.traceback()
    (align1, align2) = a.get_opt_align(path)
    print(arr1)
    print(arr2)
    print("".join(align1))
    print("".join(align2))