crunch_scripts/bwa-aln

   1 #!/usr/bin/env python
   2 # Copyright (C) The Arvados Authors. All rights reserved.
   3 #
   4 # SPDX-License-Identifier: Apache-2.0
   5
   6 import arvados
   7 import arvados_bwa
   8 import arvados_samtools
   9 import os
  10 import re
  11 import sys
  12 import subprocess
  13
  14 arvados_bwa.one_task_per_pair_input_file(if_sequence=0, and_end_task=True)
  15
  16 this_job = arvados.current_job()
  17 this_task = arvados.current_task()
  18 ref_dir = arvados.util.collection_extract(
  19     collection = this_job['script_parameters']['reference_index'],
  20     path = 'reference',
  21     decompress = False)
  22
  23 ref_basename = None
  24 for f in os.listdir(ref_dir):
  25     basename = re.sub(r'\.bwt$', '', f)
  26     if basename != f:
  27         ref_basename = os.path.join(ref_dir, basename)
  28 if ref_basename == None:
  29     raise Exception("Could not find *.bwt in reference collection.")
  30
  31 tmp_dir = arvados.current_task().tmpdir
  32
  33 class Aligner:
  34     def input_filename(self):
  35         for s in arvados.CollectionReader(self.collection).all_streams():
  36             for f in s.all_files():
  37                 return f.decompressed_name()
  38     def generate_input(self):
  39         for s in arvados.CollectionReader(self.collection).all_streams():
  40             for f in s.all_files():
  41                 for s in f.readall_decompressed():
  42                     yield s
  43     def aln(self, input_param):
  44         self.collection = this_task['parameters'][input_param]
  45         reads_filename = os.path.join(tmp_dir, self.input_filename())
  46         aln_filename = os.path.join(tmp_dir, self.input_filename() + '.sai')
  47         reads_pipe_r, reads_pipe_w = os.pipe()
  48         if os.fork() == 0:
  49             os.close(reads_pipe_r)
  50             reads_file = open(reads_filename, 'wb')
  51             for s in self.generate_input():
  52                 if len(s) != os.write(reads_pipe_w, s):
  53                     raise Exception("short write")
  54                 reads_file.write(s)
  55             reads_file.close()
  56             os.close(reads_pipe_w)
  57             sys.exit(0)
  58         os.close(reads_pipe_w)
  59
  60         aln_file = open(aln_filename, 'wb')
  61         bwa_proc = subprocess.Popen(
  62             [arvados_bwa.bwa_binary(),
  63              'aln', '-t', '16',
  64              ref_basename,
  65              '-'],
  66             stdin=os.fdopen(reads_pipe_r, 'rb', 2**20),
  67             stdout=aln_file)
  68         aln_file.close()
  69         return reads_filename, aln_filename
  70
  71 reads_1, alignments_1 = Aligner().aln('input_1')
  72 reads_2, alignments_2 = Aligner().aln('input_2')
  73 pid1, exit1 = os.wait()
  74 pid2, exit2 = os.wait()
  75 if exit1 != 0 or exit2 != 0:
  76     raise Exception("bwa aln exited non-zero (0x%x, 0x%x)" % (exit1, exit2))
  77
  78 # output alignments in sam format to pipe
  79 sam_pipe_r, sam_pipe_w = os.pipe()
  80 sam_pid = os.fork()
  81 if sam_pid != 0:
  82     # parent
  83     os.close(sam_pipe_w)
  84 else:
  85     # child
  86     os.close(sam_pipe_r)
  87     arvados_bwa.run('sampe',
  88                     [ref_basename,
  89                      alignments_1, alignments_2,
  90                      reads_1, reads_2],
  91                     stdout=os.fdopen(sam_pipe_w, 'wb', 2**20))
  92     sys.exit(0)
  93
  94 # convert sam (sam_pipe_r) to bam (bam_pipe_w)
  95 bam_pipe_r, bam_pipe_w = os.pipe()
  96 bam_pid = os.fork()
  97 if bam_pid != 0:
  98     # parent
  99     os.close(bam_pipe_w)
 100     os.close(sam_pipe_r)
 101 else:
 102     # child
 103     os.close(bam_pipe_r)
 104     arvados_samtools.run('view',
 105                          ['-S', '-b',
 106                           '-'],
 107                          stdin=os.fdopen(sam_pipe_r, 'rb', 2**20),
 108                          stdout=os.fdopen(bam_pipe_w, 'wb', 2**20))
 109     sys.exit(0)
 110
 111 # copy bam (bam_pipe_r) to Keep
 112 out_bam_filename = os.path.split(reads_1)[-1] + '.bam'
 113 out = arvados.CollectionWriter()
 114 out.start_new_stream()
 115 out.start_new_file(out_bam_filename)
 116 out.write(os.fdopen(bam_pipe_r, 'rb', 2**20))
 117
 118 # make sure everyone exited nicely
 119 pid3, exit3 = os.waitpid(sam_pid, 0)
 120 if exit3 != 0:
 121     raise Exception("bwa sampe exited non-zero (0x%x)" % exit3)
 122 pid4, exit4 = os.waitpid(bam_pid, 0)
 123 if exit4 != 0:
 124     raise Exception("samtools view exited non-zero (0x%x)" % exit4)
 125
 126 # proclaim success
 127 this_task.set_output(out.finish())