<commit_before>#!/usr/bin/env python
"""Updated version of VCF_from_FASTA.py that hopefully will simplify the
process of generating a VCF from Sanger reads aligned to a reference sequence.
Requires Biopython. Assumes the reference sequence is the first one in the
alignment. Takes one argument:
    1) FASTA multiple sequence alignment"""

import sys
import re
from datetime import date
try:
    from Bio import SeqIO
    fa_in = sys.argv[1]
except ImportError:
    sys.stderr.write('This script requires Biopython.\n')
    sys.exit(1)
except IndexError:
    sys.stderr.write(__doc__ + '\n')
    sys.exit(2)


def extract_ref_coords(alignment, ref_idx=0):
    """Extract the name and the coordinates from the reference sequence of the
    alignment. This expects the sequence to be named in SAM/BED region
    format."""
    ref = alignment[ref_idx]
    chrom = ref.name.split(':')[0]
    start = ref.name.split(':')[1].split('-')[0]
    return (chrom, start)


def check_ref_gaps(alignment, ref_idx=0):
    """Check the reference sequence for end gaps. If we find any end gaps on
    the left, we will throw an error because we cannot accurately calculate the
    positions relative to reference. If we find them on the right, we will
    throw a warning, but continue anyway."""
    ref = alignment[ref_idx]
    refseq = str(ref.seq)
    left_gap = re.compile(r'^-+[ACGTMRWSYKVHDBNacgtmrwsykvhdbn]')
    right_gap = re.compile(r'[ACGTMRWSYKVHDBNacgtmrwsykvhdbn]-+$')
    if left_gap.match(refseq):
        sys.stderr.write(
            """Error:
Reference sequence has end-gaps on the left. This will cause calculated
positions to be incorrect. Please remove the end gaps and re-run this script.
""")
        sys.exit(10)
    if right_gap.search(refseq):
        sys.stderr.write(
            """Warning:
reference sequence has end-gaps on the right. This is not an error, but some of
the variants at the end of the alignment will not be placed on the reference
sequence. You may either remove the right end-gap or remove the variants from
the resulting VCF that occur in the end-gap.\n""")
    return


def extract_variants(alignment, ref_idx=0):
    """Extract the positions of SNPs and indels in the Sanger reads aligned to
    the reference sequence."""
    snp_pos = []
    indel_pos = []
    # First, convert the alignment to a list of lists, as opposed to a list
    # of SeqRecord objects
    raw_aln = [list(s.seq) for s in alignment]
    # Transpose it so that we iterate over columns of the alignment
    t_raw_aln = zip(*raw_aln)
    # Start iterating across columns and saving positions of variant sites. We
    # keep track of the reference base and only increment the position counter
    # for when we see a non-gap character in the reference sequence.
    offset = 0
    for aln_column in t_raw_aln:
        # First, get the states that exist at this position
        states = set(aln_column)
        # Discard any 'N' bases
        states.discard('N')
        # And get the ref state
        ref_state = aln_column[ref_idx]
        # Use the ref state to get the alternate states
        alt_states = states - set(ref_state)
        # If there is a gap in this position, then we will append it to the
        # list of indel positions
        if '-' in states:
            indel_pos.append((offset, ref_state, alt_states))
        # Then, discard the gap character to look for SNPs
        states.discard('-')
        # If the length of the states is greater than 1, then we have a SNP
        if len(states) > 1:
            # We will calculate the following:
            #   Number of non-missing alleles
            #   Minor allele count
            #   Minor allele frequency
            # The reference IS included in these calculations.
            non_missing = [
                base
                for base
                in aln_column
                if base != '-'
                or base != 'N']
            acs = [aln_column.count(x) for x in states]
            afs = [c/len(non_missing) for c in acs]
            # Re-discard the gap character, just to be sure we do not count it
            # as an alternate state
            alt_states.discard('-')
            snp_pos.append(
                (offset, ref_state, alt_states, len(non_missing), min(acs),
                 min(afs)))
        # If the reference sequence is not a gap, then we increment our offset
        # counter.
        if ref_state != '-':
            offset += 1
    return (snp_pos, indel_pos)


