moshi4 / pymsaviz Goto Github PK

>Genus1 species1 strain1 | accession_number1
----M----AD---A
>Genus2 species2 strain2 | accession_number2
----M----SD---A
etc.

>accession_number1 | Genus1 species1 strain1
----M----AD---A
>accession_number2 | Genus2 species2 strain2 
----M----SD---A
etc.

def make_pymsaviz_plot(path, name, outname, min_gap_length = 5, gap_fraction=0.05, gap_char="-", variable_consensus=0.4, variable_char="x", show_count=True, show_consensus=True, color_scheme="Clustal", sorted = False):
    # create the input and output file names
    infile = os.path.join(path, name)
    outfile = os.path.join(path, outname)

    # parse the input fasta file into an array of sequences
    sequences = []
    for record in SeqIO.parse(infile, "fasta"):
        sequences.append(record.seq)

    # for every position in the alignment, count the number of gaps and variable characters
    gap_count = np.zeros(len(sequences[0]))

    for sequence in sequences:
        for i, aa in enumerate(sequence):
            if aa == gap_char:
                gap_count[i] += 1


    # get the continuous stretches of gaps
    gap_stretches = []
    start = 0
    end = 0
    for i, count in enumerate(gap_count):
        if(count/len(sequences) > (1-gap_fraction)):
            end = i
        else:
            if end > start:
                gap_stretches.append((start+2, end+1))
            start = i
            end = i
    if end > start:
        gap_stretches.append((start+2, end+1))

    # remove the stretches of gaps that are too short
    gap_stretches = [stretch for stretch in gap_stretches if stretch[1] - stretch[0] > min_gap_length]
    print(gap_stretches)

    # create a new fasta file with the gap_stretches removed
    protein_accessions = []
    genus_species = []

    with open(outfile + "_gap_trimmed.fasta", "w") as f:
        for record in SeqIO.parse(infile, "fasta"):

            # get the protein accession and genus species
            protein_accessions.append(record.description.split("|")[0].strip())
            genus_species.append(record.description.split("|")[1])

            new_seq = ""
            for i, aa in enumerate(record.seq):
                if not any([i >= stretch[0]-1 and i <= stretch[1]-1 for stretch in gap_stretches]):
                    new_seq += aa
            f.write(">" + record.description + "\n")
            f.write(new_seq + "\n")

    # create a pymsaviz object from the both the non-trimmed and trimmed fasta file
    msa = pymsaviz.MsaViz(infile, show_count=show_count, show_consensus=show_consensus, sort=sorted, color_scheme=color_scheme)
    msa_trimmed = pymsaviz.MsaViz(outfile + "_gap_trimmed.fasta", show_count=show_count, show_consensus=show_consensus, sort=sorted, color_scheme=color_scheme)

    # add annotations to the pymsaviz object
    for gap in gap_stretches:
        msa.add_text_annotation(gap, "Gap Region", text_color="black", range_color="red")
    msa.savefig(outfile + "_gap.png")

    # add variable markers to the trimmed pymsaviz object to show for highly non-conserved regions
    high_variability = []
    identity_list = msa_trimmed._get_consensus_identity_list()

    for position, identity in enumerate(identity_list, 1):
        if identity < variable_consensus:
            high_variability.append(position)
    msa_trimmed.add_markers(high_variability, marker=variable_char, color="red")
    msa_trimmed.savefig(outfile + "_gap_trimmed.png")

    return msa, msa_trimmed