diff draw_amr_matrix.py @ 0:6044ca44e9f6 draft

Uploaded

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/draw_amr_matrix.py	Tue Feb 07 22:13:57 2023 +0000
@@ -0,0 +1,117 @@
+#!/usr/bin/env python
+
+import argparse
+import csv
+import os
+
+import numpy
+import pandas
+import matplotlib.pyplot as pyplot
+
+
+def get_amr_in_feature_hits(feature_hits):
+    for k in feature_hits.keys():
+        if k.lower().find('amr') >= 0:
+            return feature_hits[k]
+    return None
+
+
+def draw_amr_matrix(feature_hits_files, amr_gene_drug_file, output_dir):
+    ofh = open('process_log', 'w')
+
+    # Read feature_hits_files.
+    feature_hits = pandas.Series(dtype=object)
+    for feature_hits_file in feature_hits_files:
+        feature_name = os.path.basename(feature_hits_file)
+        # Make sure the file is not empty.
+        if os.path.isfile(feature_hits_file) and os.path.getsize(feature_hits_file) > 0:
+            best_hits = pandas.read_csv(filepath_or_buffer=feature_hits_file, sep='\t', header=None)
+            ofh.write("\nFeature file %s will be processed\n" % os.path.basename(feature_hits_file))
+        else:
+            ofh.write("\nEmpty feature file %s will NOT be processed\n" % os.path.basename(feature_hits_file))
+            best_hits = None
+        feature_hits[feature_name] = best_hits
+
+    amr_hits = get_amr_in_feature_hits(feature_hits)
+    ofh.write("\namr_hits:\n%s\n" % str(amr_hits))
+    if amr_hits is not None:
+        amr_to_draw = pandas.DataFrame(columns=['gene', 'drug'])
+        ofh.write("\namr_to_draw:\n%s\n" % str(amr_to_draw))
+        # Read amr_drug_gene_file.
+        amr_gene_drug = pandas.read_csv(amr_gene_drug_file, index_col=None, sep='\t', quoting=csv.QUOTE_NONE, header=None)
+        ofh.write("\namr_gene_drug:\n%s\n" % str(amr_gene_drug))
+
+        # Roll up AMR gene hits.
+        ofh.write("\namr_hits.shape[0]:%s\n" % str(amr_hits.shape[0]))
+        if amr_hits.shape[0] > 0:
+            for gene_idx, gene in amr_hits.iterrows():
+                ofh.write("gene_idx:%s\n" % str(gene_idx))
+                ofh.write("gene:%s\n" % str(gene))
+                gene_name = gene[3]
+                ofh.write("gene_name: %s\n" % str(gene_name))
+                ofh.write("amr_gene_drug[0]: %s\n" % str(amr_gene_drug[0]))
+                drugs = amr_gene_drug.loc[amr_gene_drug[0] == gene_name, :][1]
+                ofh.write("drugs:%s\n" % str(drugs))
+                for drug in drugs:
+                    amr_to_draw = amr_to_draw.append(pandas.Series([gene_name, drug], name=amr_to_draw.shape[0], index=amr_to_draw.columns))
+                    ofh.write("\amr_to_draw:%s\n" % str(amr_to_draw))
+
+        amr_mutations = pandas.Series(dtype=object)
+        if amr_mutations.shape[0] > 0:
+            # Roll up potentially resistance conferring mutations.
+            # TODO: may need to populate this if we need to use call_amr_mutations.
+            for mutation_region, mutation_hits in amr_mutations.iteritems():
+                for mutation_idx, mutation_hit in mutation_hits.iterrows():
+                    mutation_name = mutation_region + ' ' + mutation_hit['REF'] + '->' + mutation_hit['ALT']
+                    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))
+
+        amr_deletions = pandas.DataFrame()
+        if amr_deletions.shape[0] > 0:
+            # Roll up deletions that might confer resistance.
+            # TODO: may need to populate this if we need to use call_insertions.
+            for deletion_idx, deleted_gene in amr_deletions.iterrows():
+                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))
+
+        if amr_to_draw.shape[0] > 1:
+            ofh.write("\nDrawing AMR matrix...\n")
+            present_genes = amr_to_draw['gene'].unique()
+            present_drugs = amr_to_draw['drug'].unique()
+            amr_matrix = pandas.DataFrame(0, index=present_genes, columns=present_drugs)
+            for hit_idx, hit in amr_to_draw.iterrows():
+                amr_matrix.loc[hit[0], hit[1]] = 1
+            amr_matrix_png = os.path.join(output_dir, 'amr_matrix.png')
+            int_matrix = amr_matrix[amr_matrix.columns].astype(int)
+            figure, axis = pyplot.subplots()
+            axis.invert_yaxis()
+            axis.set_yticks(numpy.arange(0.5, len(amr_matrix.index)), minor=False)
+            axis.set_yticklabels(int_matrix.index.values)
+            axis.set_xticks(numpy.arange(0.5, len(amr_matrix.columns)), minor=False)
+            axis.set_xticklabels(amr_matrix.columns.values, rotation=90)
+            axis.xaxis.tick_top()
+            axis.xaxis.set_label_position('top')
+            pyplot.tight_layout()
+            pyplot.savefig(amr_matrix_png, dpi=300)
+        else:
+            ofh.write("\nEmpty AMR matrix, nothing to draw...\n")
+    ofh.close()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--amr_gene_drug_file', action='store', dest='amr_gene_drug_file', help='AMR_gene_drugs tsv file')
+    parser.add_argument('--feature_hits_dir', action='store', dest='feature_hits_dir', help='Directory of tabular files containing feature hits')
+    parser.add_argument('--output_dir', action='store', dest='output_dir', help='Output directory')
+
+    args = parser.parse_args()
+
+    # Get thge collection of feature hits files.  The collection
+    # will be sorted alphabetically and will contain 2 files
+    # named something like AMR_CDS_311_2022_12_20.fasta and
+    # Incompatibility_Groups_2023_01_01.fasta.
+    feature_hits_files = []
+    for file_name in sorted(os.listdir(args.feature_hits_dir)):
+        file_path = os.path.abspath(os.path.join(args.feature_hits_dir, file_name))
+        feature_hits_files.append(file_path)
+
+    draw_amr_matrix(feature_hits_files, args.amr_gene_drug_file, args.output_dir)