ROSE_bamToGFF.py

#!/usr/bin/env python3
#bamToGFF.py

#script to grab reads from a bam that align to a .gff file
import sys
import re

import ROSE_utils


from collections import defaultdict

import os

import string



#=====================================================================
#====================MAPPING BAM READS TO GFF REGIONS=================
#=====================================================================


def mapBamToGFF(bamFile,gff,sense = 'both',extension = 200,floor = 0,rpm = False,matrix = None):

#def mapBamToGFF(bamFile,gff,sense = 'both',unique = 0,extension = 200,floor = 0,density = False,rpm = False,binSize = 25,clusterGram = None,matrix = None,raw = False,includeJxnReads = False):
    '''maps reads from a bam to a gff'''
    floor = int(floor)
    
    #USING BAM CLASS
    bam = ROSE_utils.Bam(bamFile)


    #new GFF to write to
    newGFF = []
    #millionMappedReads


    if rpm:    
        MMR= round(float(bam.getTotalReads('mapped'))/1000000,4)
    else:
        MMR = 1

    print(('using a MMR value of %s' % (MMR)))
    
    #senseTrans = maketrans('-+.','+-+') #deprecated

    if ROSE_utils.checkChrStatus(bamFile) == 1:
      print("has chr")
      hasChrFlag = 1
      #sys.exit();
    else:
      print("does not have chr")
      hasChrFlag = 0
      #sys.exit()
      
    if type(gff) == str:
        gff = ROSE_utils.parseTable(gff,'\t')
        
    #setting up a maxtrix table

    newGFF.append(['GENE_ID','locusLine'] + ['bin_'+str(n)+'_'+bamFile.split('/')[-1] for n in range(1,int(matrix)+1,1)])        

    #getting and processing reads for gff lines
    ticker = 0
    print('Number lines processed')
    for line in gff:
        line = line[0:9]
        if ticker%100 == 0:
            print(ticker)
        ticker+=1
        if not hasChrFlag:
            line[0] = re.sub(r"chr",r"",line[0])
        gffLocus = ROSE_utils.Locus(line[0],int(line[3]),int(line[4]),line[6],line[1])
        #print line[0]
        #sys.exit()
        searchLocus = ROSE_utils.makeSearchLocus(gffLocus,int(extension),int(extension))
        
        reads = bam.getReadsLocus(searchLocus,'both',False,'none')
        #now extend the reads and make a list of extended reads
        extendedReads = []
        for locus in reads:
            if locus.sense() == '+' or locus.sense() == '.':
                locus = ROSE_utils.Locus(locus.chr(),locus.start(),locus.end()+extension,locus.sense(), locus.ID())
            if locus.sense() == '-':
                locus = ROSE_utils.Locus(locus.chr(),locus.start()-extension,locus.end(),locus.sense(),locus.ID())
            extendedReads.append(locus)
        if gffLocus.sense() == '+' or gffLocus.sense == '.':
            senseReads = [x for x in extendedReads if x.sense() == '+' or x.sense() == '.']
            antiReads = [x for x in extendedReads if x.sense() == '-']
        else:
            senseReads = [x for x in extendedReads if x.sense() == '-' or x.sense() == '.']
            antiReads = [x for x in extendedReads if x.sense() == '+']

        senseHash = defaultdict(int)
        antiHash = defaultdict(int)

        #filling in the readHashes             
        if sense == '+' or sense == 'both' or sense =='.':
            for read in senseReads:
                for x in range(read.start(),read.end()+1,1):
                    senseHash[x]+=1
        if sense == '-' or sense == 'both' or sense == '.':
            #print('foo')
            for read in antiReads:
                for x in range(read.start(),read.end()+1,1):
                    antiHash[x]+=1

        #now apply flooring and filtering for coordinates
        keys = ROSE_utils.uniquify(list(senseHash.keys())+list(antiHash.keys()))
        if floor > 0:

            keys = [x for x in keys if (senseHash[x]+antiHash[x]) > floor]
        #coordinate filtering
        keys = [x for x in keys if gffLocus.start() < x < gffLocus.end()]


        #setting up the output table
        #clusterLine = [gffLocus.ID(),gffLocus.__str__()]

        # bug fix gff coordinates with same chromosomal name as BAM
        if not hasChrFlag:
            clusterLine = [gffLocus.ID(),"chr" + gffLocus.__str__()]
        else: 
           clusterLine = [gffLocus.ID(),gffLocus.__str__()]

        #getting the binsize
        binSize = (gffLocus.len()-1)/int(matrix)
        nBins = int(matrix)
        if binSize == 0:
            clusterLine+=['NA']*int(matrix)
            newGFF.append(clusterLine)
            continue
        n=0
        if gffLocus.sense() == '+' or gffLocus.sense() =='.' or gffLocus.sense() == 'both':
            i = gffLocus.start()

            while n <nBins:
                n+=1
                binKeys = [x for x in keys if i < x < i+binSize]
                binDen = float(sum([senseHash[x]+antiHash[x] for x in binKeys]))/binSize
                clusterLine+=[round(binDen/MMR,4)]
                i = i+binSize
        else:
            i = gffLocus.end()
            while n < nBins:
                n+=1
                binKeys = [x for x in keys if i-binSize < x < i]
                binDen = float(sum([senseHash[x]+antiHash[x] for x in binKeys]))/binSize
                clusterLine+=[round(binDen/MMR,4)]
                i = i-binSize
        newGFF.append(clusterLine)
        
            
    return newGFF
        
                
                