def collapse_indels(indels):
    """Collapse indels by identifying runs of consecutive integers and merging
    those into a single entry."""
    # Sort the indel bases by their coordinate
    indel_srt = sorted(indels, key=lambda x: x[0])
    # Make a list to hold our aggregated indels
    agg_indel = []
    # We will now iterate over adjacent records - if they are consecutive, then
    # merge them. Else, break it and start a new record.
    curr_indel = []
    for ind, ind_adj in zip(indel_srt, indel_srt[1:]):
        # Unpack the alleles. It's a little silly, but we have to cast the set
        # to a list to subset it.
        curr_ref = ind[1]
        curr_alt = list(ind[2])[0]
        adj_ref = ind_adj[1]
        adj_alt = list(ind_adj[2])[0]
        if not curr_indel:
            curr_indel = [ind[0], curr_ref, curr_alt]
        # If the next position is not consecutive, append it and start over
        if ind_adj[0] - ind[0] > 1:
            agg_indel.append(curr_indel)
            curr_indel = []
        else:
            curr_indel[1] += adj_ref
            curr_indel[2] += adj_alt
    # The way we are iterating through the indel list means that we will always
    # leave off the last one. Append it after the loop finishes.
    agg_indel.append(curr_indel)
    return agg_indel


def adjust_indels(indels, alignment, ref_idx=0):
    """Adjust the indel positions so that they are offset by one, as required
    by the VCF spec. This is because the reported position must be the base
    *before* any insertion/deletion polymorphism is observed."""
    spec_indels = []
    # Remove the gaps from the reference sequnce for getting the reference base
    # of the indel
    ref_seq = ''.join([base for base in alignment[ref_idx].seq if base != '-'])
    for i in indels:
        spec_pos = i[0] - 1
        spec_ref = ref_seq[spec_pos]
        spec_indel = [spec_pos, spec_ref + i[1], spec_ref + i[2]]
        spec_indels.append(spec_indel)
    return spec_indels


def print_vcf(snp_var, ind_var, refseq, offset):
    """Print a VCF from the calculated positions of the variants."""
    # Define the VCF header
    today = date.today().strftime('%Y%m%d')
    vcf_header = """##fileformat=VCFv4.1
##fileDate={filedate}
##source=VCF_from_FASTA_2.py;Morrell Lab @ UMN
##INFO=<ID=MAC,Number=1,Type=Integer,Description="Minor allele count">
##INFO=<ID=MAF,Number=1,Type=Float,Description="Minor allele frequency"><commit_msg>Fix case for no indels and add some safeguards for python2<commit_after>#!/usr/bin/env python
"""Updated version of VCF_from_FASTA.py that hopefully will simplify the
process of generating a VCF from Sanger reads aligned to a reference sequence.
Requires Biopython. Assumes the reference sequence is the first one in the
alignment. Takes one argument:
    1) FASTA multiple sequence alignment"""

import sys
import re
from datetime import date
try:
    from Bio import SeqIO
    fa_in = sys.argv[1]
except ImportError:
    sys.stderr.write('This script requires Biopython.\n')
    sys.exit(1)
except IndexError:
    sys.stderr.write(__doc__ + '\n')
    sys.exit(2)


def extract_ref_coords(alignment, ref_idx=0):
    """Extract the name and the coordinates from the reference sequence of the
    alignment. This expects the sequence to be named in SAM/BED region
    format."""
    ref = alignment[ref_idx]
    chrom = ref.name.split(':')[0]
    start = ref.name.split(':')[1].split('-')[0]
    return (chrom, start)


def check_ref_gaps(alignment, ref_idx=0):
    """Check the reference sequence for end gaps. If we find any end gaps on
    the left, we will throw an error because we cannot accurately calculate the
    positions relative to reference. If we find them on the right, we will
    throw a warning, but continue anyway."""
    ref = alignment[ref_idx]
    refseq = str(ref.seq)
    left_gap = re.compile(r'^-+[ACGTMRWSYKVHDBNacgtmrwsykvhdbn]')
    right_gap = re.compile(r'[ACGTMRWSYKVHDBNacgtmrwsykvhdbn]-+$')
    if left_gap.match(refseq):
        sys.stderr.write(
            """Error:
Reference sequence has end-gaps on the left. This will cause calculated
positions to be incorrect. Please remove the end gaps and re-run this script.
""")
        sys.exit(10)
    if right_gap.search(refseq):
        sys.stderr.write(
            """Warning:
