Mercurial > repos > iuc > vsnp_add_zero_coverage

comparison vsnp_build_tables.py @ 3:2e863710a2f0 draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
"planemo upload for repository https://github.com/galaxyproject/tools-iuc/tree/master/tools/vsnp commit 2e312886647244b416c64eca91e1a61dd1be939b"
author iuc
date Thu, 10 Dec 2020 15:26:25 +0000
parents aed013f6b13b
children 6dc6dd4666e3
comparison
equal deleted inserted replaced
2:4cc004985e27 3:2e863710a2f0
1 #!/usr/bin/env python 1 #!/usr/bin/env python
2 2
3 import argparse 3 import argparse
4 import multiprocessing
5 import os 4 import os
6 import queue
7 import re 5 import re
8 6
9 import pandas 7 import pandas
10 import pandas.io.formats.excel 8 import pandas.io.formats.excel
11 from Bio import SeqIO 9 from Bio import SeqIO
12 10
13 INPUT_JSON_AVG_MQ_DIR = 'input_json_avg_mq_dir'
14 INPUT_JSON_DIR = 'input_json_dir'
15 INPUT_NEWICK_DIR = 'input_newick_dir'
16 # Maximum columns allowed in a LibreOffice 11 # Maximum columns allowed in a LibreOffice
17 # spreadsheet is 1024. Excel allows for 12 # spreadsheet is 1024. Excel allows for
18 # 16,384 columns, but we'll set the lower 13 # 16,384 columns, but we'll set the lower
19 # number as the maximum. Some browsers 14 # number as the maximum. Some browsers
20 # (e.g., Firefox on Linux) are configured 15 # (e.g., Firefox on Linux) are configured
143 pro.index = pandas.IntervalIndex.from_arrays(pro['start'], pro['stop'], closed='both') 138 pro.index = pandas.IntervalIndex.from_arrays(pro['start'], pro['stop'], closed='both')
144 annotation_dict[chromosome] = pro 139 annotation_dict[chromosome] = pro
145 return annotation_dict 140 return annotation_dict
146 141
147 142
148 def get_base_file_name(file_path): 143 def get_sample_name(file_path):
149 base_file_name = os.path.basename(file_path) 144 base_file_name = os.path.basename(file_path)
150 if base_file_name.find(".") > 0: 145 if base_file_name.find(".") > 0:
151 # Eliminate the extension. 146 # Eliminate the extension.
152 return os.path.splitext(base_file_name)[0] 147 return os.path.splitext(base_file_name)[0]
153 elif base_file_name.find("_") > 0: 148 return base_file_name
154 # The dot extension was likely changed to
155 # the " character.
156 items = base_file_name.split("_")
157 return "_".join(items[0:-1])
158 else:
159 return base_file_name
160 149
161 150
162 def output_cascade_table(cascade_order, mqdf, group, annotation_dict): 151 def output_cascade_table(cascade_order, mqdf, group, annotation_dict):
163 cascade_order_mq = pandas.concat([cascade_order, mqdf], join='inner') 152 cascade_order_mq = pandas.concat([cascade_order, mqdf], join='inner')
164 output_table(cascade_order_mq, "cascade", group, annotation_dict) 153 output_table(cascade_order_mq, "cascade", group, annotation_dict)
167 def output_excel(df, type_str, group, annotation_dict, count=None): 156 def output_excel(df, type_str, group, annotation_dict, count=None):
168 # Output the temporary json file that 157 # Output the temporary json file that
169 # is used by the excel_formatter. 158 # is used by the excel_formatter.
170 if count is None: 159 if count is None:
171 if group is None: 160 if group is None:
172 json_file_name = "%s_order_mq.json" % type_str 161 json_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_order_mq.json" % type_str)
173 excel_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_table.xlsx" % type_str) 162 excel_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_table.xlsx" % type_str)
174 else: 163 else:
175 json_file_name = "%s_%s_order_mq.json" % (group, type_str) 164 json_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_%s_order_mq.json" % (group, type_str))
176 excel_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_%s_table.xlsx" % (group, type_str)) 165 excel_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_%s_table.xlsx" % (group, type_str))
177 else: 166 else:
167 # The table has more columns than is allowed by the
168 # MAXCOLS setting, so multiple files will be produced
169 # as an output collection.
178 if group is None: 170 if group is None:
179 json_file_name = "%s_order_mq_%d.json" % (type_str, count) 171 json_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_order_mq_%d.json" % (type_str, count))
180 excel_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_table_%d.xlsx" % (type_str, count)) 172 excel_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_table_%d.xlsx" % (type_str, count))
181 else: 173 else:
182 json_file_name = "%s_%s_order_mq_%d.json" % (group, type_str, count) 174 json_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_%s_order_mq_%d.json" % (group, type_str, count))
183 excel_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_%s_table_%d.xlsx" % (group, type_str, count)) 175 excel_file_name = os.path.join(OUTPUT_EXCEL_DIR, "%s_%s_table_%d.xlsx" % (group, type_str, count))
184 df.to_json(json_file_name, orient='split') 176 df.to_json(json_file_name, orient='split')
185 # Output the Excel file. 177 # Output the Excel file.
186 excel_formatter(json_file_name, excel_file_name, group, annotation_dict) 178 excel_formatter(json_file_name, excel_file_name, group, annotation_dict)
187 179
227 output_excel(df_of_type, type_str, group_str, annotation_dict, count=count) 219 output_excel(df_of_type, type_str, group_str, annotation_dict, count=count)
228 else: 220 else:
229 output_excel(df, type_str, group_str, annotation_dict) 221 output_excel(df, type_str, group_str, annotation_dict)
230 222
231 223
232 def preprocess_tables(task_queue, annotation_dict, timeout): 224 def preprocess_tables(newick_file, json_file, json_avg_mq_file, annotation_dict):
233 while True: 225 avg_mq_series = pandas.read_json(json_avg_mq_file, typ='series', orient='split')
234 try: 226 # Map quality to dataframe.
235 tup = task_queue.get(block=True, timeout=timeout) 227 mqdf = avg_mq_series.to_frame(name='MQ')
236 except queue.Empty: 228 mqdf = mqdf.T
237 break 229 # Get the group.
238 newick_file, json_file, json_avg_mq_file = tup 230 group = get_sample_name(newick_file)
239 avg_mq_series = pandas.read_json(json_avg_mq_file, typ='series', orient='split') 231 snps_df = pandas.read_json(json_file, orient='split')
240 # Map quality to dataframe. 232 with open(newick_file, 'r') as fh:
241 mqdf = avg_mq_series.to_frame(name='MQ') 233 for line in fh:
242 mqdf = mqdf.T 234 line = re.sub('[:,]', '\n', line)
243 # Get the group. 235 line = re.sub('[)(]', '', line)
244 group = get_base_file_name(newick_file) 236 line = re.sub(r'[0-9].*\.[0-9].*\n', '', line)
245 snps_df = pandas.read_json(json_file, orient='split') 237 line = re.sub('root\n', '', line)
246 with open(newick_file, 'r') as fh: 238 sample_order = line.split('\n')
247 for line in fh: 239 sample_order = list([_f for _f in sample_order if _f])
248 line = re.sub('[:,]', '\n', line) 240 sample_order.insert(0, 'root')
249 line = re.sub('[)(]', '', line) 241 tree_order = snps_df.loc[sample_order]
250 line = re.sub(r'[0-9].*\.[0-9].*\n', '', line) 242 # Count number of SNPs in each column.
251 line = re.sub('root\n', '', line) 243 snp_per_column = []
252 sample_order = line.split('\n') 244 for column_header in tree_order:
253 sample_order = list([_f for _f in sample_order if _f]) 245 count = 0
254 sample_order.insert(0, 'root') 246 column = tree_order[column_header]
255 tree_order = snps_df.loc[sample_order] 247 for element in column:
256 # Count number of SNPs in each column. 248 if element != column[0]:
257 snp_per_column = [] 249 count = count + 1
258 for column_header in tree_order: 250 snp_per_column.append(count)
259 count = 0 251 row1 = pandas.Series(snp_per_column, tree_order.columns, name="snp_per_column")
260 column = tree_order[column_header] 252 # Count number of SNPS from the
261 for element in column: 253 # top of each column in the table.
262 if element != column[0]: 254 snp_from_top = []
263 count = count + 1 255 for column_header in tree_order:
264 snp_per_column.append(count) 256 count = 0
265 row1 = pandas.Series(snp_per_column, tree_order.columns, name="snp_per_column") 257 column = tree_order[column_header]
266 # Count number of SNPS from the 258 # for each element in the column
267 # top of each column in the table. 259 # skip the first element
268 snp_from_top = [] 260 for element in column[1:]:
269 for column_header in tree_order: 261 if element == column[0]:
270 count = 0 262 count = count + 1
271 column = tree_order[column_header] 263 else:
272 # for each element in the column 264 break
273 # skip the first element 265 snp_from_top.append(count)
274 for element in column[1:]: 266 row2 = pandas.Series(snp_from_top, tree_order.columns, name="snp_from_top")
275 if element == column[0]: 267 tree_order = tree_order.append([row1])
276 count = count + 1 268 tree_order = tree_order.append([row2])
277 else: 269 # In pandas=0.18.1 even this does not work:
278 break 270 # abc = row1.to_frame()
279 snp_from_top.append(count) 271 # abc = abc.T --> tree_order.shape (5, 18), abc.shape (1, 18)
280 row2 = pandas.Series(snp_from_top, tree_order.columns, name="snp_from_top") 272 # tree_order.append(abc)
281 tree_order = tree_order.append([row1]) 273 # Continue to get error: "*** ValueError: all the input arrays must have same number of dimensions"
282 tree_order = tree_order.append([row2]) 274 tree_order = tree_order.T
283 # In pandas=0.18.1 even this does not work: 275 tree_order = tree_order.sort_values(['snp_from_top', 'snp_per_column'], ascending=[True, False])
284 # abc = row1.to_frame() 276 tree_order = tree_order.T
285 # abc = abc.T --> tree_order.shape (5, 18), abc.shape (1, 18) 277 # Remove snp_per_column and snp_from_top rows.
286 # tree_order.append(abc) 278 cascade_order = tree_order[:-2]
287 # Continue to get error: "*** ValueError: all the input arrays must have same number of dimensions" 279 # Output the cascade table.
288 tree_order = tree_order.T 280 output_cascade_table(cascade_order, mqdf, group, annotation_dict)
289 tree_order = tree_order.sort_values(['snp_from_top', 'snp_per_column'], ascending=[True, False]) 281 # Output the sorted table.
290 tree_order = tree_order.T 282 output_sort_table(cascade_order, mqdf, group, annotation_dict)
291 # Remove snp_per_column and snp_from_top rows.
292 cascade_order = tree_order[:-2]
293 # Output the cascade table.
294 output_cascade_table(cascade_order, mqdf, group, annotation_dict)
295 # Output the sorted table.
296 output_sort_table(cascade_order, mqdf, group, annotation_dict)
297 task_queue.task_done()
298
299
300 def set_num_cpus(num_files, processes):
301 num_cpus = int(multiprocessing.cpu_count())
302 if num_files < num_cpus and num_files < processes:
303 return num_files
304 if num_cpus < processes:
305 half_cpus = int(num_cpus / 2)
306 if num_files < half_cpus:
307 return num_files
308 return half_cpus
309 return processes
310 283
311 284
312 if __name__ == '__main__': 285 if __name__ == '__main__':
313 parser = argparse.ArgumentParser() 286 parser = argparse.ArgumentParser()
314 287
315 parser.add_argument('--input_avg_mq_json', action='store', dest='input_avg_mq_json', required=False, default=None, help='Average MQ json file')
316 parser.add_argument('--input_newick', action='store', dest='input_newick', required=False, default=None, help='Newick file')
317 parser.add_argument('--input_snps_json', action='store', dest='input_snps_json', required=False, default=None, help='SNPs json file')
318 parser.add_argument('--gbk_file', action='store', dest='gbk_file', required=False, default=None, help='Optional gbk file'), 288 parser.add_argument('--gbk_file', action='store', dest='gbk_file', required=False, default=None, help='Optional gbk file'),
319 parser.add_argument('--processes', action='store', dest='processes', type=int, help='User-selected number of processes to use for job splitting') 289 parser.add_argument('--input_avg_mq_json', action='store', dest='input_avg_mq_json', help='Average MQ json file')
290 parser.add_argument('--input_newick', action='store', dest='input_newick', help='Newick file')
291 parser.add_argument('--input_snps_json', action='store', dest='input_snps_json', help='SNPs json file')
320 292
321 args = parser.parse_args() 293 args = parser.parse_args()
322 294
323 if args.gbk_file is not None: 295 if args.gbk_file is not None:
324 # Create the annotation_dict for annotating 296 # Create the annotation_dict for annotating
325 # the Excel tables. 297 # the Excel tables.
326 annotation_dict = get_annotation_dict(args.gbk_file) 298 annotation_dict = get_annotation_dict(args.gbk_file)
327 else: 299 else:
328 annotation_dict = None 300 annotation_dict = None
329 301
330 # The assumption here is that the list of files 302 preprocess_tables(args.input_newick, args.input_snps_json, args.input_avg_mq_json, annotation_dict)
331 # in both INPUT_NEWICK_DIR and INPUT_JSON_DIR are
332 # named such that they are properly matched if
333 # the directories contain more than 1 file (i.e.,
334 # hopefully the newick file names and json file names
335 # will be something like Mbovis-01D6_* so they can be
336 # sorted and properly associated with each other).
337 if args.input_newick is not None:
338 newick_files = [args.input_newick]
339 else:
340 newick_files = []
341 for file_name in sorted(os.listdir(INPUT_NEWICK_DIR)):
342 file_path = os.path.abspath(os.path.join(INPUT_NEWICK_DIR, file_name))
343 newick_files.append(file_path)
344 if args.input_snps_json is not None:
345 json_files = [args.input_snps_json]
346 else:
347 json_files = []
348 for file_name in sorted(os.listdir(INPUT_JSON_DIR)):
349 file_path = os.path.abspath(os.path.join(INPUT_JSON_DIR, file_name))
350 json_files.append(file_path)
351 if args.input_avg_mq_json is not None:
352 json_avg_mq_files = [args.input_avg_mq_json]
353 else:
354 json_avg_mq_files = []
355 for file_name in sorted(os.listdir(INPUT_JSON_AVG_MQ_DIR)):
356 file_path = os.path.abspath(os.path.join(INPUT_JSON_AVG_MQ_DIR, file_name))
357 json_avg_mq_files.append(file_path)
358
359 multiprocessing.set_start_method('spawn')
360 queue1 = multiprocessing.JoinableQueue()
361 queue2 = multiprocessing.JoinableQueue()
362 num_files = len(newick_files)
363 cpus = set_num_cpus(num_files, args.processes)
364 # Set a timeout for get()s in the queue.
365 timeout = 0.05
366
367 for i, newick_file in enumerate(newick_files):
368 json_file = json_files[i]
369 json_avg_mq_file = json_avg_mq_files[i]
370 queue1.put((newick_file, json_file, json_avg_mq_file))
371
372 # Complete the preprocess_tables task.
373 processes = [multiprocessing.Process(target=preprocess_tables, args=(queue1, annotation_dict, timeout, )) for _ in range(cpus)]
374 for p in processes:
375 p.start()
376 for p in processes:
377 p.join()
378 queue1.join()
379
380 if queue1.empty():
381 queue1.close()
382 queue1.join_thread()