#=====================================================================
#============================MAIN METHOD==============================
#=====================================================================
        

def main():
    from optparse import OptionParser
    usage = "usage: %prog [options] -b [SORTED BAMFILE] -i [INPUTFILE] -o [OUTPUTFILE]"
    parser = OptionParser(usage = usage)
    #required flags
    parser.add_option("-b","--bam", dest="bam",nargs = 1, default=None,
                      help = "Enter .bam file to be processed.")
    parser.add_option("-i","--input", dest="input",nargs = 1, default=None,
                      help = "Enter .gff or ENRICHED REGION file to be processed.")
    #output flag
    parser.add_option("-o","--output", dest="output",nargs = 1, default=None,
                      help = "Enter the output filename.")
    #additional options
    parser.add_option("-s","--sense", dest="sense",nargs = 1, default='both',
                      help = "Map to '+','-' or 'both' strands. Default maps to both.")


    parser.add_option("-f","--floor", dest="floor",nargs =1, default=0,
                      help = "Sets a read floor threshold necessary to count towards density")    
    parser.add_option("-e","--extension", dest="extension",nargs = 1, default=200,
                      help = "Extends reads by n bp. Default value is 200bp")
    parser.add_option("-r","--rpm", dest="rpm",action = 'store_true', default=False,
                      help = "Normalizes density to reads per million (rpm)")


    parser.add_option("-m","--matrix", dest="matrix",nargs = 1, default=None,
                      help = "Outputs a variable bin sized matrix. User must specify number of bins.")

    (options,args) = parser.parse_args()

    print(options)
    print(args)

    if options.bam:
        bamFile = options.bam
        fullPath = os.path.abspath(bamFile)
        bamName = fullPath.split('/')[-1].split('.')[0]
        pathFolder = '/'.join(fullPath.split('/')[0:-1])
        fileList = os.listdir(pathFolder)
        hasBai = False
        for fileName in fileList:
            if fileName.count(bamName) == 1 and fileName.count('.bai') == 1:
                hasBai = True

        if not hasBai:
            print('ERROR: no associated .bai file found with bam. Must use a sorted bam with accompanying index file')
            parser.print_help()
            exit()
   
    if options.sense:
        if ['+','-','.','both'].count(options.sense) == 0:
            print('ERROR: sense flag must be followed by +,-,.,both')
            parser.print_help()
            exit()


    if options.matrix:
        try:
            int(options.matrix)
        except:
            print('ERROR: User must specify an integer bin number for matrix (try 50)')
            parser.print_help()
            exit()
            

    
    
    if options.input and options.bam:
        inputFile = options.input
        gffFile = inputFile

        bamFile = options.bam
        
        if options.output == None:
            output = os.getcwd() + inputFile.split('/')[-1]+'.mapped'
        else:
            output = options.output
        if options.matrix:
            print('mapping to GFF and making a matrix with fixed bin number')

            newGFF = mapBamToGFF(bamFile,gffFile,options.sense,int(options.extension),options.floor,options.rpm,options.matrix)

            
        ROSE_utils.unParseTable(newGFF,output,'\t')
    else:
        parser.print_help()
                
 
if __name__ == "__main__":
    main()