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