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