reference sequence has end-gaps on the right. This is not an error, but some of
the variants at the end of the alignment will not be placed on the reference
sequence. You may either remove the right end-gap or remove the variants from
the resulting VCF that occur in the end-gap.\n""")
    return


def extract_variants(alignment, ref_idx=0):
    """Extract the positions of SNPs and indels in the Sanger reads aligned to
    the reference sequence."""
    snp_pos = []
    indel_pos = []
    # First, convert the alignment to a list of lists, as opposed to a list
    # of SeqRecord objects
    raw_aln = [list(s.seq) for s in alignment]
    # Transpose it so that we iterate over columns of the alignment
    t_raw_aln = zip(*raw_aln)
    # Start iterating across columns and saving positions of variant sites. We
    # keep track of the reference base and only increment the position counter
    # for when we see a non-gap character in the reference sequence.
    offset = 0
    for aln_column in t_raw_aln:
        # First, get the states that exist at this position
        states = set(aln_column)
        # Discard any 'N' bases
        states.discard('N')
        # And get the ref state
        ref_state = aln_column[ref_idx]
        # Use the ref state to get the alternate states
        alt_states = states - set(ref_state)
        # If there is a gap in this position, then we will append it to the
        # list of indel positions
        if '-' in states:
            indel_pos.append((offset, ref_state, alt_states))
        # Then, discard the gap character to look for SNPs
        states.discard('-')
        # If the length of the states is greater than 1, then we have a SNP
        if len(states) > 1:
            # We will calculate the following:
            #   Number of non-missing alleles
            #   Minor allele count
            #   Minor allele frequency
            # The reference IS included in these calculations.
            non_missing = [
                base
                for base
                in aln_column
                if base != '-'
                or base != 'N']
            acs = [aln_column.count(x) for x in states]
            afs = [float(c)/len(non_missing) for c in acs]
            # Re-discard the gap character, just to be sure we do not count it
            # as an alternate state
            alt_states.discard('-')
            snp_pos.append(
                (offset, ref_state, alt_states, len(non_missing), min(acs),
                 min(afs)))
        # If the reference sequence is not a gap, then we increment our offset
        # counter.
        if ref_state != '-':
            offset += 1
    return (snp_pos, indel_pos)


def collapse_indels(indels):
    """Collapse indels by identifying runs of consecutive integers and merging
    those into a single entry."""
    # Sort the indel bases by their coordinate
    indel_srt = sorted(indels, key=lambda x: x[0])
    # Make a list to hold our aggregated indels
    agg_indel = []
    # We will now iterate over adjacent records - if they are consecutive, then
    # merge them. Else, break it and start a new record.
    curr_indel = []
    for ind, ind_adj in zip(indel_srt, indel_srt[1:]):
        # Unpack the alleles. It's a little silly, but we have to cast the set
        # to a list to subset it.
        curr_ref = ind[1]
        curr_alt = list(ind[2])[0]
        adj_ref = ind_adj[1]
        adj_alt = list(ind_adj[2])[0]
        if not curr_indel:
            curr_indel = [ind[0], curr_ref, curr_alt]
        # If the next position is not consecutive, append it and start over
        if ind_adj[0] - ind[0] > 1:
            agg_indel.append(curr_indel)
            curr_indel = []
        else:
            curr_indel[1] += adj_ref
            curr_indel[2] += adj_alt
    # The way we are iterating through the indel list means that we will always
    # leave off the last one. Append it after the loop finishes, but only if
    # the current indel exists
    if curr_indel:
        agg_indel.append(curr_indel)
    return agg_indel


def adjust_indels(indels, alignment, ref_idx=0):
    """Adjust the indel positions so that they are offset by one, as required
    by the VCF spec. This is because the reported position must be the base
    *before* any insertion/deletion polymorphism is observed."""
    spec_indels = []
    # Remove the gaps from the reference sequnce for getting the reference base
    # of the indel
    ref_seq = ''.join([base for base in alignment[ref_idx].seq if base != '-'])
    for i in indels:
        spec_pos = i[0] - 1
        spec_ref = ref_seq[spec_pos]
        spec_indel = [spec_pos, spec_ref + i[1], spec_ref + i[2]]
        spec_indels.append(spec_indel)
    return spec_indels


def print_vcf(snp_var, ind_var, refseq, offset):
    """Print a VCF from the calculated positions of the variants."""
    # Define the VCF header
    today = date.today().strftime('%Y%m%d')
    vcf_header = """##fileformat=VCFv4.1
##fileDate={filedate}
##source=VCF_from_FASTA_2.py;Morrell Lab @ UMN
##INFO=<ID=MAC,Number=1,Type=Integer,Description="Minor allele count">
##INFO=<ID=MAF,Number=1,Type=Float,Description="Minor allele frequency">
{'fraction_non_alphanumeric': 0.08480615735461802, 'fraction_numerical': 0.005772519954389966, 'mean_word_length': 2.62422520661157, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3580810', 'n_tokens_mistral': 4291, 'n_tokens_neox': 4148, 'n_words': 1602}