PageRenderTime 47ms CodeModel.GetById 14ms app.highlight 28ms RepoModel.GetById 1ms app.codeStats 0ms

/tools/sr_mapping/bwa_wrapper.py

https://bitbucket.org/cistrome/cistrome-harvard/
Python | 342 lines | 325 code | 5 blank | 12 comment | 26 complexity | 3a9aba466fb4a8a8029ac3c4fcccf412 MD5 | raw file
  1#!/usr/bin/env python
  2
  3"""
  4Runs BWA on single-end or paired-end data.
  5Produces a SAM file containing the mappings.
  6Works with BWA version 0.5.9.
  7
  8usage: bwa_wrapper.py [options]
  9
 10See below for options
 11"""
 12
 13import optparse, os, shutil, subprocess, sys, tempfile
 14
 15def stop_err( msg ):
 16    sys.stderr.write( '%s\n' % msg )
 17    sys.exit()
 18
 19def check_is_double_encoded( fastq ):
 20    # check that first read is bases, not one base followed by numbers
 21    bases = [ 'A', 'C', 'G', 'T', 'a', 'c', 'g', 't', 'N' ]
 22    nums = [ '0', '1', '2', '3' ]
 23    for line in file( fastq, 'rb'):
 24        if not line.strip() or line.startswith( '@' ):
 25            continue
 26        if len( [ b for b in line.strip() if b in nums ] ) > 0:
 27            return False
 28        elif line.strip()[0] in bases and len( [ b for b in line.strip() if b in bases ] ) == len( line.strip() ):
 29            return True
 30        else:
 31            raise Exception, 'First line in first read does not appear to be a valid FASTQ read in either base-space or color-space'
 32    raise Exception, 'There is no non-comment and non-blank line in your FASTQ file'
 33
 34def __main__():
 35    #Parse Command Line
 36    parser = optparse.OptionParser()
 37    parser.add_option( '-t', '--threads', dest='threads', help='The number of threads to use' )
 38    parser.add_option( '-c', '--color-space', dest='color_space', action='store_true', help='If the input files are SOLiD format' )
 39    parser.add_option( '-r', '--ref', dest='ref', help='The reference genome to use or index' )
 40    parser.add_option( '-f', '--input1', dest='fastq', help='The (forward) fastq file to use for the mapping' )
 41    parser.add_option( '-F', '--input2', dest='rfastq', help='The reverse fastq file to use for mapping if paired-end data' )
 42    parser.add_option( '-u', '--output', dest='output', help='The file to save the output (SAM format)' )
 43    parser.add_option( '-g', '--genAlignType', dest='genAlignType', help='The type of pairing (single or paired)' )
 44    parser.add_option( '-p', '--params', dest='params', help='Parameter setting to use (pre_set or full)' )
 45    parser.add_option( '-s', '--fileSource', dest='fileSource', help='Whether to use a previously indexed reference sequence or one form history (indexed or history)' )
 46    parser.add_option( '-n', '--maxEditDist', dest='maxEditDist', help='Maximum edit distance if integer' )
 47    parser.add_option( '-m', '--fracMissingAligns', dest='fracMissingAligns', help='Fraction of missing alignments given 2% uniform base error rate if fraction' )
 48    parser.add_option( '-o', '--maxGapOpens', dest='maxGapOpens', help='Maximum number of gap opens' )
 49    parser.add_option( '-e', '--maxGapExtens', dest='maxGapExtens', help='Maximum number of gap extensions' )
 50    parser.add_option( '-d', '--disallowLongDel', dest='disallowLongDel', help='Disallow a long deletion within specified bps' )
 51    parser.add_option( '-i', '--disallowIndel', dest='disallowIndel', help='Disallow indel within specified bps' )
 52    parser.add_option( '-l', '--seed', dest='seed', help='Take the first specified subsequences' )
 53    parser.add_option( '-k', '--maxEditDistSeed', dest='maxEditDistSeed', help='Maximum edit distance to the seed' )
 54    parser.add_option( '-M', '--mismatchPenalty', dest='mismatchPenalty', help='Mismatch penalty' )
 55    parser.add_option( '-O', '--gapOpenPenalty', dest='gapOpenPenalty', help='Gap open penalty' )
 56    parser.add_option( '-E', '--gapExtensPenalty', dest='gapExtensPenalty', help='Gap extension penalty' )
 57    parser.add_option( '-R', '--suboptAlign', dest='suboptAlign', help='Proceed with suboptimal alignments even if the top hit is a repeat' )
 58    parser.add_option( '-N', '--noIterSearch', dest='noIterSearch', help='Disable iterative search' )
 59    parser.add_option( '-T', '--outputTopN', dest='outputTopN', help='Maximum number of alignments to output in the XA tag for reads paired properly' )
 60    parser.add_option( '', '--outputTopNDisc', dest='outputTopNDisc', help='Maximum number of alignments to output in the XA tag for disconcordant read pairs (excluding singletons)' )
 61    parser.add_option( '-S', '--maxInsertSize', dest='maxInsertSize', help='Maximum insert size for a read pair to be considered mapped good' )
 62    parser.add_option( '-P', '--maxOccurPairing', dest='maxOccurPairing', help='Maximum occurrences of a read for pairings' )
 63    parser.add_option( '', '--rgid', dest='rgid', help='Read group identifier' )
 64    parser.add_option( '', '--rgcn', dest='rgcn', help='Sequencing center that produced the read' )
 65    parser.add_option( '', '--rgds', dest='rgds', help='Description' )
 66    parser.add_option( '', '--rgdt', dest='rgdt', help='Date that run was produced (ISO8601 format date or date/time, like YYYY-MM-DD)' )
 67    parser.add_option( '', '--rgfo', dest='rgfo', help='Flow order' )
 68    parser.add_option( '', '--rgks', dest='rgks', help='The array of nucleotide bases that correspond to the key sequence of each read' )
 69    parser.add_option( '', '--rglb', dest='rglb', help='Library name' )
 70    parser.add_option( '', '--rgpg', dest='rgpg', help='Programs used for processing the read group' )
 71    parser.add_option( '', '--rgpi', dest='rgpi', help='Predicted median insert size' )
 72    parser.add_option( '', '--rgpl', dest='rgpl', choices=[ 'CAPILLARY', 'LS454', 'ILLUMINA', 'SOLID', 'HELICOS', 'IONTORRENT' and 'PACBIO' ], help='Platform/technology used to produce the reads' )
 73    parser.add_option( '', '--rgpu', dest='rgpu', help='Platform unit (e.g. flowcell-barcode.lane for Illumina or slide for SOLiD)' )
 74    parser.add_option( '', '--rgsm', dest='rgsm', help='Sample' )
 75    parser.add_option( '-D', '--dbkey', dest='dbkey', help='Dbkey for reference genome' )
 76    parser.add_option( '-X', '--do_not_build_index', dest='do_not_build_index', action='store_true', help="Don't build index" )
 77    parser.add_option( '-H', '--suppressHeader', dest='suppressHeader', help='Suppress header' )
 78    parser.add_option( '-I', '--illumina1.3', dest='illumina13qual', help='Input FASTQ files have Illuina 1.3 quality scores' )
 79    (options, args) = parser.parse_args()
 80
 81    # output version # of tool
 82    try:
 83        tmp = tempfile.NamedTemporaryFile().name
 84        tmp_stdout = open( tmp, 'wb' )
 85        proc = subprocess.Popen( args='bwa 2>&1', shell=True, stdout=tmp_stdout )
 86        tmp_stdout.close()
 87        returncode = proc.wait()
 88        stdout = None
 89        for line in open( tmp_stdout.name, 'rb' ):
 90            if line.lower().find( 'version' ) >= 0:
 91                stdout = line.strip()
 92                break
 93        if stdout:
 94            sys.stdout.write( 'BWA %s\n' % stdout )
 95        else:
 96            raise Exception
 97    except:
 98        sys.stdout.write( 'Could not determine BWA version\n' )
 99
