annotate bwa_wrapper.py @ 3:0e0ff9e9c761 default tip

fix inputs clean name

1 #!/usr/bin/env python
2 
3 """
4 Runs BWA on single-end or paired-end data.
5 Produces a SAM file containing the mappings.
6 Works with BWA version 0.5.9.
7 
8 usage: bwa_wrapper.py [options]
9 
10 See below for options
11 """
12 
13 import optparse, os, shutil, subprocess, sys, tempfile
14 import glob 
15 import gzip, zipfile, tarfile
16 
17 def stop_err( msg ):
18     sys.stderr.write( '%s\n' % msg )
19     sys.exit()
20 
21 def check_is_double_encoded( fastq ):
22     # check that first read is bases, not one base followed by numbers
23     bases = [ 'A', 'C', 'G', 'T', 'a', 'c', 'g', 't', 'N' ]
24     nums = [ '0', '1', '2', '3' ]
25     for line in file( fastq, 'rb'):
26         if not line.strip() or line.startswith( '@' ):
27             continue
28         if len( [ b for b in line.strip() if b in nums ] ) > 0:
29             return False
30         elif line.strip()[0] in bases and len( [ b for b in line.strip() if b in bases ] ) == len( line.strip() ):
31             return True
32         else:
33             raise Exception, 'First line in first read does not appear to be a valid FASTQ read in either base-space or color-space'
34     raise Exception, 'There is no non-comment and non-blank line in your FASTQ file'
35 
36 def __main__():
37     #Parse Command Line
38     parser = optparse.OptionParser()
39     parser.add_option( '-t', '--threads', dest='threads', help='The number of threads to use' )
40     parser.add_option( '-c', '--color-space', dest='color_space', action='store_true', help='If the input files are SOLiD format' )
41     parser.add_option( '-r', '--ref', dest='ref', help='The reference genome to use or index' )
42     parser.add_option( '-f', '--input1', dest='fastq', help='The (forward) fastq file to use for the mapping' )
43     parser.add_option( '-u', '--output', dest='output', help='The file to save the output (SAM format)' )
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', default=None, 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     tmp_input_dir = tempfile.mkdtemp()
82     tmp_output_dir= tempfile.mkdtemp()
83 
84 
85     myarchive = zipfile.ZipFile(options.fastq, 'r', allowZip64=True)
86     myarchive.extractall(tmp_input_dir)
87 
88     for fastq in glob.glob(tmp_input_dir+'/*'):
89 
90 	    sam_output_file=tmp_output_dir+'/'+os.path.splitext(os.path.basename(fastq))[0]+'.sam'
91 	    create_sam=open(sam_output_file, "w")
92 	    create_sam.close()
93 
94 	    # output version # of tool
95 	    try:
96 		tmp = tempfile.NamedTemporaryFile().name
97 		tmp_stdout = open( tmp, 'wb' )
98 		proc = subprocess.Popen( args='bwa 2>&1', shell=True, stdout=tmp_stdout )
99 		tmp_stdout.close()
00 		returncode = proc.wait()
01 		stdout = None
02 		for line in open( tmp_stdout.name, 'rb' ):
03 		    if line.lower().find( 'version' ) >= 0:
04 			stdout = line.strip()
05 			break
06 		if stdout:
07 		    sys.stdout.write( 'BWA %s\n' % stdout )
08 		else:
09 		    raise Exception
10 	    except:
11 		sys.stdout.write( 'Could not determine BWA version\n' )
12 
13 	    # check for color space fastq that's not double-encoded and exit if appropriate
14 	    if options.color_space:
15 		if not check_is_double_encoded( options.fastq ):
16 		    stop_err( 'Your file must be double-encoded (it must be converted from "numbers" to "bases"). See the help section for details' )
17 		#if options.genAlignType == 'paired':
18 		    #if not check_is_double_encoded( options.rfastq ):
19 			#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' )
20 
21 	    #fastq = options.fastq
22 	    #if options.rfastq:
23 		 #rfastq = options.rfastq
24 
25 	    # set color space variable
26 	    if options.color_space:
27 		color_space = '-c'
28 	    else:
29 		color_space = ''
30 
31 	    # make temp directory for placement of indices
32 	    tmp_index_dir = tempfile.mkdtemp()
33 	    tmp_dir = tempfile.mkdtemp()
34 	    # index if necessary
35 	    if options.fileSource == 'history' and not options.do_not_build_index:
36 		ref_file = tempfile.NamedTemporaryFile( dir=tmp_index_dir )
37 		ref_file_name = ref_file.name
38 		ref_file.close()
39 		os.symlink( options.ref, ref_file_name )
40 		# determine which indexing algorithm to use, based on size
41 		try:
42 		    size = os.stat( options.ref ).st_size
43 		    if size <= 2**30: 
44 			indexingAlg = 'is'
45 		    else:
46 			indexingAlg = 'bwtsw'
47 		except:
48 		    indexingAlg = 'is'
49 		indexing_cmds = '%s -a %s' % ( color_space, indexingAlg )
50 		cmd1 = 'bwa index %s %s' % ( indexing_cmds, ref_file_name )
51 		try:
52 		    tmp = tempfile.NamedTemporaryFile( dir=tmp_index_dir ).name
53 		    tmp_stderr = open( tmp, 'wb' )
54 		    proc = subprocess.Popen( args=cmd1, shell=True, cwd=tmp_index_dir, stderr=tmp_stderr.fileno() )
55 		    returncode = proc.wait()
56 		    tmp_stderr.close()
57 		    # get stderr, allowing for case where it's very large
58 		    tmp_stderr = open( tmp, 'rb' )
59 		    stderr = ''
60 		    buffsize = 1048576
61 		    try:
62 			while True:
63 			    stderr += tmp_stderr.read( buffsize )
64 			    if not stderr or len( stderr ) % buffsize != 0:
65 				break
66 		    except OverflowError:
67 			pass
68 		    tmp_stderr.close()
69 		    if returncode != 0:
70 			raise Exception, stderr
71 		except Exception, e:
72 		    # clean up temp dirs
73 		    if os.path.exists( tmp_index_dir ):
74 			shutil.rmtree( tmp_index_dir )
75 		    if os.path.exists( tmp_dir ):
76 			shutil.rmtree( tmp_dir )
77 		    stop_err( 'Error indexing reference sequence. ' + str( e ) )
78 	    else:
79 		ref_file_name = options.ref
80 	    if options.illumina13qual:
81 		illumina_quals = "-I"
82 	    else:
83 		illumina_quals = ""
84 
85 	    # set up aligning and generate aligning command options
86 	    if options.params == 'pre_set':
87 		aligning_cmds = '-t %s %s %s' % ( options.threads, color_space, illumina_quals )
88 		gen_alignment_cmds = ''
89 	    else:
90 		if options.maxEditDist != '0':
91 		    editDist = options.maxEditDist
92 		else:
93 		    editDist = options.fracMissingAligns
94 		if options.seed != '-1':
95 		    seed = '-l %s' % options.seed
96 		else:
97 		    seed = ''
98 		if options.suboptAlign:
99 		    suboptAlign = '-R "%s"' % ( options.suboptAlign )
00 		else:
01 		    suboptAlign = ''
02 		if options.noIterSearch == 'true':
03 		    noIterSearch = '-N'
04 		else:
05 		    noIterSearch = ''
06 		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' % \
07 				( editDist, options.maxGapOpens, options.maxGapExtens, options.disallowLongDel,
08 				  options.disallowIndel, seed, options.maxEditDistSeed, options.threads,
09 				  options.mismatchPenalty, options.gapOpenPenalty, options.gapExtensPenalty,
10 				  suboptAlign, noIterSearch, color_space, illumina_quals )
11 		#if options.genAlignType == 'paired':
12 		    #gen_alignment_cmds = '-a %s -o %s' % ( options.maxInsertSize, options.maxOccurPairing )
13 		    #if options.outputTopNDisc:
14 			#gen_alignment_cmds += ' -N %s' % options.outputTopNDisc
15 
16 		gen_alignment_cmds = ''
17 		if options.rgid:
18 		    if not options.rglb or not options.rgpl or not options.rgsm:
19 			stop_err( 'If you want to specify read groups, you must include the ID, LB, PL, and SM tags.' )
20 		    readGroup = '@RG\tID:%s\tLB:%s\tPL:%s\tSM:%s' % ( options.rgid, options.rglb, options.rgpl, options.rgsm )
21 		    if options.rgcn:
22 			readGroup += '\tCN:%s' % options.rgcn
23 		    if options.rgds:
24 			readGroup += '\tDS:%s' % options.rgds
25 		    if options.rgdt:
26 			readGroup += '\tDT:%s' % options.rgdt
27 		    if options.rgfo:
28 			readGroup += '\tFO:%s' % options.rgfo
29 		    if options.rgks:
30 			readGroup += '\tKS:%s' % options.rgks
31 		    if options.rgpg:
32 			readGroup += '\tPG:%s' % options.rgpg
33 		    if options.rgpi:
34 			readGroup += '\tPI:%s' % options.rgpi
35 		    if options.rgpu:
36 			readGroup += '\tPU:%s' % options.rgpu
37 		    gen_alignment_cmds += ' -r "%s"' % readGroup
38 		if options.outputTopN:
39 		    gen_alignment_cmds += ' -n %s' % options.outputTopN
40 	    # set up output files
41 	    tmp_align_out = tempfile.NamedTemporaryFile( dir=tmp_dir )
42 	    tmp_align_out_name = tmp_align_out.name
43 	    tmp_align_out.close()
44 	    tmp_align_out2 = tempfile.NamedTemporaryFile( dir=tmp_dir )
45 	    tmp_align_out2_name = tmp_align_out2.name
46 	    tmp_align_out2.close()
47 	    # prepare actual aligning and generate aligning commands
48 	    cmd2 = 'bwa aln %s %s %s > %s' % ( aligning_cmds, ref_file_name, fastq, tmp_align_out_name )
49 	    cmd2b = ''
50 	    cmd3 = 'bwa samse %s %s %s %s >> %s' % ( gen_alignment_cmds, ref_file_name, tmp_align_out_name, fastq, sam_output_file )
51 	    # perform alignments
52 	    buffsize = 1048576
53 	    try:
54 		# need to nest try-except in try-finally to handle 2.4
55 		try:
56 		    # align
57 		    try:
58 			tmp = tempfile.NamedTemporaryFile( dir=tmp_dir ).name
59 			tmp_stderr = open( tmp, 'wb' )
60 			proc = subprocess.Popen( args=cmd2, shell=True, cwd=tmp_dir, stderr=tmp_stderr.fileno() )
61 			returncode = proc.wait()
62 			tmp_stderr.close()
63 			# get stderr, allowing for case where it's very large
64 			tmp_stderr = open( tmp, 'rb' )
65 			stderr = ''
66 			try:
67 			    while True:
68 				stderr += tmp_stderr.read( buffsize )
69 				if not stderr or len( stderr ) % buffsize != 0:
70 				    break
71 			except OverflowError:
72 			    pass
73 			tmp_stderr.close()
74 			if returncode != 0:
75 			    raise Exception, stderr
76 		    except Exception, e:
77 			raise Exception, 'Error aligning sequence. ' + str( e )
78 		    # and again if paired data
79 		    try:
80 			if cmd2b:
81 			    tmp = tempfile.NamedTemporaryFile( dir=tmp_dir ).name
82 			    tmp_stderr = open( tmp, 'wb' )
83 			    proc = subprocess.Popen( args=cmd2b, shell=True, cwd=tmp_dir, stderr=tmp_stderr.fileno() )
84 			    returncode = proc.wait()
85 			    tmp_stderr.close()
86 			    # get stderr, allowing for case where it's very large
87 			    tmp_stderr = open( tmp, 'rb' )
88 			    stderr = ''
89 			    try:
90 				while True:
91 				    stderr += tmp_stderr.read( buffsize )
92 				    if not stderr or len( stderr ) % buffsize != 0:
93 					break
94 			    except OverflowError:
95 				pass
96 			    tmp_stderr.close()
97 			    if returncode != 0:
98 				raise Exception, stderr
99 		    except Exception, e:
00 			raise Exception, 'Error aligning second sequence. ' + str( e )
01 		    # generate align
02 		    try:
03 			tmp = tempfile.NamedTemporaryFile( dir=tmp_dir ).name
04 			tmp_stderr = open( tmp, 'wb' )
05 			proc = subprocess.Popen( args=cmd3, shell=True, cwd=tmp_dir, stderr=tmp_stderr.fileno() )
06 			returncode = proc.wait()
07 			tmp_stderr.close()
08 			# get stderr, allowing for case where it's very large
09 			tmp_stderr = open( tmp, 'rb' )
10 			stderr = ''
11 			try:
12 			    while True:
13 				stderr += tmp_stderr.read( buffsize )
14 				if not stderr or len( stderr ) % buffsize != 0:
15 				    break
16 			except OverflowError:
17 			    pass
18 			tmp_stderr.close()
19 			if returncode != 0:
20 			    raise Exception, stderr
21 		    except Exception, e:
22 			raise Exception, 'Error generating alignments. ' + str( e ) 
23 		    # remove header if necessary
24 		    if options.suppressHeader == 'true':
25 			tmp_out = tempfile.NamedTemporaryFile( dir=tmp_dir)
26 			tmp_out_name = tmp_out.name
27 			tmp_out.close()
28 			try:
29 			    shutil.move( sam_output_file, tmp_out_name )
30 			except Exception, e:
31 			    raise Exception, 'Error moving output file before removing headers. ' + str( e )
32 			fout = file( sam_output_file, 'w' )
33 			for line in file( tmp_out.name, 'r' ):
34 			    if not ( line.startswith( '@HD' ) or line.startswith( '@SQ' ) or line.startswith( '@RG' ) or line.startswith( '@PG' ) or line.startswith( '@CO' ) ):
35 				fout.write( line )
36 			fout.close()
37 		    # check that there are results in the output file
38 		    if os.path.getsize( sam_output_file ) > 0:
39 			sys.stdout.write( 'BWA run on single-end data')
40 		    else:
41 			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.'
42 		except Exception, e:
43 		    stop_err( 'The alignment failed.\n' + str( e ) )
44 	    finally:
45 		# clean up temp dir
46 		if os.path.exists( tmp_index_dir ):
47 		    shutil.rmtree( tmp_index_dir )
48 		if os.path.exists( tmp_dir ):
49 		    shutil.rmtree( tmp_dir )	
50 
51     # put all in an archive
52     mytotalzipfile=zipfile.ZipFile(options.output, 'w', allowZip64=True)
53     os.chdir(tmp_output_dir)
54     for samfile in glob.glob(tmp_output_dir+'/*'):
55 	mytotalzipfile.write(os.path.basename(samfile))
56 
57 
58 if __name__=="__main__": __main__()