Generate consensus fasta from bam file

bam2consensus.py

from collections import Counter, OrderedDict
import pysam
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord

def bam2consensus(
        ref_fname, bam_fname, ac_threshold, af_threshold, out_fname):

    consensus = OrderedDict()
    for record in SeqIO.parse(ref_fname, "fasta"):
        consensus[record.id] = ["N"] * len(record)

    with pysam.AlignmentFile(bam_fname, "rb") as bam:
        allele_counter = Counter()
        for pileup_column in bam.pileup():
            chrom = pileup_column.reference_name
            pos = pileup_column.reference_pos
            assert consensus[chrom][pos] == "N"

            allele_counter.clear()
            for pileup_read in pileup_column.pileups:
                if pileup_read.is_del:
                    allele = "-"
                else:
                    allele = pileup_read.alignment.query_sequence[
                        pileup_read.query_position]
                allele_counter[allele] += 1

            max_allele = "N"
            max_count, total_count = 0, 0
            for allele, count in allele_counter.items():
                if count > max_count:
                    max_count = count
                    max_allele = allele
                total_count += count

            assert max_allele in "ACGTN-"
            if (max_count >= ac_threshold and
                max_count / total_count >= af_threshold):
                consensus[chrom][pos] = max_allele

    records = []
    for chrom, seq in consensus.items():
        records.append(
            SeqRecord(Seq("".join(seq)), id=chrom, description=""))

    SeqIO.write(records, out_fname, "fasta")