Mercurial > repos > greg > draw_amr_matrix

comparison draw_amr_matrix.py @ 7:389c98d344ce draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Uploaded
author greg
date Fri, 03 Mar 2023 22:03:09 +0000
parents caf554e039b2
children 7fe8ea50a81d
comparison
equal deleted inserted replaced
6:caf554e039b2 7:389c98d344ce
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, amr_mutations_file, amr_mutation_regions_file, amr_gene_drug_file, reference, reference_size, region_mutations_output_file, mutations_dir, output_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):
62 ofh = open('process_log', 'w') 62 ofh = open('process_log', 'w')
63 63
64 # Read amr_feature_hits_files. 64 # Read amr_feature_hits_files.
65 amr_feature_hits = pandas.Series(dtype=object) 65 amr_feature_hits = pandas.Series(dtype=object)
66 for amr_feature_hits_file in amr_feature_hits_files: 66 for amr_feature_hits_file in amr_feature_hits_files:
96 ofh.write("drugs: %s\n" % str(drugs)) 96 ofh.write("drugs: %s\n" % str(drugs))
97 for drug in drugs: 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)) 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 ofh.write("\amr_to_draw: %s\n" % str(amr_to_draw)) 99 ofh.write("\amr_to_draw: %s\n" % str(amr_to_draw))
100 100
101 ofh.write("\namr_mutations_file si None: %s\n" % str(amr_mutations_file == 'None')) 101 ofh.write("\nvarscan_vcf_file si None: %s\n" % str(varscan_vcf_file == 'None'))
102 if amr_mutations_file not in [None, 'None'] and os.path.getsize(amr_mutations_file) > 0: 102 if varscan_vcf_file not in [None, 'None'] and os.path.getsize(varscan_vcf_file) > 0:
103 amr_mutations = pandas.Series(dtype=object) 103 amr_mutations = pandas.Series(dtype=object)
104 if amr_mutation_regions_file is not None: 104 if amr_mutation_regions_bed_file is not None:
105 mutation_regions = pandas.read_csv(amr_mutation_regions_file, header=0, sep='\t', index_col=False) 105 mutation_regions = pandas.read_csv(amr_mutation_regions_bed_file, header=0, sep='\t', index_col=False)
106 if mutation_regions.shape[1] != 7: 106 if mutation_regions.shape[1] != 7:
107 ofh.write("\nMutation regions should be a six column file.\n") 107 ofh.write("\nMutation regions should be a six column file.\n")
108 elif mutation_regions.shape[0] == 0: 108 elif mutation_regions.shape[0] == 0:
109 ofh.write("\nNo rows in mutation regions file.\n") 109 ofh.write("\nNo rows in mutation regions file.\n")
110 else: 110 else:
111 """
112 TODO: move this coder to the pima_report tool...
111 # Make sure the positions in the BED file fall within the chromosomes provided in the reference sequence. 113 # 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]): 114 for mutation_region in range(mutation_regions.shape[0]):
113 mutation_region = mutation_regions.iloc[mutation_region, :] 115 mutation_region = mutation_regions.iloc[mutation_region, :]
114 if not (mutation_region[0] in reference): 116 if not (mutation_region[0] in reference):
115 ofh.write("\nMutation region :%s not found in reference genome.\n" % ' '.join(mutation_region.astype(str))) 117 ofh.write("\nMutation region: %s not found in reference genome.\n" % ' '.join(mutation_region.astype(str)))
116 continue 118 continue
117 if not isinstance(mutation_region[1], numpy.int64): 119 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])) 120 ofh.write("\nNon-integer found in mutation region start (column 2): %s.\n" % str(mutation_region[1]))
119 break 121 break
120 elif not isinstance(mutation_region[2], numpy.int64): 122 elif not isinstance(mutation_region[2], numpy.int64):
122 break 124 break
123 if mutation_region[1] <= 0 or mutation_region[2] <= 0: 125 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))) 126 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): 127 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))) 128 ofh.write("\nMutation region %s ends after the reference sequence.\n" % ' '.join(mutation_region.astype(str)))
129 """
127 for region_i in range(mutation_regions.shape[0]): 130 for region_i in range(mutation_regions.shape[0]):
128 region = mutation_regions.iloc[region_i, :] 131 region = mutation_regions.iloc[region_i, :]
129 if not region.get('type', default='No Type') in ['snp', 'small-indel', 'any']: 132 if not region.get('type', default='No Type') in ['snp', 'small-indel', 'any']:
130 continue 133 continue
131 ofh.write("\nFinding AMR mutations for %s.\n" % str(region['name'])) 134 ofh.write("\nFinding AMR mutations for %s.\n" % str(region['name']))
132 region_dir = os.path.join(mutations_dir, 'region_' + str(region_i)) 135 region_bed = 'region_%s.bed' % 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 mutation_regions.loc[[region_i], ].to_csv(path_or_buf=region_bed, sep='\t', header=False, index=False)
137 region_mutations_tsv = os.path.join(mutation_regions_dir, 'region_%s_mutations.tsv' % region_i)
136 cmd = ' '.join(['bedtools intersect -nonamecheck -wb -a', 138 cmd = ' '.join(['bedtools intersect -nonamecheck -wb -a',
137 region_bed, 139 region_bed,
138 '-b', 140 '-b',
139 amr_mutations_file, 141 varscan_vcf_file,
140 ' | awk \'BEGIN{getline < "' + amr_mutation_regions_file + '";printf $0"\\t";', 142 ' | awk \'BEGIN{getline < "' + amr_mutation_regions_bed_file + '";printf $0"\\t";',
141 'getline < "' + amr_mutations_file + '"; getline < "' + amr_mutations_file + '";print $0}{print}\'', 143 'getline < "' + varscan_vcf_file + '"; getline < "' + varscan_vcf_file + '";print $0}{print}\'',
142 '1>' + region_mutations_output_file]) 144 '1>' + region_mutations_tsv])
143 ofh.write("\ncmd:\n%s\n" % cmd) 145 ofh.write("\ncmd:\n%s\n" % cmd)
144 run_command(cmd) 146 run_command(cmd)
147 """
148 TODO: move this coder to the pima_report tool...
145 try: 149 try:
146 region_mutations = pandas.read_csv(region_mutations_output_file, sep='\t', header=0, index_col=False) 150 region_mutations = pandas.read_csv(region_mutations_tsv, sep='\t', header=0, index_col=False)
147 except Exception: 151 except Exception:
148 region_mutations = pandas.DataFrame() 152 region_mutations = pandas.DataFrame()
149 if region_mutations.shape[0] == 0: 153 if region_mutations.shape[0] == 0:
150 continue 154 continue
151 # Figure out what kind of mutations are in this region. 155 # Figure out what kind of mutations are in this region.
154 region_mutation_drugs = pandas.Series(region['drug'] * region_mutations.shape[0], name='DRUG', index=region_mutations.index) 158 region_mutation_drugs = pandas.Series(region['drug'] * region_mutations.shape[0], name='DRUG', index=region_mutations.index)
155 region_notes = pandas.Series(region['note'] * region_mutations.shape[0], name='NOTE', index=region_mutations.index) 159 region_notes = pandas.Series(region['note'] * region_mutations.shape[0], name='NOTE', index=region_mutations.index)
156 region_mutations = pandas.concat([region_mutations, region_mutation_types, region_mutation_drugs, region_notes], axis=1) 160 region_mutations = pandas.concat([region_mutations, region_mutation_types, region_mutation_drugs, region_notes], axis=1)
157 region_mutations = region_mutations[['#CHROM', 'POS', 'TYPE', 'REF', 'ALT', 'DRUG', 'NOTE']] 161 region_mutations = region_mutations[['#CHROM', 'POS', 'TYPE', 'REF', 'ALT', 'DRUG', 'NOTE']]
158 amr_mutations[region['name']] = region_mutations 162 amr_mutations[region['name']] = region_mutations
163 """
159 else: 164 else:
160 ofh.write("\nMutation region BED not received.\n") 165 ofh.write("\nMutation region BED not received.\n")
161 # Roll up potentially resistance conferring mutations. 166 # Roll up potentially resistance conferring mutations.
162 for mutation_region, mutation_hits in amr_mutations.iteritems(): 167 for mutation_region, mutation_hits in amr_mutations.iteritems():
163 for mutation_idx, mutation_hit in mutation_hits.iterrows(): 168 for mutation_idx, mutation_hit in mutation_hits.iterrows():
181 present_genes = amr_to_draw['gene'].unique() 186 present_genes = amr_to_draw['gene'].unique()
182 present_drugs = amr_to_draw['drug'].unique() 187 present_drugs = amr_to_draw['drug'].unique()
183 amr_matrix = pandas.DataFrame(0, index=present_genes, columns=present_drugs) 188 amr_matrix = pandas.DataFrame(0, index=present_genes, columns=present_drugs)
184 for hit_idx, hit in amr_to_draw.iterrows(): 189 for hit_idx, hit in amr_to_draw.iterrows():
185 amr_matrix.loc[hit[0], hit[1]] = 1 190 amr_matrix.loc[hit[0], hit[1]] = 1
186 amr_matrix_png = os.path.join(output_dir, 'amr_matrix.png') 191 amr_matrix_png = os.path.join(amr_matrix_png_dir, 'amr_matrix.png')
187 int_matrix = amr_matrix[amr_matrix.columns].astype(int) 192 int_matrix = amr_matrix[amr_matrix.columns].astype(int)
188 figure, axis = pyplot.subplots() 193 figure, axis = pyplot.subplots()
189 heatmap = axis.pcolor(int_matrix, cmap=pyplot.cm.Blues, linewidth=0) 194 heatmap = axis.pcolor(int_matrix, cmap=pyplot.cm.Blues, linewidth=0)
190 axis.invert_yaxis() 195 axis.invert_yaxis()
191 axis.set_yticks(numpy.arange(0.5, len(amr_matrix.index)), minor=False) 196 axis.set_yticks(numpy.arange(0.5, len(amr_matrix.index)), minor=False)
204 if __name__ == '__main__': 209 if __name__ == '__main__':
205 parser = argparse.ArgumentParser() 210 parser = argparse.ArgumentParser()
206 211
207 parser.add_argument('--amr_feature_hits_dir', action='store', dest='amr_feature_hits_dir', help='Directory of tabular files containing feature hits') 212 parser.add_argument('--amr_feature_hits_dir', action='store', dest='amr_feature_hits_dir', help='Directory of tabular files containing feature hits')
208 parser.add_argument('--amr_deletions_file', action='store', dest='amr_deletions_file', default=None, help='AMR deletions BED file') 213 parser.add_argument('--amr_deletions_file', action='store', dest='amr_deletions_file', default=None, help='AMR deletions BED file')
209 parser.add_argument('--amr_mutations_file', action='store', dest='amr_mutations_file', default=None, help='AMR mutations TSV file') 214 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')
210 parser.add_argument('--amr_mutation_regions_file', action='store', dest='amr_mutation_regions_file', default=None, help='AMR mutation regions BED file') 215 parser.add_argument('--amr_mutation_regions_bed_file', action='store', dest='amr_mutation_regions_bed_file', default=None, help='AMR mutation regions BED file')
211 parser.add_argument('--amr_gene_drug_file', action='store', dest='amr_gene_drug_file', help='AMR_gene_drugs tsv file') 216 parser.add_argument('--amr_gene_drug_file', action='store', dest='amr_gene_drug_file', help='AMR_gene_drugs tsv file')
212 parser.add_argument('--reference_genome', action='store', dest='reference_genome', help='Reference genome fasta file') 217 parser.add_argument('--reference_genome', action='store', dest='reference_genome', help='Reference genome fasta file')
213 parser.add_argument('--region_mutations_output_file', action='store', dest='region_mutations_output_file', default=None, help='Region mutations TSV output file') 218 parser.add_argument('--mutation_regions_dir', action='store', dest='mutation_regions_dir', help='Directory for mutation regions TSV files produced by this tool')
214 parser.add_argument('--mutations_dir', action='store', dest='mutations_dir', help='Mutations directory') 219 parser.add_argument('--amr_matrix_png_dir', action='store', dest='amr_matrix_png_dir', help='Directory for PNG files produced by this tool')
215 parser.add_argument('--output_dir', action='store', dest='output_dir', help='Output directory')
216 220
217 args = parser.parse_args() 221 args = parser.parse_args()
218 222
219 # Get the collection of feature hits files. The collection 223 # Get the collection of feature hits files. The collection
220 # will be sorted alphabetically and will contain 2 files 224 # will be sorted alphabetically and will contain 2 files
229 reference = load_fasta(args.reference_genome) 233 reference = load_fasta(args.reference_genome)
230 reference_size = 0 234 reference_size = 0
231 for i in reference: 235 for i in reference:
232 reference_size += len(i.seq) 236 reference_size += len(i.seq)
233 237
234 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) 238 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)