Calculate number of minor alleles (not in consensus sequence).

naive_nshared.py

import sys
from readfq import readfq
from itertools import combinations
from datetime import datetime

def num_shared(seq_a, seq_b, consensus_seq):
    """
    Given two alignment sequences in multiple alignment FASTA format,
    calculate the number of shared SNPs (for minor alleles only, not 
    in consensus).
    """
    assert len(seq_a) == len(seq_b) == len(consensus_seq)
    # return sum([seq_a[i] == seq_b[i] for i in
    #             range(len(seq_a)) if seq_a[i] != "-" and seq_b[i] != "-"])
    nshared = 0
    for i in xrange(len(seq_a)):
        if "-" in (seq_a[i], seq_b[i], consensus_seq[i]):
           continue
        nshared += (seq_a[i] == seq_b[i]) and seq_a[i] != consensus_seq[i]
    return nshared

if __name__ == "__main__":
    if len(sys.argv) < 2:
        sys.exit("usage: consensus.fa <alignments.fa>")

    n = 100000 # increment for timing messages 
    alns = dict()
    with open(sys.argv[1]) as consensus_file:
        _, consensus_seq, _ = next(readfq(consensus_file))
    
    if len(sys.argv) > 2 and sys.argv[2] != "-":
        aln_stream = open(sys.argv[2])
    else:
        aln_stream = sys.stdin
    for name, seq, _ in readfq(aln_stream):
        alns[name] = seq
    
    results = dict()
    i = 0
    time = datetime.now()
    for aln_a, aln_b in combinations(alns, 2):
        pair_key = (aln_a, aln_b)
        seq_a, seq_b = alns[aln_a], alns[aln_b]
        num_ab_share = num_shared(seq_a, seq_b, consensus_seq)
        #results[pair_key] = num_ab_share
        if i > 0 and i % n == 0:
            elapsed = (datetime.now() - time).total_seconds()
            sys.stderr.write("\ttotal pairs counted: %d, %d in %f seconds\r" % (i, n, elapsed))
            time = datetime.now()
        i += 1
        print "\t".join([aln_a, aln_b, str(num_ab_share)])