initial commit v0.1.2

import os

import pysam

import argparse

import multiprocessing as mp

import regex

#from Bio import SeqIO

#make the lookup dictionary for reverse complementing

complement = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A', 'N': 'N'}

def reverse_complement(seq):

    bases = list(seq)

    bases = reversed([complement.get(base,base) for base in bases])

    bases = ''.join(bases)

    return bases

def complementseq(seq):

    bases = list(seq)

    bases = [complement.get(base,base) for base in bases]

    bases = ''.join(bases)

    return bases

def idx_bam(bam, threads):

    threads = str(threads)

    try:

        pysam.index("-@"+threads,bam)

    except:

        outcome = 'idxerror'

    else:

        outcome = 'idxsuccess'

    if outcome == 'idxerror':

        print("indexing failed, trying to sort bam file...")

        inbam = bam

        bam = bam+".sorted.bam"

        pysam.sort("-@"+threads,"-o", bam, inbam)

        print("indexing bam file...")

        pysam.index(bam)

    return bam

def makeBAMheader(args, v):

    bam = pysam.AlignmentFile(args.bam, 'rb')

    hdr = bam.header.to_dict()

    bam.close()

    cmdlinecall = 'pyTSOfilter --bam '+args.bam+' --out '+args.out+' --fa '+args.fa+' --p '+str(args.p)+' --n_mismatch '+str(args.n_mismatch)

    if args.ggg:

        cmdlinecall = cmdlinecall+' --ggg'

    pg = {'ID': 'pyTSOfilter', 'PN': 'pyTSOfilter',

          'CL': cmdlinecall, 'VN': v}

    if 'PG' in hdr:

        pglines = hdr['PG']

        pglines.append(pg)

    else:

        pglines = [pg]

    hdr['PG'] = pglines

    return hdr

def is_fiveprime_read(read):

    fivep = False

    if read.has_tag('UB'):

        UMIseq = read.get_tag('UB')

        if UMIseq != "":

            fivep = True

    return fivep

def get_gene(read):

    if read.has_tag("GE"):

        gene = read.get_tag("GE")

    else:

        if read.has_tag("GI"):

            gene = read.get_tag("GI")

        else:

            gene = "none"

    return gene

def collect_bam_chunks(inpath, chrs, outpath):

    allpaths = [inpath+".tmp."+c+".bam" for c in chrs[:-1]]

    allpaths.append(inpath+".tmp."+"unmapped"+".bam")

    cat_args = ['-o', outpath]+allpaths

    pysam.cat(*cat_args)

    x = [os.remove(f) for f in allpaths]

def filterTSO(inpath, threads, fastapath, chr, mm, ggg, outheader):

    fa = pysam.FastaFile(fastapath)

    if chr == '*':

        chrlabel = 'unmapped'

    else:

        chrlabel = chr

        #reflen = fa.get_reference_length(chr)

    #open in/out files

    outpath = inpath+".tmp."+chrlabel+".bam"

    inp = pysam.AlignmentFile(inpath, 'rb', threads = threads)

    out = pysam.AlignmentFile(outpath, 'wb', header = outheader, threads = threads)

    for read in inp.fetch(chr):

        isfivep = is_fiveprime_read(read)

        gene = get_gene(read)

        if chr in fa.references and isfivep and gene != "none":

            if read.is_reverse:

                strand = "-"

                rloffset = read.query_length

            else:

                strand = "+"

            if read.is_paired:

                readtype = "PE"

                if read.is_read1:

                    startpos = read.reference_start

                else:

                    if read.mate_is_unmapped:

                        continue

                    startpos = read.next_reference_start

                    if read.mate_is_reverse:

                        strand = "-"

                        rloffset = read.query_length-22

                    else:

                        strand = "+"

            else:

                startpos = read.reference_start

                readtype = "SE"

            UMIseq = read.get_tag('UB')

            if ggg:

                UMIseq = UMIseq+"GGG"

            if strand == "-":

                #if(startpos+rloffset+21 > reflen):

                #    upstream_seq = fa.fetch(chr, startpos+rloffset, reflen)

                #else:

                #    upstream_seq = fa.fetch(chr, startpos+rloffset, startpos+rloffset+21)

                #if the fetch is longer than reflen, it just returns empty string but no error

                upstream_seq = fa.fetch(chr, startpos+rloffset, startpos+rloffset+21)

                upstream_seq = reverse_complement(upstream_seq)

            else:

                if startpos-20 < 0:

                    startpos = startpos + 20-startpos

                upstream_seq = fa.fetch(chr, startpos-20, startpos+1)

            pattern = r"(" + UMIseq+ "){s<="+ str(mm) + "}"

            m = regex.search(pattern, upstream_seq)

            if m is not None:

                continue

        out.write(read)

    inp.close()

    out.close()

def main():

    parser = argparse.ArgumentParser(add_help=True)

    parser.add_argument('--bam', type=str, metavar='FILENAME',

                        help='Path to input BAM file', required = True)

    parser.add_argument('--out', type=str, metavar='FILENAME',

                        help='Path to output bam file', required = True)

    parser.add_argument('--fa', type=str, metavar='FILENAME',

                        help='Path to genome reference fasta', required = True)

    parser.add_argument('--p', type=int, default = 10,

                        help='Number of processes to use')

    parser.add_argument('--n_mismatch', type=int, default = 1,

                        help='Number of mismatches allowed when matching TSO')

    parser.add_argument('--ggg', action = 'store_true',

                        help='Require template-switching Gs as part of sequence match')

    args = parser.parse_args()

    v = '0.1.2'

    print("pyTSOfilter.py v"+v)

    print("WARNING: This script only works correctly when using stranded gene assignment via zUMIs.")

    bampath = args.bam

    print("Loading reference sequence...")

    #reffa = '/home/chris/projects/pyTSOfilter/hg38.primary_assembly.fa'

    reffa = args.fa

    #fa_dict = SeqIO.to_dict(SeqIO.parse(reffa, "fasta"))

    #check if fasta file is indexed

    try:

        fa = pysam.FastaFile(reffa)

    except ValueError:

        print("Error: Reference fasta file is not indexed!")

        print("Please run: samtools faidx "+args.fa)

        quit()

    fa.close()

    if not os.path.exists(bampath+'.bai'):

        print("input bam index not found, indexing...")

        bampath = idx_bam(bampath,args.p)

    chrs = pysam.idxstats(bampath).split('\n')

    chrs = [c.split('\t')[0] for c in chrs[:-1]]

    #chrs.append('*')

    if args.p > 20:

        pysam_workers = 2

        n_jobs = int(args.p/2)

    else:

        pysam_workers = 1

        n_jobs = args.p

    #Construct the new bam header to work with

    bamheader = makeBAMheader(args, v)

    print("Working...")

    pool = mp.Pool(n_jobs)

    results = [pool.apply_async(filterTSO, (args.bam,pysam_workers,args.fa,chr,args.n_mismatch,args.ggg,bamheader,  )) for chr in chrs]

    x = [r.get() for r in results]

    print("Creating final output .bam file...")

    collect_bam_chunks(inpath = bampath, chrs = chrs, outpath = args.out)

    print("Indexing final output .bam file...")

    y = idx_bam(args.out,args.p)

    print("Done!")

if __name__ == "__main__":

    main()