comparison draw_amr_matrix.py @ 9:70073df30a06 draft

Uploaded
author greg
date Tue, 21 Mar 2023 14:06:11 +0000
parents 7fe8ea50a81d
children 03240ffe969a
comparison
equal deleted inserted replaced
8:7fe8ea50a81d 9:70073df30a06
56 def stop_err(msg): 56 def stop_err(msg):
57 sys.stderr.write(msg) 57 sys.stderr.write(msg)
58 sys.exit(1) 58 sys.exit(1)
59 59
60 60
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): 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):
62 efh = open(errors, 'w')
62 ofh = open('process_log', 'w') 63 ofh = open('process_log', 'w')
63 64
64 # Read amr_feature_hits_files. 65 # Read amr_feature_hits_files.
65 amr_feature_hits = pandas.Series(dtype=object) 66 amr_feature_hits = pandas.Series(dtype=object)
66 for amr_feature_hits_file in amr_feature_hits_files: 67 for amr_feature_hits_file in amr_feature_hits_files:
96 ofh.write("drugs: %s\n" % str(drugs)) 97 ofh.write("drugs: %s\n" % str(drugs))
97 for drug in drugs: 98 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)) 99 amr_to_draw = amr_to_draw.append(pandas.Series([gene_name, drug], name=amr_to_draw.shape[0], index=amr_to_draw.columns))
99 ofh.write("\amr_to_draw: %s\n" % str(amr_to_draw)) 100 ofh.write("\amr_to_draw: %s\n" % str(amr_to_draw))
100 101
101 ofh.write("\nvarscan_vcf_file si None: %s\n" % str(varscan_vcf_file == 'None')) 102 ofh.write("\nvarscan_vcf_file is None: %s\n" % str(varscan_vcf_file == 'None'))
102 if varscan_vcf_file not in [None, 'None'] and os.path.getsize(varscan_vcf_file) > 0: 103 if varscan_vcf_file not in [None, 'None'] and os.path.getsize(varscan_vcf_file) > 0:
103 amr_mutations = pandas.Series(dtype=object) 104 amr_mutations = pandas.Series(dtype=object)
104 if amr_mutation_regions_bed_file is not None: 105 if amr_mutation_regions_bed_file is not None:
105 mutation_regions = pandas.read_csv(amr_mutation_regions_bed_file, header=0, sep='\t', index_col=False) 106 mutation_regions = pandas.read_csv(amr_mutation_regions_bed_file, header=0, sep='\t', index_col=False)
107 # Validate mutation regions.
106 if mutation_regions.shape[1] != 7: 108 if mutation_regions.shape[1] != 7:
107 ofh.write("\nMutation regions should be a six column file.\n") 109 efh.write("\nThe selected mutations regions BED file is invalid, it should be a six column file.\n")
108 elif mutation_regions.shape[0] == 0: 110 elif mutation_regions.shape[0] == 0:
109 ofh.write("\nNo rows in mutation regions file.\n") 111 efh.write("\nThere are no rows in the selected mutation regions file.\n")
110 else: 112 else:
111 for region_i in range(mutation_regions.shape[0]): 113 for region_i in range(mutation_regions.shape[0]):
112 region = mutation_regions.iloc[region_i, :] 114 region = mutation_regions.iloc[region_i, :]
115 if region[0] not in reference:
116 efh.write("\nMutation region '%s' not found in reference genome.\n" % str(region))
117 break
118 if not isinstance(region[1], numpy.int64):
119 efh.write("\nNon-integer found in mutation region start (column 2): %s.\n" % str(region[1]))
120 break
121 if not isinstance(region[2], numpy.int64):
122 efh.write("\nNon-integer found in mutation region start (column 3): %s.\n" % str(region[2]))
123 break
124 if region[1] <= 0 or region[2] <= 0:
125 efh.write("\nMutation region '%s' starts before the reference sequence.\n" % str(region))
126 if region[1] > len(reference[region[0]].seq) or region[2] > len(reference[region[0]].seq):
127 efh.write("\nMutation region '%s' ends after the reference sequence.\n" % str(region))
113 if not region.get('type', default='No Type') in ['snp', 'small-indel', 'any']: 128 if not region.get('type', default='No Type') in ['snp', 'small-indel', 'any']:
129 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'))))
114 continue 130 continue
115 ofh.write("\nFinding AMR mutations for %s.\n" % str(region['name'])) 131 ofh.write("\nFinding AMR mutations for %s.\n" % str(region['name']))
116 region_bed = 'region_%s.bed' % region_i 132 region_bed = 'region_%s.bed' % region_i
117 mutation_regions.loc[[region_i], ].to_csv(path_or_buf=region_bed, sep='\t', header=False, index=False) 133 mutation_regions.loc[[region_i], ].to_csv(path_or_buf=region_bed, sep='\t', header=False, index=False)
118 region_mutations_tsv = os.path.join(mutation_regions_dir, 'region_%s_mutations.tsv' % region_i) 134 region_mutations_tsv = os.path.join(mutation_regions_dir, 'region_%s_mutations.tsv' % region_i)
165 axis.xaxis.set_label_position('top') 181 axis.xaxis.set_label_position('top')
166 pyplot.tight_layout() 182 pyplot.tight_layout()
167 pyplot.savefig(amr_matrix_png, dpi=300) 183 pyplot.savefig(amr_matrix_png, dpi=300)
168 else: 184 else:
169 ofh.write("\nEmpty AMR matrix, nothing to draw...\n") 185 ofh.write("\nEmpty AMR matrix, nothing to draw...\n")
186 efh.close()
170 ofh.close() 187 ofh.close()
171 188
172 189
173 if __name__ == '__main__': 190 if __name__ == '__main__':
174 parser = argparse.ArgumentParser() 191 parser = argparse.ArgumentParser()
179 parser.add_argument('--amr_mutation_regions_bed_file', action='store', dest='amr_mutation_regions_bed_file', default=None, help='AMR mutation regions BED file') 196 parser.add_argument('--amr_mutation_regions_bed_file', action='store', dest='amr_mutation_regions_bed_file', default=None, help='AMR mutation regions BED file')
180 parser.add_argument('--amr_gene_drug_file', action='store', dest='amr_gene_drug_file', help='AMR_gene_drugs tsv file') 197 parser.add_argument('--amr_gene_drug_file', action='store', dest='amr_gene_drug_file', help='AMR_gene_drugs tsv file')
181 parser.add_argument('--reference_genome', action='store', dest='reference_genome', help='Reference genome fasta file') 198 parser.add_argument('--reference_genome', action='store', dest='reference_genome', help='Reference genome fasta file')
182 parser.add_argument('--mutation_regions_dir', action='store', dest='mutation_regions_dir', help='Directory for mutation regions TSV files produced by this tool') 199 parser.add_argument('--mutation_regions_dir', action='store', dest='mutation_regions_dir', help='Directory for mutation regions TSV files produced by this tool')
183 parser.add_argument('--amr_matrix_png_dir', action='store', dest='amr_matrix_png_dir', help='Directory for PNG files produced by this tool') 200 parser.add_argument('--amr_matrix_png_dir', action='store', dest='amr_matrix_png_dir', help='Directory for PNG files produced by this tool')
201 parser.add_argument('--errors', action='store', dest='errors', help='Output file containing errors')
184 202
185 args = parser.parse_args() 203 args = parser.parse_args()
186 204
187 # Get the collection of feature hits files. The collection 205 # Get the collection of feature hits files. The collection
188 # will be sorted alphabetically and will contain 2 files 206 # will be sorted alphabetically and will contain 2 files
197 reference = load_fasta(args.reference_genome) 215 reference = load_fasta(args.reference_genome)
198 reference_size = 0 216 reference_size = 0
199 for i in reference: 217 for i in reference:
200 reference_size += len(i.seq) 218 reference_size += len(i.seq)
201 219
202 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) 220 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)