annotate draw_amr_matrix.py @ 11:da1c9c1be421 draft

Uploaded
author greg
date Mon, 27 Mar 2023 19:41:30 +0000
parents 03240ffe969a
children aa2b737102dc
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
1 #!/usr/bin/env python
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
2
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
3 import argparse
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
4 import csv
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
5 import os
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
6 import subprocess
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
7 import sys
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
8 import tempfile
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
9
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
10 import Bio.SeqIO
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
11 import numpy
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
12 import pandas
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
13 import matplotlib.pyplot as pyplot
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
14
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
15
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
16 def get_amr_in_feature_hits(amr_feature_hits):
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
17 for k in amr_feature_hits.keys():
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
18 if k.lower().find('amr') >= 0:
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
19 return amr_feature_hits[k]
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
20 return None
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
21
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
22
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
23 def load_fasta(fasta_file):
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
24 sequence = pandas.Series(dtype=object)
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
25 for contig in Bio.SeqIO.parse(fasta_file, 'fasta'):
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
26 sequence[contig.id] = contig
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
27 return sequence
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
28
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
29
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
30 def run_command(cmd):
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
31 try:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
32 tmp_name = tempfile.NamedTemporaryFile(dir=".").name
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
33 tmp_stderr = open(tmp_name, 'wb')
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
34 proc = subprocess.Popen(args=cmd, shell=True, stderr=tmp_stderr.fileno())
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
35 returncode = proc.wait()
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
36 tmp_stderr.close()
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
37 if returncode != 0:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
38 # Get stderr, allowing for case where it's very large.
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
39 tmp_stderr = open(tmp_name, 'rb')
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
40 stderr = ''
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
41 buffsize = 1048576
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
42 try:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
43 while True:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
44 stderr += tmp_stderr.read(buffsize)
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
45 if not stderr or len(stderr) % buffsize != 0:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
46 break
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
47 except OverflowError:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
48 pass
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
49 tmp_stderr.close()
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
50 os.remove(tmp_name)
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
51 stop_err(stderr)
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
52 except Exception as e:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
53 stop_err('Command:\n%s\n\nended with error:\n%s\n\n' % (cmd, str(e)))
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
54
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
55
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
56 def stop_err(msg):
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
57 sys.stderr.write(msg)
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
58 sys.exit(1)
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
59
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
60
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
61 def draw_amr_matrix(amr_feature_hits_files, amr_deletions_file, varscan_vcf_file, amr_mutation_regions_bed_file, amr_gene_drug_file, reference, reference_size, mutation_regions_dir, amr_matrix_png_dir, errors):
70073df30a06 Uploaded
greg
parents: 8
diff changeset
62 efh = open(errors, 'w')
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
63 ofh = open('process_log', 'w')
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
64
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
65 # Read amr_feature_hits_files.
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
66 amr_feature_hits = pandas.Series(dtype=object)
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
67 for amr_feature_hits_file in amr_feature_hits_files:
11
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
68 ofh.write("\namr_feature_hits_file: %s\n" % amr_feature_hits_file)
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
69 feature_name = os.path.basename(amr_feature_hits_file)
11
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
70 ofh.write("\nfeature_name: %s\n" % feature_name)
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
71 # Make sure the file is not empty.
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
72 if os.path.isfile(amr_feature_hits_file) and os.path.getsize(amr_feature_hits_file) > 0:
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
73 best_hits = pandas.read_csv(filepath_or_buffer=amr_feature_hits_file, sep='\t', header=None)
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
74 ofh.write("\nFeature file %s will be processed\n" % os.path.basename(amr_feature_hits_file))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
75 else:
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
76 ofh.write("\nEmpty feature file %s will NOT be processed\n" % os.path.basename(amr_feature_hits_file))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
77 best_hits = None
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
78 amr_feature_hits[feature_name] = best_hits
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
79
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
80 amr_hits = get_amr_in_feature_hits(amr_feature_hits)
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
81 ofh.write("\namr_hits:\n%s\n" % str(amr_hits))
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
82 if amr_hits is not None:
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
83 amr_to_draw = pandas.DataFrame(columns=['gene', 'drug'])
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
84 ofh.write("\namr_to_draw:\n%s\n" % str(amr_to_draw))
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
85 # Read amr_drug_gene_file.
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
86 amr_gene_drug = pandas.read_csv(amr_gene_drug_file, index_col=None, sep='\t', quoting=csv.QUOTE_NONE, header=None)
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
87 ofh.write("\namr_gene_drug:\n%s\n" % str(amr_gene_drug))
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
88
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
89 # Roll up AMR gene hits.
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
90 ofh.write("\namr_hits.shape[0]: %s\n" % str(amr_hits.shape[0]))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
91 if amr_hits.shape[0] > 0:
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
92 for gene_idx, gene in amr_hits.iterrows():
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
93 ofh.write("gene_idx: %s\n" % str(gene_idx))
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
94 ofh.write("gene: %s\n" % str(gene))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
95 gene_name = gene[3]
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
96 ofh.write("gene_name: %s\n" % str(gene_name))
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
97 ofh.write("amr_gene_drug[0]: %s\n" % str(amr_gene_drug[0]))
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
98 drugs = amr_gene_drug.loc[amr_gene_drug[0] == gene_name, :][1]
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
99 ofh.write("drugs: %s\n" % str(drugs))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
100 for drug in drugs:
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
101 amr_to_draw = amr_to_draw.append(pandas.Series([gene_name, drug], name=amr_to_draw.shape[0], index=amr_to_draw.columns))
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
102 ofh.write("\amr_to_draw: %s\n" % str(amr_to_draw))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
103
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
104 ofh.write("\nvarscan_vcf_file is None: %s\n" % str(varscan_vcf_file == 'None'))
7
389c98d344ce Uploaded
greg
parents: 6
diff changeset
105 if varscan_vcf_file not in [None, 'None'] and os.path.getsize(varscan_vcf_file) > 0:
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
106 amr_mutations = pandas.Series(dtype=object)
7
389c98d344ce Uploaded
greg
parents: 6
diff changeset
107 if amr_mutation_regions_bed_file is not None:
389c98d344ce Uploaded
greg
parents: 6
diff changeset
108 mutation_regions = pandas.read_csv(amr_mutation_regions_bed_file, header=0, sep='\t', index_col=False)
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
109 # Validate mutation regions.
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
110 if mutation_regions.shape[1] != 7:
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
111 efh.write("The selected mutations regions BED file is invalid, it should be a six column file.\n")
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
112 elif mutation_regions.shape[0] == 0:
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
113 efh.write("There are no rows in the selected mutation regions file.\n")
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
114 else:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
115 for region_i in range(mutation_regions.shape[0]):
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
116 region = mutation_regions.iloc[region_i, :]
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
117 if region[0] not in reference:
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
118 efh.write("Mutation region '%s' not found in reference genome.\n" % str(region))
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
119 break
70073df30a06 Uploaded
greg
parents: 8
diff changeset
120 if not isinstance(region[1], numpy.int64):
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
121 efh.write("Non-integer found in mutation region start (column 2): %s.\n" % str(region[1]))
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
122 break
70073df30a06 Uploaded
greg
parents: 8
diff changeset
123 if not isinstance(region[2], numpy.int64):
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
124 efh.write("Non-integer found in mutation region start (column 3): %s.\n" % str(region[2]))
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
125 break
70073df30a06 Uploaded
greg
parents: 8
diff changeset
126 if region[1] <= 0 or region[2] <= 0:
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
127 efh.write("Mutation region '%s' starts before the reference sequence.\n" % str(region))
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
128 if region[1] > len(reference[region[0]].seq) or region[2] > len(reference[region[0]].seq):
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
129 efh.write("Mutation region '%s' ends after the reference sequence.\n" % str(region))
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
130 if not region.get('type', default='No Type') in ['snp', 'small-indel', 'any']:
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
131 ofh.write("\n\nSkipping mutation region '%s' with invalid type '%s', valid types are 'snp', 'small-indel', 'any'.\n\n" % (str(region), str(region.get('type', default='No Type'))))
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
132 continue
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
133 ofh.write("\nFinding AMR mutations for %s.\n" % str(region['name']))
7
389c98d344ce Uploaded
greg
parents: 6
diff changeset
134 region_bed = 'region_%s.bed' % region_i
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
135 ofh.write("region_bed: %s\n" % str(region_bed))
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
136 mutation_regions.loc[[region_i], ].to_csv(path_or_buf=region_bed, sep='\t', header=False, index=False)
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
137 ofh.write("mutation_regions.loc[[region_i], ]:\n%s\n" % str(mutation_regions.loc[[region_i], ]))
7
389c98d344ce Uploaded
greg
parents: 6
diff changeset
138 region_mutations_tsv = os.path.join(mutation_regions_dir, 'region_%s_mutations.tsv' % region_i)
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
139 ofh.write("region_mutations_tsv: %s\n" % str(region_mutations_tsv))
03240ffe969a Uploaded
greg
parents: 9
diff changeset
140 cmd = ' '.join(['bedtools intersect',
03240ffe969a Uploaded
greg
parents: 9
diff changeset
141 '-nonamecheck',
03240ffe969a Uploaded
greg
parents: 9
diff changeset
142 '-wb',
03240ffe969a Uploaded
greg
parents: 9
diff changeset
143 '-a', region_bed,
03240ffe969a Uploaded
greg
parents: 9
diff changeset
144 '-b', varscan_vcf_file,
7
389c98d344ce Uploaded
greg
parents: 6
diff changeset
145 ' | awk \'BEGIN{getline < "' + amr_mutation_regions_bed_file + '";printf $0"\\t";',
389c98d344ce Uploaded
greg
parents: 6
diff changeset
146 'getline < "' + varscan_vcf_file + '"; getline < "' + varscan_vcf_file + '";print $0}{print}\'',
389c98d344ce Uploaded
greg
parents: 6
diff changeset
147 '1>' + region_mutations_tsv])
6
caf554e039b2 Uploaded
greg
parents: 4
diff changeset
148 ofh.write("\ncmd:\n%s\n" % cmd)
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
149 run_command(cmd)
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
150 try:
03240ffe969a Uploaded
greg
parents: 9
diff changeset
151 ofh.write("After running command, os.path.getsize((region_mutations_tsv): %s\n" % str(os.path.getsize(region_mutations_tsv)))
03240ffe969a Uploaded
greg
parents: 9
diff changeset
152 region_mutations = pandas.read_csv(region_mutations_tsv, sep='\t', header=0, index_col=False)
11
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
153 ofh.write("\nregion_mutations: %s\n" % region_mutations)
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
154 except Exception:
03240ffe969a Uploaded
greg
parents: 9
diff changeset
155 continue
03240ffe969a Uploaded
greg
parents: 9
diff changeset
156 # Figure out what kind of mutations are in this region.
03240ffe969a Uploaded
greg
parents: 9
diff changeset
157 region_mutation_types = pandas.Series(['snp'] * region_mutations.shape[0], name='TYPE', index=region_mutations.index)
03240ffe969a Uploaded
greg
parents: 9
diff changeset
158 region_mutation_types[region_mutations['REF'].str.len() != region_mutations['ALT'].str.len()] = 'small-indel'
03240ffe969a Uploaded
greg
parents: 9
diff changeset
159 region_mutation_drugs = pandas.Series(region['drug'] * region_mutations.shape[0], name='DRUG', index=region_mutations.index)
03240ffe969a Uploaded
greg
parents: 9
diff changeset
160 region_notes = pandas.Series(region['note'] * region_mutations.shape[0], name='NOTE', index=region_mutations.index)
03240ffe969a Uploaded
greg
parents: 9
diff changeset
161 region_mutations = pandas.concat([region_mutations, region_mutation_types, region_mutation_drugs, region_notes], axis=1)
03240ffe969a Uploaded
greg
parents: 9
diff changeset
162 region_mutations = region_mutations[['#CHROM', 'POS', 'TYPE', 'REF', 'ALT', 'DRUG', 'NOTE']]
03240ffe969a Uploaded
greg
parents: 9
diff changeset
163 amr_mutations[region['name']] = region_mutations
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
164 else:
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
165 ofh.write("\nMutation region BED file not received.\n")
11
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
166 ofh.write("\nAfter processing mutations, amr_mutations: %s\n" % str(amr_mutations))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
167 # Roll up potentially resistance conferring mutations.
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
168 ofh.write("\n##### Rolling up potentially resistance conferring mutations..\n")
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
169 for mutation_region, mutation_hits in amr_mutations.iteritems():
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
170 ofh.write("mutation_region: %s\n" % str(mutation_region))
03240ffe969a Uploaded
greg
parents: 9
diff changeset
171 ofh.write("mutation_hits: %s\n" % str(mutation_hits))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
172 for mutation_idx, mutation_hit in mutation_hits.iterrows():
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
173 ofh.write("mutation_idx: %s\n" % str(mutation_idx))
03240ffe969a Uploaded
greg
parents: 9
diff changeset
174 ofh.write("mutation_hit: %s\n" % str(mutation_hit))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
175 mutation_name = mutation_region + ' ' + mutation_hit['REF'] + '->' + mutation_hit['ALT']
10
03240ffe969a Uploaded
greg
parents: 9
diff changeset
176 ofh.write("mutation_name: %s\n" % str(mutation_name))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
177 amr_to_draw = amr_to_draw.append(pandas.Series([mutation_name, mutation_hit['DRUG']], name=amr_to_draw.shape[0], index=amr_to_draw.columns))
11
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
178 ofh.write("\nAfter processing mutations, amr_to_draw: %s\n" % str(amr_to_draw))
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
179 ofh.write("\nAfter processing mutations, amr_to_draw.shape[0]: %s\n" % str(amr_to_draw.shape[0]))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
180
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
181 if amr_deletions_file not in [None, 'None'] and os.path.getsize(amr_deletions_file) > 0:
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
182 # Roll up deletions that might confer resistance.
3
7d7884f2d921 Uploaded
greg
parents: 1
diff changeset
183 try:
7d7884f2d921 Uploaded
greg
parents: 1
diff changeset
184 amr_deletions = pandas.read_csv(filepath_or_buffer=amr_deletions_file, sep='\t', header=None)
7d7884f2d921 Uploaded
greg
parents: 1
diff changeset
185 except Exception:
7d7884f2d921 Uploaded
greg
parents: 1
diff changeset
186 amr_deletions = pandas.DataFrame()
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
187 if amr_deletions.shape[0] > 0:
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
188 amr_deletions.columns = ['contig', 'start', 'stop', 'name', 'type', 'drug', 'note']
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
189 amr_deletions = amr_deletions.loc[amr_deletions['type'].isin(['large-deletion', 'any']), :]
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
190 for deletion_idx, deleted_gene in amr_deletions.iterrows():
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
191 amr_to_draw = amr_to_draw.append(pandas.Series(['\u0394' + deleted_gene[3], deleted_gene[5]], name=amr_to_draw.shape[0], index=amr_to_draw.columns))
11
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
192 ofh.write("\nAfter processing deletions, amr_to_draw: %s\n" % str(amr_to_draw))
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
193
11
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
194 ofh.write("\namr_to_draw.shape[0]: %s\n" % str(amr_to_draw.shape[0]))
da1c9c1be421 Uploaded
greg
parents: 10
diff changeset
195 if amr_to_draw.shape[0] > 0:
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
196 ofh.write("\nDrawing AMR matrix...\n")
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
197 present_genes = amr_to_draw['gene'].unique()
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
198 present_drugs = amr_to_draw['drug'].unique()
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
199 amr_matrix = pandas.DataFrame(0, index=present_genes, columns=present_drugs)
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
200 for hit_idx, hit in amr_to_draw.iterrows():
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
201 amr_matrix.loc[hit[0], hit[1]] = 1
7
389c98d344ce Uploaded
greg
parents: 6
diff changeset
202 amr_matrix_png = os.path.join(amr_matrix_png_dir, 'amr_matrix.png')
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
203 int_matrix = amr_matrix[amr_matrix.columns].astype(int)
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
204 figure, axis = pyplot.subplots()
6
caf554e039b2 Uploaded
greg
parents: 4
diff changeset
205 heatmap = axis.pcolor(int_matrix, cmap=pyplot.cm.Blues, linewidth=0)
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
206 axis.invert_yaxis()
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
207 axis.set_yticks(numpy.arange(0.5, len(amr_matrix.index)), minor=False)
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
208 axis.set_yticklabels(int_matrix.index.values)
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
209 axis.set_xticks(numpy.arange(0.5, len(amr_matrix.columns)), minor=False)
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
210 axis.set_xticklabels(amr_matrix.columns.values, rotation=90)
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
211 axis.xaxis.tick_top()
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
212 axis.xaxis.set_label_position('top')
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
213 pyplot.tight_layout()
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
214 pyplot.savefig(amr_matrix_png, dpi=300)
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
215 else:
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
216 ofh.write("\nEmpty AMR matrix, nothing to draw...\n")
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
217 efh.close()
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
218 ofh.close()
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
219
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
220
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
221 if __name__ == '__main__':
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
222 parser = argparse.ArgumentParser()
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
223
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
224 parser.add_argument('--amr_feature_hits_dir', action='store', dest='amr_feature_hits_dir', help='Directory of tabular files containing feature hits')
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
225 parser.add_argument('--amr_deletions_file', action='store', dest='amr_deletions_file', default=None, help='AMR deletions BED file')
7
389c98d344ce Uploaded
greg
parents: 6
diff changeset
226 parser.add_argument('--varscan_vcf_file', action='store', dest='varscan_vcf_file', default=None, help='Varscan VCF file produced by the call_amr_mutations tool')
389c98d344ce Uploaded
greg
parents: 6
diff changeset
227 parser.add_argument('--amr_mutation_regions_bed_file', action='store', dest='amr_mutation_regions_bed_file', default=None, help='AMR mutation regions BED file')
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
228 parser.add_argument('--amr_gene_drug_file', action='store', dest='amr_gene_drug_file', help='AMR_gene_drugs tsv file')
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
229 parser.add_argument('--reference_genome', action='store', dest='reference_genome', help='Reference genome fasta file')
7
389c98d344ce Uploaded
greg
parents: 6
diff changeset
230 parser.add_argument('--mutation_regions_dir', action='store', dest='mutation_regions_dir', help='Directory for mutation regions TSV files produced by this tool')
389c98d344ce Uploaded
greg
parents: 6
diff changeset
231 parser.add_argument('--amr_matrix_png_dir', action='store', dest='amr_matrix_png_dir', help='Directory for PNG files produced by this tool')
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
232 parser.add_argument('--errors', action='store', dest='errors', help='Output file containing errors')
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
233
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
234 args = parser.parse_args()
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
235
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
236 # Get the collection of feature hits files. The collection
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
237 # will be sorted alphabetically and will contain 2 files
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
238 # named something like AMR_CDS_311_2022_12_20.fasta and
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
239 # Incompatibility_Groups_2023_01_01.fasta.
1
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
240 amr_feature_hits_files = []
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
241 for file_name in sorted(os.listdir(args.amr_feature_hits_dir)):
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
242 file_path = os.path.abspath(os.path.join(args.amr_feature_hits_dir, file_name))
5c923c77cf5f Uploaded
greg
parents: 0
diff changeset
243 amr_feature_hits_files.append(file_path)
0
6044ca44e9f6 Uploaded
greg
parents:
diff changeset
244
4
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
245 # Load the reference genome into memory.
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
246 reference = load_fasta(args.reference_genome)
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
247 reference_size = 0
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
248 for i in reference:
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
249 reference_size += len(i.seq)
33a0ea992043 Uploaded
greg
parents: 3
diff changeset
250
9
70073df30a06 Uploaded
greg
parents: 8
diff changeset
251 draw_amr_matrix(amr_feature_hits_files, args.amr_deletions_file, args.varscan_vcf_file, args.amr_mutation_regions_bed_file, args.amr_gene_drug_file, reference, reference_size, args.mutation_regions_dir, args.amr_matrix_png_dir, args.errors)