/tools/sr_mapping/bwa_wrapper.py

https://bitbucket.org/cistrome/cistrome-harvard/ · Python · 342 lines · 296 code · 12 blank · 34 comment · 113 complexity · 3a9aba466fb4a8a8029ac3c4fcccf412 MD5 · raw file

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