Add 'build/' from commit '555b039609a3c8700c27767c255fdfe00eb42063'
[arvados.git] / crunch_scripts / GATK2-realign
1 #!/usr/bin/env python
2
3 import os
4 import re
5 import arvados
6 import arvados_gatk2
7 import arvados_picard
8 import arvados_samtools
9 from arvados_ipc import *
10
11 class InvalidArgumentError(Exception):
12     pass
13
14 arvados_samtools.one_task_per_bam_file(if_sequence=0, and_end_task=True)
15
16 this_job = arvados.current_job()
17 this_task = arvados.current_task()
18 tmpdir = arvados.current_task().tmpdir
19 arvados.util.clear_tmpdir()
20
21 known_sites_files = arvados.getjobparam(
22     'known_sites',
23     ['dbsnp_137.b37.vcf',
24      'Mills_and_1000G_gold_standard.indels.b37.vcf',
25      ])
26 bundle_dir = arvados.util.collection_extract(
27     collection = this_job['script_parameters']['gatk_bundle'],
28     files = [
29         'human_g1k_v37.dict',
30         'human_g1k_v37.fasta',
31         'human_g1k_v37.fasta.fai'
32         ] + known_sites_files + [v + '.idx' for v in known_sites_files],
33     path = 'gatk_bundle')
34 ref_fasta_files = [os.path.join(bundle_dir, f)
35                    for f in os.listdir(bundle_dir)
36                    if re.search(r'\.fasta(\.gz)?$', f)]
37 regions_args = []
38 if 'regions' in this_job['script_parameters']:
39     regions_dir = arvados.util.collection_extract(
40         collection = this_job['script_parameters']['regions'],
41         path = 'regions')
42     region_padding = int(this_job['script_parameters']['region_padding'])
43     for f in os.listdir(regions_dir):
44         if re.search(r'\.bed$', f):
45             regions_args += [
46                 '--intervals', os.path.join(regions_dir, f),
47                 '--interval_padding', str(region_padding)
48                 ]
49
50 input_collection = this_task['parameters']['input']
51 input_dir = arvados.util.collection_extract(
52     collection = input_collection,
53     path = os.path.join(this_task.tmpdir, 'input'))
54 input_bam_files = []
55 for f in arvados.util.listdir_recursive(input_dir):
56     if re.search(r'\.bam$', f):
57         input_stream_name, input_file_name = os.path.split(f)
58         input_bam_files += [os.path.join(input_dir, f)]
59 if len(input_bam_files) != 1:
60     raise InvalidArgumentError("Expected exactly one bam file per task.")
61
62 known_sites_args = []
63 for f in known_sites_files:
64     known_sites_args += ['-known', os.path.join(bundle_dir, f)]
65
66 children = {}
67 pipes = {}
68
69 arvados_gatk2.run(
70     args=[
71         '-nt', arvados_gatk2.cpus_per_task(),
72         '-T', 'RealignerTargetCreator',
73         '-R', ref_fasta_files[0],
74         '-I', input_bam_files[0],
75         '-o', os.path.join(tmpdir, 'intervals.list')
76         ] + known_sites_args + regions_args)
77
78 pipe_setup(pipes, 'IndelRealigner')
79 if 0 == named_fork(children, 'IndelRealigner'):
80     pipe_closeallbut(pipes, ('IndelRealigner', 'w'))
81     arvados_gatk2.run(
82         args=[
83         '-T', 'IndelRealigner',
84         '-R', ref_fasta_files[0],
85         '-targetIntervals', os.path.join(tmpdir, 'intervals.list'),
86         '-I', input_bam_files[0],
87         '-o', '/dev/fd/' + str(pipes['IndelRealigner','w']),
88         '--disable_bam_indexing',
89         ] + known_sites_args + regions_args,
90         close_fds=False)
91     os._exit(0)
92 os.close(pipes.pop(('IndelRealigner','w'), None))
93
94 pipe_setup(pipes, 'bammanifest')
95 pipe_setup(pipes, 'bam')
96 if 0==named_fork(children, 'bammanifest'):
97     pipe_closeallbut(pipes,
98                      ('IndelRealigner', 'r'),
99                      ('bammanifest', 'w'),
100                      ('bam', 'w'))
101     out = arvados.CollectionWriter()
102     out.start_new_stream(input_stream_name)
103     out.start_new_file(input_file_name)
104     while True:
105         buf = os.read(pipes['IndelRealigner','r'], 2**20)
106         if len(buf) == 0:
107             break
108         os.write(pipes['bam','w'], buf)
109         out.write(buf)
110     os.write(pipes['bammanifest','w'], out.manifest_text())
111     os.close(pipes['bammanifest','w'])
112     os._exit(0)
113
114 pipe_setup(pipes, 'index')
115 if 0==named_fork(children, 'index'):
116     pipe_closeallbut(pipes, ('bam', 'r'), ('index', 'w'))
117     arvados_picard.run(
118         'BuildBamIndex',
119         params={
120             'i': '/dev/fd/' + str(pipes['bam','r']),
121             'o': '/dev/fd/' + str(pipes['index','w']),
122             'quiet': 'true',
123             'validation_stringency': 'LENIENT'
124             },
125         close_fds=False)
126     os._exit(0)
127
128 pipe_setup(pipes, 'indexmanifest')
129 if 0==named_fork(children, 'indexmanifest'):
130     pipe_closeallbut(pipes, ('index', 'r'), ('indexmanifest', 'w'))
131     out = arvados.CollectionWriter()
132     out.start_new_stream(input_stream_name)
133     out.start_new_file(re.sub('\.bam$', '.bai', input_file_name))
134     while True:
135         buf = os.read(pipes['index','r'], 2**20)
136         if len(buf) == 0:
137             break
138         out.write(buf)
139     os.write(pipes['indexmanifest','w'], out.manifest_text())
140     os.close(pipes['indexmanifest','w'])
141     os._exit(0)
142
143 pipe_closeallbut(pipes, ('bammanifest', 'r'), ('indexmanifest', 'r'))
144 outmanifest = ''
145 for which in ['bammanifest', 'indexmanifest']:
146     with os.fdopen(pipes[which,'r'], 'rb', 2**20) as f:
147         while True:
148             buf = f.read()
149             if buf == '':
150                 break
151             outmanifest += buf
152
153 all_ok = True
154 for (childname, pid) in children.items():
155     all_ok = all_ok and waitpid_and_check_exit(pid, childname)
156
157 if all_ok:
158     this_task.set_output(outmanifest)
159 else:
160     sys.exit(1)