Local_alignment_tool.py

###NAME: GIULIO RANDAZZO
###MATRICOLA: 231391


import numpy as np
from enum import IntEnum
import argparse

#using argparse to handle inputs arguments
#sequence 1 and 2 must to be given following the -s1 and -s2 notation, respectively

def main():
    parser = argparse.ArgumentParser(
        description="DNA alignment tool"
    )

    parser.add_argument(
        "-s1", "--sequence1",
        type=str,
        required=True,
        help="First DNA sequence (for instance, 'AGCT')"
    )
    parser.add_argument(
        "-s2", "--sequence2",
        type=str,
        required=True,
        help="Second DNA sequence (for instance, 'TGCA')"
    )

    # Optional argument for match score
    parser.add_argument(
        "-m", "--match",
        type=int,
        default=3,
        help="Match score (default: 3)"
    )

    # Optional argument for mismatch score
    parser.add_argument(
        "-mm", "--mismatch",
        type=int,
        default=-3,
        help="Mismatch score (default: -3)"
    )

    # Optional argument for gap score
    parser.add_argument(
        "-g", "--gap",
        type=int,
        default=-2,
        help="Gap score (default: -2)"
    )

    parser.add_argument(
        "-nba", "--numOfBestAlignments",
        type=int,
        default=10000000000000,
        help="Number of best aligments outputted (default: 10e13)"
    )

    parser.add_argument(
        "-fsh", "--filterScoreHigherThan",
        type=int,
        default=-1,
        help="Filter for alignments with score higher than the value specified (default: -1)"
    )

    parser.add_argument(
        "-fsl", "--filterScoreLowerThan",
        type=int,
        default=1000000000000,
        help="Filter for alignments with score lower than the value specified (default: 10e12)"
    )
    parser.add_argument(
        "-flh", "--filterLengthHigher",
        type=int,
        default=-1,
        help="Filter for alignments with length higher than the value specified (default: -1)"
    )
    parser.add_argument(
        "-fll", "--filterLengthLower",
        type=int,
        default=1000000000000,
        help="Filter for alignments with length lower than the value specified (default: 10e12)"
    )
    parser.add_argument(
        "-fmh", "--filterMatchHigher",
        type=int,
        default=-1,
        help="Filter for alignments with number of matches higher than the value specified (default: -1)"
    )
    parser.add_argument(
        "-fml", "--filterMatchLower",
        type=int,
        default=1000000000000,
        help="Filter for alignments with number of matches lower than the value specified (default: 10e12)"
    )
    parser.add_argument(
        "-fmmh", "--filterMismatchHigher",
        type=int,
        default=-1,
        help="Filter for alignments with number of mismatches higher than the value specified (default: -1)"
    )
    parser.add_argument(
        "-fmml", "--filterMismatchLower",
        type=int,
        default=1000000000000,
        help="Filter for alignments with number of mismatches lower than the value specified (default: 10e12)"
    )
    parser.add_argument(
        "-fgh", "--filterGapHigher",
        type=int,
        default=0,
        help="Filter for alignments with number of gaps higher than the value specified (default: -1)"
    )
    parser.add_argument(
        "-fgl", "--filterGapLower",
        type=int,
        default=1000000000000,
        help="Filter for alignments with number of gaps higher than the value specified (default: 10e12)"
    )
    # Argument parsing
    args = parser.parse_args()

    seq1 = args.sequence1
    seq2 = args.sequence2

    # Print sequences and parameters as a check 
    print("First sequence:", args.sequence1)
    print("Second sequence:", args.sequence2)
    print("Match score:", args.match)
    print("Mismatch score:", args.mismatch)
    print("Gap score:", args.gap)


    # Assigning the constants for the scores
    class Score(IntEnum):
        MATCH = args.match
        MISMATCH = args.mismatch
        GAP = args.gap

    # Assigning the constant values for the traceback
    class Trace(IntEnum):
        STOP = 0
        LEFT = 1 
        UP = 2
        DIAGONAL = 3

    def smith_waterman(seq1, seq2):
        # Generating the empty matrices for storing scores and tracing
        row = len(seq1) + 1
        col = len(seq2) + 1
        matrix = np.zeros(shape=(row, col), dtype=np.int32)  
        tracing_matrix = np.zeros(shape=(row, col), dtype=np.int32)  
        
        # Calculating the scores for all cells in the matrix
        for i in range(1, row):
            for j in range(1, col):
                # Calculating the diagonal score (match score)
                match_value = Score.MATCH if seq1[i - 1] == seq2[j - 1] else Score.MISMATCH
                diagonal_score = matrix[i - 1, j - 1] + match_value
                
                # Calculating the vertical gap score
                vertical_score = matrix[i - 1, j] + Score.GAP
                
                # Calculating the horizontal gap score
                horizontal_score = matrix[i, j - 1] + Score.GAP
                
                # Taking the highest score 
                matrix[i, j] = max(0, diagonal_score, vertical_score, horizontal_score)
                
                 # Tracking where the cell's value is coming from    
                if matrix[i, j] == 0: 
                    tracing_matrix[i, j] = Trace.STOP
                    
                elif matrix[i, j] == vertical_score: 
                    tracing_matrix[i, j] = Trace.UP

                elif matrix[i, j] == horizontal_score: 
                    tracing_matrix[i, j] = Trace.LEFT
                    
                elif matrix[i, j] == diagonal_score: 
                    tracing_matrix[i, j] = Trace.DIAGONAL

        list_of_scores = []
        for i in range(1, row):
            for j in range(1, col):
                if int(matrix[i,j]) !=0:
                    list_of_scores.append((int(matrix[i,j]), i, j))

        # Ordinamento per score in ordine decrescente
        sorted_tuples = sorted(list_of_scores, key=lambda x: x[0], reverse=True)

        # Initialising the variables for tracing
        aligned_seq1 = ""
        aligned_seq2 = ""   
        current_aligned_seq1 = ""   
        current_aligned_seq2 = ""

        #List of already tracked cells
        already_tracked_cells = []
        
        # Tracing and computing the pathway with the local alignment
        list_of_aligments = []
        for tuple in sorted_tuples:
            max_i = int(tuple[1])
            max_j = int(tuple[2])
            score = tuple[0]
            count = 0
            if (int(max_i),int(max_j)) in already_tracked_cells:
                continue
            else:
                already_tracked_cells.append((int(max_i),int(max_j)))
            while tracing_matrix[max_i, max_j] != Trace.STOP:
                if tracing_matrix[max_i, max_j] == Trace.DIAGONAL:
                    current_aligned_seq1 = seq1[max_i - 1]
                    current_aligned_seq2 = seq2[max_j - 1]
                    if current_aligned_seq1 != current_aligned_seq2 and count == 0:
                        break
                    count += 1
                    max_i = max_i - 1
                    max_j = max_j - 1
                    if (int(max_i), int(max_j)) in already_tracked_cells:
                        aligned_seq1 = ""
                        aligned_seq2 = ""
                        break
                    already_tracked_cells.append((int(max_i),int(max_j)))
                
                elif tracing_matrix[max_i, max_j] == Trace.UP:
                    if count == 0:
                        break
                    else:
                        current_aligned_seq1 = seq1[max_i - 1]
                        current_aligned_seq2 = '_'
                        max_i = max_i - 1
                        if (int(max_i), int(max_j)) in already_tracked_cells:
                            aligned_seq1 = ""
                            aligned_seq2 = ""
                            break    
                        already_tracked_cells.append((int(max_i),int(max_j)))

                elif tracing_matrix[max_i, max_j] == Trace.LEFT:
                    if count == 0:
                        break
                    else:
                        current_aligned_seq1 = '_'
                        current_aligned_seq2 = seq2[max_j - 1]
                        max_j = max_j - 1
                        if (int(max_i), int(max_j)) in already_tracked_cells:
                            aligned_seq1 = ""
                            aligned_seq2 = ""
                            break
                        already_tracked_cells.append((int(max_i),int(max_j)))


                aligned_seq1 = aligned_seq1 + current_aligned_seq1
                aligned_seq2 = aligned_seq2 + current_aligned_seq2
            
            # Reversing the order of the sequences
            aligned_seq1 = aligned_seq1[::-1]
            aligned_seq2 = aligned_seq2[::-1]
            num_matches = 0
            num_mismatches = 0
            num_gaps = 0
            align_len = len(aligned_seq1)
            for c in range(len(aligned_seq1)):
                if aligned_seq1[c] == aligned_seq2[c]:
                    num_matches += 1
                elif aligned_seq1[c] != aligned_seq2[c] and (aligned_seq1[c] != "_" and aligned_seq2[c] != "_"):
                    num_mismatches +=1
                elif aligned_seq1[c] != aligned_seq2[c] and (aligned_seq1[c] == "_" or aligned_seq2[c] == "_"):
                    num_gaps +=1
            list_of_aligments.append((aligned_seq1, aligned_seq2, score, num_matches, num_mismatches, num_gaps, align_len))
            # Re-initialising the variables for tracing
            num_matches = 0
            num_mismatches = 0
            num_gaps = 0
            align_len = 0
            aligned_seq1 = ""
            aligned_seq2 = ""   
            current_aligned_seq1 = ""   
            current_aligned_seq2 = ""  
        return list_of_aligments

    output = smith_waterman(seq1, seq2)
    
    max_score = int(output[0][2])
    list_al = sorted(output, key=lambda x: x[6], reverse=True)

    if list_al != []:
        output_string=""
        if args.numOfBestAlignments > 1:
            output_string += "\n----------------------------------\nAlignments are shown in decreasing order of length\n----------------------------------\n"
        best_align_counter = 0
        for i in range(len(output)):
            align1 = list_al[i][0]
            align2 = list_al[i][1]
            align_score = list_al[i][2]
            num_matches = list_al[i][3]
            num_mismatches = list_al[i][4]
            num_gaps = list_al[i][5]
            align_len = list_al[i][6]
            match_string = ""
            cons_mismatch = 0
            for t in range(len(align1)):
                if align1[t] == align2[t] and (align1[t] != "_" and align2[t] != "_"):
                    cons_mismatch += 1
                    if cons_mismatch == 3:
                        break
                else:
                    cons_mismatch = 0

            if float(max_score)*0.6 >= float(align_score):     #filter criteria
                continue
            elif cons_mismatch <= 2:
                continue
            elif align1 == "" or align2 == "":
                continue
            elif args.filterScoreLowerThan <= float(align_score):
                continue
            elif args.filterLengthHigher >= float(align_len):
                continue
            elif args.filterLengthLower <= float(align_len):
                continue
            elif args.filterMatchHigher >= float(num_matches):
                continue
            elif args.filterMatchLower <= float(num_mismatches):
                continue
            elif args.filterMismatchHigher >= float(num_mismatches):
                continue
            elif args.filterMismatchLower <= float(num_mismatches):
                continue
            elif args.filterGapHigher >= float(num_gaps):
                continue
            elif args.filterGapLower <= float(num_gaps):
                continue
            else:
                best_align_counter += 1
                if best_align_counter > args.numOfBestAlignments:
                    break
                else:
                    for c in range(len(align1)):
                        if align1[c] == align2[c]:
                            match_string += "|"
                        else:
                            match_string += " "
                    output_string += "\n"+"-"*20+"\n" + str(align1) + "\n" + match_string + "\n" + str(align2) + "\nScore of the alignment: " + str(align_score) +"\n"
            output_string += "Alignment length: {}\nNumber of matches: {}\nNumber of mismatches: {}\nNumber of gaps: {}\n".format(align_len, num_matches, num_mismatches, num_gaps) +"-"*20+"\n"
        if output_string == "\n----------------------------------\nAlignments are shown in decreasing order of score\n----------------------------------\n":
            print("\n"+"-"*20+"\nNo alignments satisfying the criteria\n"+"-"*20)
        else:
            print(output_string)
            #print(output)
        
        
    else:
        print("No alignments found!")

if __name__ == "__main__":
    main()