100    # check for color space fastq that's not double-encoded and exit if appropriate
101    if options.color_space:
102        if not check_is_double_encoded( options.fastq ):
103            stop_err( 'Your file must be double-encoded (it must be converted from "numbers" to "bases"). See the help section for details' )
104        if options.genAlignType == 'paired':
105            if not check_is_double_encoded( options.rfastq ):
106                stop_err( 'Your reverse reads file must also be double-encoded (it must be converted from "numbers" to "bases"). See the help section for details' )
107
108    fastq = options.fastq
109    if options.rfastq:
110         rfastq = options.rfastq
111
112    # set color space variable
113    if options.color_space:
114        color_space = '-c'
115    else:
116        color_space = ''
117
118    # make temp directory for placement of indices
119    tmp_index_dir = tempfile.mkdtemp()
120    tmp_dir = tempfile.mkdtemp()
121    # index if necessary
122    if options.fileSource == 'history' and not options.do_not_build_index:
123        ref_file = tempfile.NamedTemporaryFile( dir=tmp_index_dir )
124        ref_file_name = ref_file.name
125        ref_file.close()
126        os.symlink( options.ref, ref_file_name )
127        # determine which indexing algorithm to use, based on size
128        try:
129            size = os.stat( options.ref ).st_size
130            if size <= 2**30: 
131                indexingAlg = 'is'
132            else:
133                indexingAlg = 'bwtsw'
134        except:
135            indexingAlg = 'is'
136        indexing_cmds = '%s -a %s' % ( color_space, indexingAlg )
137        cmd1 = 'bwa index %s %s' % ( indexing_cmds, ref_file_name )
138        try:
139            tmp = tempfile.NamedTemporaryFile( dir=tmp_index_dir ).name
140            tmp_stderr = open( tmp, 'wb' )
141            proc = subprocess.Popen( args=cmd1, shell=True, cwd=tmp_index_dir, stderr=tmp_stderr.fileno() )
142            returncode = proc.wait()
143            tmp_stderr.close()
144            # get stderr, allowing for case where it's very large
145            tmp_stderr = open( tmp, 'rb' )
146            stderr = ''
147            buffsize = 1048576
148            try:
149                while True:
150                    stderr += tmp_stderr.read( buffsize )
151                    if not stderr or len( stderr ) % buffsize != 0:
152                        break
153            except OverflowError:
154                pass
155            tmp_stderr.close()
156            if returncode != 0:
157                raise Exception, stderr
158        except Exception, e:
159            # clean up temp dirs
160            if os.path.exists( tmp_index_dir ):
161                shutil.rmtree( tmp_index_dir )
162            if os.path.exists( tmp_dir ):
163                shutil.rmtree( tmp_dir )
164            stop_err( 'Error indexing reference sequence. ' + str( e ) )
165    else:
166        ref_file_name = options.ref
167    if options.illumina13qual:
168        illumina_quals = "-I"
169    else:
170        illumina_quals = ""
171
172    # set up aligning and generate aligning command options
173    if options.params == 'pre_set':
174        aligning_cmds = '-t %s %s %s' % ( options.threads, color_space, illumina_quals )
175        gen_alignment_cmds = ''
176    else:
177        if options.maxEditDist != '0':
178            editDist = options.maxEditDist
179        else:
180            editDist = options.fracMissingAligns
181        if options.seed != '-1':
182            seed = '-l %s' % options.seed
183        else:
184            seed = ''
185        if options.suboptAlign == 'true':
186            suboptAlign = '-R'
187        else:
188            suboptAlign = ''
189        if options.noIterSearch == 'true':
190            noIterSearch = '-N'
191        else:
192            noIterSearch = ''
193        aligning_cmds = '-n %s -o %s -e %s -d %s -i %s %s -k %s -t %s -M %s -O %s -E %s %s %s %s %s' % \
194                        ( editDist, options.maxGapOpens, options.maxGapExtens, options.disallowLongDel,
195                          options.disallowIndel, seed, options.maxEditDistSeed, options.threads,
196                          options.mismatchPenalty, options.gapOpenPenalty, options.gapExtensPenalty,
197                          suboptAlign, noIterSearch, color_space, illumina_quals )
198        if options.genAlignType == 'paired':
199            gen_alignment_cmds = '-a %s -o %s' % ( options.maxInsertSize, options.maxOccurPairing )
200            if options.outputTopNDisc:
201                gen_alignment_cmds += ' -N %s' % options.outputTopNDisc
202        else:
203            gen_alignment_cmds = ''
204        if options.rgid:
205            if not options.rglb or not options.rgpl or not options.rgsm:
206                stop_err( 'If you want to specify read groups, you must include the ID, LB, PL, and SM tags.' )
207            readGroup = '@RG\tID:%s\tLB:%s\tPL:%s\tSM:%s' % ( options.rgid, options.rglb, options.rgpl, options.rgsm )
208            if options.rgcn:
209                readGroup += '\tCN:%s' % options.rgcn
210            if options.rgds:
211                readGroup += '\tDS:%s' % options.rgds
212            if options.rgdt:
213                readGroup += '\tDT:%s' % options.rgdt
214            if options.rgfo:
215                readGroup += '\tFO:%s' % options.rgfo
216            if options.rgks:
217                readGroup += '\tKS:%s' % options.rgks
218            if options.rgpg:
219                readGroup += '\tPG:%s' % options.rgpg
220            if options.rgpi:
221                readGroup += '\tPI:%s' % options.rgpi
222            if options.rgpu:
223                readGroup += '\tPU:%s' % options.rgpu
224            gen_alignment_cmds += ' -r "%s"' % readGroup
225        if options.outputTopN:
226            gen_alignment_cmds += ' -n %s' % options.outputTopN
227    # set up output files
228    tmp_align_out = tempfile.NamedTemporaryFile( dir=tmp_dir )
229    tmp_align_out_name = tmp_align_out.name
230    tmp_align_out.close()
231    tmp_align_out2 = tempfile.NamedTemporaryFile( dir=tmp_dir )
232    tmp_align_out2_name = tmp_align_out2.name
233    tmp_align_out2.close()
234    # prepare actual aligning and generate aligning commands
235    cmd2 = 'bwa aln %s %s %s > %s' % ( aligning_cmds, ref_file_name, fastq, tmp_align_out_name )
236    cmd2b = ''
237    if options.genAlignType == 'paired':
238        cmd2b = 'bwa aln %s %s %s > %s' % ( aligning_cmds, ref_file_name, rfastq, tmp_align_out2_name )
239        cmd3 = 'bwa sampe %s %s %s %s %s %s >> %s' % ( gen_alignment_cmds, ref_file_name, tmp_align_out_name, tmp_align_out2_name, fastq, rfastq, options.output )
240    else:
241        cmd3 = 'bwa samse %s %s %s %s >> %s' % ( gen_alignment_cmds, ref_file_name, tmp_align_out_name, fastq, options.output )
242    # perform alignments
243    buffsize = 1048576
244    try:
245        # need to nest try-except in try-finally to handle 2.4
246        try:
247            # align
248            try:
249                tmp = tempfile.NamedTemporaryFile( dir=tmp_dir ).name
250                tmp_stderr = open( tmp, 'wb' )
251                proc = subprocess.Popen( args=cmd2, shell=True, cwd=tmp_dir, stderr=tmp_stderr.fileno() )
252                returncode = proc.wait()
253                tmp_stderr.close()
254                # get stderr, allowing for case where it's very large
255                tmp_stderr = open( tmp, 'rb' )
256                stderr = ''
257                try:
258                    while True:
259                        stderr += tmp_stderr.read( buffsize )
260                        if not stderr or len( stderr ) % buffsize != 0:
261                            break
262                except OverflowError:
263                    pass
264                tmp_stderr.close()
265                if returncode != 0:
266                    raise Exception, stderr
267            except Exception, e:
268                raise Exception, 'Error aligning sequence. ' + str( e )
269            # and again if paired data
270            try:
271                if cmd2b:
272                    tmp = tempfile.NamedTemporaryFile( dir=tmp_dir ).name
273                    tmp_stderr = open( tmp, 'wb' )
274                    proc = subprocess.Popen( args=cmd2b, shell=True, cwd=tmp_dir, stderr=tmp_stderr.fileno() )
275                    returncode = proc.wait()
276                    tmp_stderr.close()
277                    # get stderr, allowing for case where it's very large
278                    tmp_stderr = open( tmp, 'rb' )
279                    stderr = ''
280                    try:
281                        while True:
282                            stderr += tmp_stderr.read( buffsize )
283                            if not stderr or len( stderr ) % buffsize != 0:
284                                break
285                    except OverflowError:
286                        pass
287                    tmp_stderr.close()
288                    if returncode != 0:
289                        raise Exception, stderr
290            except Exception, e:
291                raise Exception, 'Error aligning second sequence. ' + str( e )
292            # generate align
293            try:
294                tmp = tempfile.NamedTemporaryFile( dir=tmp_dir ).name
295                tmp_stderr = open( tmp, 'wb' )
296                proc = subprocess.Popen( args=cmd3, shell=True, cwd=tmp_dir, stderr=tmp_stderr.fileno() )
297                returncode = proc.wait()
298                tmp_stderr.close()
299                # get stderr, allowing for case where it's very large
300                tmp_stderr = open( tmp, 'rb' )
301                stderr = ''
302                try:
303                    while True:
304                        stderr += tmp_stderr.read( buffsize )
305                        if not stderr or len( stderr ) % buffsize != 0:
306                            break
307                except OverflowError:
308                    pass
309                tmp_stderr.close()
310                if returncode != 0:
311                    raise Exception, stderr
312            except Exception, e:
313                raise Exception, 'Error generating alignments. ' + str( e ) 
314            # remove header if necessary
315            if options.suppressHeader == 'true':
316                tmp_out = tempfile.NamedTemporaryFile( dir=tmp_dir)
317                tmp_out_name = tmp_out.name
318                tmp_out.close()
319                try:
320                    shutil.move( options.output, tmp_out_name )
321                except Exception, e:
322                    raise Exception, 'Error moving output file before removing headers. ' + str( e )
323                fout = file( options.output, 'w' )
324                for line in file( tmp_out.name, 'r' ):
325                    if not ( line.startswith( '@HD' ) or line.startswith( '@SQ' ) or line.startswith( '@RG' ) or line.startswith( '@PG' ) or line.startswith( '@CO' ) ):
326                        fout.write( line )
327                fout.close()
328            # check that there are results in the output file
329            if os.path.getsize( options.output ) > 0:
330                sys.stdout.write( 'BWA run on %s-end data' % options.genAlignType )
331            else:
332                raise Exception, 'The output file is empty. You may simply have no matches, or there may be an error with your input file or settings.'
333        except Exception, e:
334            stop_err( 'The alignment failed.\n' + str( e ) )
335    finally:
336        # clean up temp dir
337        if os.path.exists( tmp_index_dir ):
338            shutil.rmtree( tmp_index_dir )
339        if os.path.exists( tmp_dir ):
340            shutil.rmtree( tmp_dir )
341
342if __name__=="__main__": __main__()