kkonganti@1: #!/usr/bin/env python3 kkonganti@1: kkonganti@1: # Kranti Konganti kkonganti@1: kkonganti@1: import argparse kkonganti@1: import glob kkonganti@1: import inspect kkonganti@1: import json kkonganti@1: import logging kkonganti@1: import os kkonganti@1: import pickle kkonganti@1: import pprint kkonganti@1: import re kkonganti@1: from collections import defaultdict kkonganti@1: kkonganti@1: import yaml kkonganti@1: kkonganti@1: kkonganti@1: # Multiple inheritence for pretty printing of help text. kkonganti@1: class MultiArgFormatClasses(argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter): kkonganti@1: pass kkonganti@1: kkonganti@1: kkonganti@1: # Main kkonganti@1: def main() -> None: kkonganti@1: """ kkonganti@1: The succesful execution of this script requires access to bettercallsal formatted kkonganti@1: db flat files. On raven2, they are at /hpc/db/bettercallsall/PDGXXXXXXXXXX.XXXXX kkonganti@1: kkonganti@1: It takes the ACC2SERO.pickle file and *.reference_target.cluster_list.tsv file kkonganti@1: for that particular NCBI Pathogens release from the db directory mentioned with kkonganti@1: -db option and a root parent directory of the `salmon quant` results mentioned kkonganti@1: with -sal option and generates a final results table with number of reads kkonganti@1: mapped and a .json file to be used with MultiQC to generate a stacked bar plot. kkonganti@1: kkonganti@1: Using -url option optionally adds an extra column of NCBI Pathogens Isolates kkonganti@1: Browser, which directly links out to NCBI Pathogens Isolates SNP viewer tool. kkonganti@1: """ kkonganti@1: # Set logging. kkonganti@1: logging.basicConfig( kkonganti@1: format="\n" + "=" * 55 + "\n%(asctime)s - %(levelname)s\n" + "=" * 55 + "\n%(message)s\n\n", kkonganti@1: level=logging.DEBUG, kkonganti@1: ) kkonganti@1: kkonganti@1: # Debug print. kkonganti@1: ppp = pprint.PrettyPrinter(width=55) kkonganti@1: prog_name = inspect.stack()[0].filename kkonganti@1: kkonganti@1: parser = argparse.ArgumentParser( kkonganti@1: prog=prog_name, description=main.__doc__, formatter_class=MultiArgFormatClasses kkonganti@1: ) kkonganti@1: kkonganti@1: required = parser.add_argument_group("required arguments") kkonganti@1: kkonganti@1: required.add_argument( kkonganti@1: "-sal", kkonganti@1: dest="salmon_res_dir", kkonganti@1: default=False, kkonganti@1: required=True, kkonganti@1: help="Absolute UNIX path to the parent directory that contains the\n" kkonganti@1: + "`salmon quant` results directory. For example, if path to\n" kkonganti@1: + "`quant.sf` is in /hpc/john_doe/test/salmon_res/quant.sf, then\n" kkonganti@1: + "use this command-line option as:\n" kkonganti@1: + "-sal /hpc/john_doe/test", kkonganti@1: ) kkonganti@1: required.add_argument( kkonganti@1: "-snp", kkonganti@1: dest="rtc", kkonganti@1: default=False, kkonganti@1: required=True, kkonganti@1: help="Absolute UNIX Path to the PDG SNP reference target cluster\n" kkonganti@1: + "metadata file. On raven2, these are located at\n" kkonganti@1: + "/hpc/db/bettercallsal/PDGXXXXXXXXXX.XXXXX\n" kkonganti@1: + "Required if -sal is on.", kkonganti@1: ) kkonganti@1: required.add_argument( kkonganti@1: "-pickle", kkonganti@1: dest="acc2sero", kkonganti@1: default=False, kkonganti@1: required=True, kkonganti@1: help="Absolute UNIX Path to the *ACC2SERO.pickle\n" kkonganti@1: + "metadata file. On raven2, these are located at\n" kkonganti@1: + "/hpc/db/bettercallsal/PDGXXXXXXXXXX.XXXXX\n" kkonganti@1: + "Required if -sal is on.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-op", kkonganti@1: dest="out_prefix", kkonganti@1: default="bettercallsal.tblsum", kkonganti@1: required=False, kkonganti@1: help="Set the output file(s) prefix for output(s) generated\n" + "by this program.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-url", kkonganti@1: dest="show_snp_clust_info", kkonganti@1: default=False, kkonganti@1: required=False, kkonganti@1: action="store_true", kkonganti@1: help="Show SNP cluster participation information of the final genome hit.\n" kkonganti@1: + "This may be useful to see a relative placement of your sample in\n" kkonganti@1: + "NCBI Isolates SNP Tree Viewer based on genome similarity but however\n" kkonganti@1: + "due to rapid nature of the updates at NCBI Pathogen Detection Project,\n" kkonganti@1: + "the placement may be in an outdated cluster.", kkonganti@1: ) kkonganti@1: kkonganti@1: args = parser.parse_args() kkonganti@1: salmon_res_dir = args.salmon_res_dir kkonganti@1: out_prefix = args.out_prefix kkonganti@1: show_snp_clust_col = args.show_snp_clust_info kkonganti@1: rtc = args.rtc kkonganti@1: pickled_sero = args.acc2sero kkonganti@1: no_hit = "No genome hit" kkonganti@1: bcs_sal_yn_prefix = "bettercallsal_salyn" kkonganti@1: sal_y = "Detected" kkonganti@1: sal_n = "Not detected" kkonganti@1: ncbi_pathogens_base_url = "https://www.ncbi.nlm.nih.gov/pathogens/" kkonganti@1: kkonganti@1: sample2salmon, snp_clusters, multiqc_salmon_counts, seen_sero, sal_yn = ( kkonganti@1: defaultdict(defaultdict), kkonganti@1: defaultdict(defaultdict), kkonganti@1: defaultdict(defaultdict), kkonganti@1: defaultdict(int), kkonganti@1: defaultdict(int), kkonganti@1: ) kkonganti@1: kkonganti@1: cell_colors_yml = { kkonganti@1: bcs_sal_yn_prefix: {sal_y: "#c8e6c9 !important;", sal_n: "#ffcdd2 !important;"} kkonganti@1: } kkonganti@1: kkonganti@1: salmon_comb_res = os.path.join(os.getcwd(), out_prefix + ".txt") kkonganti@1: bcs_sal_yn = re.sub(out_prefix, bcs_sal_yn_prefix + ".tblsum", salmon_comb_res) kkonganti@1: cell_colors_yml_file = re.sub( kkonganti@1: out_prefix + ".txt", bcs_sal_yn_prefix + ".cellcolors.yml", salmon_comb_res kkonganti@1: ) kkonganti@1: salmon_comb_res_mqc = os.path.join(os.getcwd(), str(out_prefix).split(".")[0] + "_mqc.json") kkonganti@1: salmon_res_files = glob.glob(os.path.join(salmon_res_dir, "*", "quant.sf"), recursive=True) kkonganti@1: salmon_res_file_failed = glob.glob(os.path.join(salmon_res_dir, "BCS_NO_CALLS.txt")) kkonganti@1: kkonganti@1: if rtc and (not os.path.exists(rtc) or not os.path.getsize(rtc) > 0): kkonganti@1: logging.error( kkonganti@1: "The reference target cluster metadata file,\n" kkonganti@1: + f"{os.path.basename(rtc)} does not exist or is empty!" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if rtc and (not salmon_res_dir or not pickled_sero): kkonganti@1: logging.error("When -rtc is on, -sal and -ps are also required.") kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if pickled_sero and (not os.path.exists(pickled_sero) or not os.path.getsize(pickled_sero)): kkonganti@1: logging.error( kkonganti@1: "The pickle file,\n" + f"{os.path.basename(pickled_sero)} does not exist or is empty!" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if salmon_res_dir: kkonganti@1: if not os.path.isdir(salmon_res_dir): kkonganti@1: logging.error("UNIX path\n" + f"{salmon_res_dir}\n" + "does not exist!") kkonganti@1: exit(1) kkonganti@1: if len(salmon_res_files) <= 0: kkonganti@1: # logging.error( kkonganti@1: # "Parent directory,\n" kkonganti@1: # + f"{salmon_res_dir}" kkonganti@1: # + "\ndoes not seem to have any directories that contain\n" kkonganti@1: # + "the `quant.sf` file(s)." kkonganti@1: # ) kkonganti@1: # exit(1) kkonganti@1: with open(salmon_comb_res, "w") as salmon_comb_res_fh: kkonganti@1: salmon_comb_res_fh.write(f"Sample\n{no_hit}s in any samples\n") kkonganti@1: salmon_comb_res_fh.close() kkonganti@1: exit(0) kkonganti@1: kkonganti@1: if rtc and os.path.exists(rtc) and os.path.getsize(rtc) > 0: kkonganti@1: kkonganti@1: # pdg_release = re.match(r"(^PDG\d+\.\d+)\..+", os.path.basename(rtc))[1] + "/" kkonganti@1: acc2sero = pickle.load(file=open(pickled_sero, "rb")) kkonganti@1: kkonganti@1: with open(rtc, "r") as rtc_fh: kkonganti@1: kkonganti@1: for line in rtc_fh: kkonganti@1: cols = line.strip().split("\t") kkonganti@1: kkonganti@1: if len(cols) < 4: kkonganti@1: logging.error( kkonganti@1: f"The file {os.path.basename(rtc)} seems to\n" kkonganti@1: + "be malformed. It contains less than required 4 columns." kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: elif cols[3] != "NULL": kkonganti@1: snp_clusters[cols[0]].setdefault("assembly_accs", []).append(cols[3]) kkonganti@1: snp_clusters[cols[3]].setdefault("snp_clust_id", []).append(cols[0]) kkonganti@1: snp_clusters[cols[3]].setdefault("pathdb_acc_id", []).append(cols[1]) kkonganti@1: if len(snp_clusters[cols[3]]["snp_clust_id"]) > 1: kkonganti@1: logging.error( kkonganti@1: f"There is a duplicate reference accession [{cols[3]}]" kkonganti@1: + f"in the metadata file{os.path.basename(rtc)}!" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: rtc_fh.close() kkonganti@1: kkonganti@1: for salmon_res_file in salmon_res_files: kkonganti@1: sample_name = re.match( kkonganti@1: r"(^.+?)((\_salmon\_res)|(\.salmon))$", kkonganti@1: os.path.basename(os.path.dirname(salmon_res_file)), kkonganti@1: )[1] kkonganti@1: salmon_meta_json = os.path.join( kkonganti@1: os.path.dirname(salmon_res_file), "aux_info", "meta_info.json" kkonganti@1: ) kkonganti@1: kkonganti@1: if not os.path.exists(salmon_meta_json) or not os.path.getsize(salmon_meta_json) > 0: kkonganti@1: logging.error( kkonganti@1: "The file\n" kkonganti@1: + f"{salmon_meta_json}\ndoes not exist or is empty!\n" kkonganti@1: + "Did `salmon quant` fail?" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if not os.path.exists(salmon_res_file) or not os.path.getsize(salmon_res_file): kkonganti@1: logging.error( kkonganti@1: "The file\n" kkonganti@1: + f"{salmon_res_file}\ndoes not exist or is empty!\n" kkonganti@1: + "Did `salmon quant` fail?" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: with open(salmon_res_file, "r") as salmon_res_fh: kkonganti@1: for line in salmon_res_fh.readlines(): kkonganti@1: if re.match(r"^Name.+", line): kkonganti@1: continue kkonganti@1: cols = line.strip().split("\t") kkonganti@1: ref_acc = "_".join(cols[0].split("_")[:2]) kkonganti@1: ( kkonganti@1: sample2salmon[sample_name] kkonganti@1: .setdefault(acc2sero[cols[0]], []) kkonganti@1: .append(int(round(float(cols[4]), 2))) kkonganti@1: ) kkonganti@1: ( kkonganti@1: sample2salmon[sample_name] kkonganti@1: .setdefault("snp_clust_ids", {}) kkonganti@1: .setdefault("".join(snp_clusters[ref_acc]["snp_clust_id"]), []) kkonganti@1: .append("".join(snp_clusters[ref_acc]["pathdb_acc_id"])) kkonganti@1: ) kkonganti@1: seen_sero[acc2sero[cols[0]]] = 1 kkonganti@1: kkonganti@1: salmon_meta_json_read = json.load(open(salmon_meta_json, "r")) kkonganti@1: ( kkonganti@1: sample2salmon[sample_name] kkonganti@1: .setdefault("tot_reads", []) kkonganti@1: .append(salmon_meta_json_read["num_processed"]) kkonganti@1: ) kkonganti@1: kkonganti@1: with open(salmon_comb_res, "w") as salmon_comb_res_fh: kkonganti@1: kkonganti@1: # snp_clust_col_header = ( kkonganti@1: # "\tSNP Cluster(s) by Genome Hit\n" if show_snp_clust_col else "\n" kkonganti@1: # ) kkonganti@1: snp_clust_col_header = ( kkonganti@1: "\tNCBI Pathogens Isolate Browser\n" if show_snp_clust_col else "\n" kkonganti@1: ) kkonganti@1: serotypes = sorted(seen_sero.keys()) kkonganti@1: formatted_serotypes = [ kkonganti@1: re.sub(r"\,antigen_formula=", " | ", s) kkonganti@1: for s in [re.sub(r"serotype=", "", s) for s in serotypes] kkonganti@1: ] kkonganti@1: salmon_comb_res_fh.write( kkonganti@1: "Sample\t" + "\t".join(formatted_serotypes) + snp_clust_col_header kkonganti@1: ) kkonganti@1: # sample_snp_relation = ( kkonganti@1: # ncbi_pathogens_base_url kkonganti@1: # + pdg_release kkonganti@1: # + "".join(snp_clusters[ref_acc]["snp_clust_id"]) kkonganti@1: # + "?accessions=" kkonganti@1: # ) kkonganti@1: sample_snp_relation = ncbi_pathogens_base_url + "isolates/#" kkonganti@1: kkonganti@1: if len(salmon_res_file_failed) == 1: kkonganti@1: with (open("".join(salmon_res_file_failed), "r")) as no_calls_fh: kkonganti@1: for line in no_calls_fh.readlines(): kkonganti@1: if line in ["\n", "\n\r", "\r"]: kkonganti@1: continue kkonganti@1: salmon_comb_res_fh.write(line.strip()) kkonganti@1: sal_yn[line.strip()] += 0 kkonganti@1: for serotype in serotypes: kkonganti@1: salmon_comb_res_fh.write("\t-") kkonganti@1: salmon_comb_res_fh.write( kkonganti@1: "\t-\n" kkonganti@1: ) if show_snp_clust_col else salmon_comb_res_fh.write("\n") kkonganti@1: no_calls_fh.close() kkonganti@1: kkonganti@1: for sample, counts in sorted(sample2salmon.items()): kkonganti@1: salmon_comb_res_fh.write(sample) kkonganti@1: snp_cluster_res_col = list() kkonganti@1: kkonganti@1: for snp_clust_id in sample2salmon[sample]["snp_clust_ids"].keys(): kkonganti@1: # print(snp_clust_id) kkonganti@1: # print(",".join(sample2salmon[sample]["snp_clust_ids"][snp_clust_id])) kkonganti@1: # ppp.pprint(sample2salmon[sample]["snp_clust_ids"]) kkonganti@1: # ppp.pprint(sample2salmon[sample]["snp_clust_ids"][snp_clust_id]) kkonganti@1: # final_url_text = ",".join( kkonganti@1: # sample2salmon[sample]["snp_clust_ids"][snp_clust_id] kkonganti@1: # ) kkonganti@1: # final_url_text_to_show = snp_clust_id kkonganti@1: # snp_cluster_res_col.append( kkonganti@1: # "".join( kkonganti@1: # [ kkonganti@1: # f'{snp_clust_id}', kkonganti@1: # ] kkonganti@1: # ) kkonganti@1: # ) kkonganti@1: final_url_text_to_show = " ".join( kkonganti@1: sample2salmon[sample]["snp_clust_ids"][snp_clust_id] kkonganti@1: ) kkonganti@1: snp_cluster_res_col.append( kkonganti@1: "".join( kkonganti@1: [ kkonganti@1: f'{final_url_text_to_show}', kkonganti@1: ] kkonganti@1: ) kkonganti@1: ) kkonganti@1: kkonganti@1: per_serotype_counts = 0 kkonganti@1: for serotype in serotypes: kkonganti@1: kkonganti@1: if serotype in sample2salmon[sample].keys(): kkonganti@1: # ppp.pprint(counts) kkonganti@1: sample_perc_mapped = round( kkonganti@1: sum(counts[serotype]) / sum(counts["tot_reads"]) * 100, 2 kkonganti@1: ) kkonganti@1: salmon_comb_res_fh.write( kkonganti@1: f"\t{sum(counts[serotype])} ({sample_perc_mapped}%)" kkonganti@1: ) kkonganti@1: multiqc_salmon_counts[sample].setdefault( kkonganti@1: re.match(r"^serotype=(.+?)\,antigen_formula.*", serotype)[1], kkonganti@1: sum(counts[serotype]), kkonganti@1: ) kkonganti@1: per_serotype_counts += sum(counts[serotype]) kkonganti@1: sal_yn[sample] += 1 kkonganti@1: else: kkonganti@1: salmon_comb_res_fh.write(f"\t-") kkonganti@1: sal_yn[sample] += 0 kkonganti@1: kkonganti@1: multiqc_salmon_counts[sample].setdefault( kkonganti@1: no_hit, sum(counts["tot_reads"]) - per_serotype_counts kkonganti@1: ) kkonganti@1: snp_clust_col_val = ( kkonganti@1: f'\t{" ".join(snp_cluster_res_col)}\n' if show_snp_clust_col else "\n" kkonganti@1: ) kkonganti@1: # ppp.pprint(multiqc_salmon_counts) kkonganti@1: salmon_comb_res_fh.write(snp_clust_col_val) kkonganti@1: kkonganti@1: with open(bcs_sal_yn, "w") as bcs_sal_yn_fh: kkonganti@1: bcs_sal_yn_fh.write("Sample\tSalmonella Presence\tNo. of Serotypes\n") kkonganti@1: for sample in sal_yn.keys(): kkonganti@1: if sal_yn[sample] > 0: kkonganti@1: bcs_sal_yn_fh.write(f"{sample}\tDetected\t{sal_yn[sample]}\n") kkonganti@1: else: kkonganti@1: bcs_sal_yn_fh.write(f"{sample}\tNot detected\t{sal_yn[sample]}\n") kkonganti@1: kkonganti@1: with open(cell_colors_yml_file, "w") as cell_colors_fh: kkonganti@1: yaml.dump(cell_colors_yml, cell_colors_fh, default_flow_style=False) kkonganti@1: kkonganti@1: salmon_plot_json(salmon_comb_res_mqc, multiqc_salmon_counts, no_hit) kkonganti@1: kkonganti@1: salmon_comb_res_fh.close() kkonganti@1: bcs_sal_yn_fh.close() kkonganti@1: cell_colors_fh.close() kkonganti@1: kkonganti@1: kkonganti@1: def salmon_plot_json(file: None, sample_salmon_counts: None, no_hit: None) -> None: kkonganti@1: """ kkonganti@1: This method will take a dictionary of salmon counts per sample kkonganti@1: and will dump a JSON that will be used by MultiQC. kkonganti@1: """ kkonganti@1: kkonganti@1: if file is None or sample_salmon_counts is None: kkonganti@1: logging.error( kkonganti@1: "Neither an output file to dump the JSON for MultiQC or the" kkonganti@1: + "dictionary holding the salmon counts was not passed." kkonganti@1: ) kkonganti@1: kkonganti@1: # Credit: http://phrogz.net/tmp/24colors.html kkonganti@1: # Will cycle through 20 distinct colors. kkonganti@1: distinct_color_palette = [ kkonganti@1: "#FF0000", kkonganti@1: "#FFFF00", kkonganti@1: "#00EAFF", kkonganti@1: "#AA00FF", kkonganti@1: "#FF7F00", kkonganti@1: "#BFFF00", kkonganti@1: "#0095FF", kkonganti@1: "#FF00AA", kkonganti@1: "#FFD400", kkonganti@1: "#6AFF00", kkonganti@1: "#0040FF", kkonganti@1: "#EDB9B9", kkonganti@1: "#B9D7ED", kkonganti@1: "#E7E9B9", kkonganti@1: "#DCB9ED", kkonganti@1: "#B9EDE0", kkonganti@1: "#8F2323", kkonganti@1: "#23628F", kkonganti@1: "#8F6A23", kkonganti@1: "#6B238F", kkonganti@1: "#4F8F23", kkonganti@1: ] kkonganti@1: kkonganti@1: no_hit_color = "#434348" kkonganti@1: col_count = 0 kkonganti@1: serotypes = set() kkonganti@1: salmon_counts = defaultdict(defaultdict) kkonganti@1: salmon_counts["id"] = "BETTERCALLSAL_SALMON_COUNTS" kkonganti@1: salmon_counts["section_name"] = "Salmon read counts" kkonganti@1: salmon_counts["description"] = ( kkonganti@1: "This section shows the read counts from running salmon " kkonganti@1: + "in --meta mode using SE, merged PE or concatenated PE reads against " kkonganti@1: + "an on-the-fly salmon index generated from the genome hits " kkonganti@1: + "of kma." kkonganti@1: ) kkonganti@1: salmon_counts["plot_type"] = "bargraph" kkonganti@1: salmon_counts["pconfig"]["id"] = "bettercallsal_salmon_counts_plot" kkonganti@1: salmon_counts["pconfig"]["title"] = "Salmon: Read counts" kkonganti@1: salmon_counts["pconfig"]["ylab"] = "Number of reads" kkonganti@1: salmon_counts["pconfig"]["xDecimals"] = "false" kkonganti@1: salmon_counts["pconfig"]["cpswitch_counts_label"] = "Number of reads (Counts)" kkonganti@1: salmon_counts["pconfig"]["cpswitch_percent_label"] = "Number of reads (Percentages)" kkonganti@1: kkonganti@1: for sample in sorted(sample_salmon_counts.keys()): kkonganti@1: serotypes.update(list(sample_salmon_counts[sample].keys())) kkonganti@1: salmon_counts["data"][sample] = sample_salmon_counts[sample] kkonganti@1: kkonganti@1: for serotype in sorted(serotypes): kkonganti@1: if serotype == no_hit: kkonganti@1: continue kkonganti@1: if col_count == len(distinct_color_palette) - 1: kkonganti@1: col_count = 0 kkonganti@1: kkonganti@1: col_count += 1 kkonganti@1: salmon_counts["categories"][serotype] = {"color": distinct_color_palette[col_count]} kkonganti@1: kkonganti@1: salmon_counts["categories"][no_hit] = {"color": no_hit_color} kkonganti@1: json.dump(salmon_counts, open(file, "w")) kkonganti@1: kkonganti@1: kkonganti@1: if __name__ == "__main__": kkonganti@1: main()