Mercurial > repos > kkonganti > cfsan_bettercallsal

annotate 0.6.1/bin/sourmash_sim_matrix.py @ 14:b0a37e88ecb5

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
"planemo upload"
author kkonganti
date Thu, 07 Sep 2023 10:13:31 -0400
parents 749faef1caa9
children
rev   line source
kkonganti@11 1 #!/usr/bin/env python3
kkonganti@11 2
kkonganti@11 3 # Kranti Konganti
kkonganti@11 4
kkonganti@11 5 import os
kkonganti@11 6 import argparse
kkonganti@11 7 import inspect
kkonganti@11 8 import logging
kkonganti@11 9 import re
kkonganti@11 10 import pickle
kkonganti@11 11 import pprint
kkonganti@11 12 import json
kkonganti@11 13 from collections import defaultdict
kkonganti@11 14
kkonganti@11 15 # Set logging.
kkonganti@11 16 logging.basicConfig(
kkonganti@11 17 format="\n" + "=" * 55 + "\n%(asctime)s - %(levelname)s\n" + "=" * 55 + "\n%(message)s\n\n",
kkonganti@11 18 level=logging.DEBUG,
kkonganti@11 19 )
kkonganti@11 20
kkonganti@11 21 # Debug print.
kkonganti@11 22 ppp = pprint.PrettyPrinter(width=50, indent=4)
kkonganti@11 23
kkonganti@11 24 # Multiple inheritence for pretty printing of help text.
kkonganti@11 25 class MultiArgFormatClasses(argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter):
kkonganti@11 26 pass
kkonganti@11 27
kkonganti@11 28
kkonganti@11 29 def main() -> None:
kkonganti@11 30 """
kkonganti@11 31 This script works only in the context of `bettercallsal` Nextflow workflow.
kkonganti@11 32 It takes:
kkonganti@11 33 1. A CSV file containing a similarity matrix or dissimilarity matrix where
kkonganti@11 34 the header row contains the names.
kkonganti@11 35 3. It takes indexed NCBI Pathogen metadata in pickle format and converts
kkonganti@11 36 accessions to serotype names in the final distance matrix output.
kkonganti@11 37 """
kkonganti@11 38
kkonganti@11 39 prog_name = os.path.basename(inspect.stack()[0].filename)
kkonganti@11 40
kkonganti@11 41 parser = argparse.ArgumentParser(
kkonganti@11 42 prog=prog_name, description=main.__doc__, formatter_class=MultiArgFormatClasses
kkonganti@11 43 )
kkonganti@11 44
kkonganti@11 45 required = parser.add_argument_group("required arguments")
kkonganti@11 46
kkonganti@11 47 required.add_argument(
kkonganti@11 48 "-csv",
kkonganti@11 49 dest="mat",
kkonganti@11 50 default=False,
kkonganti@11 51 required=True,
kkonganti@11 52 help="Absolute UNIX path to .csv file containing similarity\n"
kkonganti@11 53 + "or dissimilarity matrix from `sourmash compare`.",
kkonganti@11 54 )
kkonganti@11 55 required.add_argument(
kkonganti@11 56 "-pickle",
kkonganti@11 57 dest="acc2sero",
kkonganti@11 58 default=False,
kkonganti@11 59 required=True,
kkonganti@11 60 help="Absolute UNIX Path to the *ACC2SERO.pickle\n"
kkonganti@11 61 + "metadata file. On raven2, these are located at\n"
kkonganti@11 62 + "/hpc/db/bettercallsal/PDGXXXXXXXXXX.XXXXX/",
kkonganti@11 63 )
kkonganti@11 64 required.add_argument(
kkonganti@11 65 "-labels",
kkonganti@11 66 dest="labels",
kkonganti@11 67 default=False,
kkonganti@11 68 required=True,
kkonganti@11 69 help="Absolute UNIX Path to the *.labels.txt\n"
kkonganti@11 70 + "file from `sourmash compare`. The accessions\n"
kkonganti@11 71 + "will be renanamed to serotype names.",
kkonganti@11 72 )
kkonganti@11 73
kkonganti@11 74 args = parser.parse_args()
kkonganti@11 75 csv = args.mat
kkonganti@11 76 labels = args.labels
kkonganti@11 77 pickled_sero = args.acc2sero
kkonganti@11 78 row_names = list()
kkonganti@11 79 distance_mat = defaultdict(defaultdict)
kkonganti@11 80 out_csv = os.path.join(os.getcwd(), "bcs_sourmash_matrix.tblsum.txt")
kkonganti@11 81 out_json = os.path.join(os.getcwd(), "bcs_sourmash_matrix_mqc.json")
kkonganti@11 82
kkonganti@11 83 # Prepare dictionary to be dumped as JSON.
kkonganti@11 84 distance_mat["id"] = "BETTERCALLSAL_CONTAINMENT_INDEX"
kkonganti@11 85 distance_mat["section_name"] = "Containment index"
kkonganti@11 86 distance_mat["description"] = (
kkonganti@11 87 "This section shows the containment index between a sample and the genomes"
kkonganti@11 88 + "by running <code>sourmash gather</code> "
kkonganti@11 89 + "using <code>--containment</code> option."
kkonganti@11 90 )
kkonganti@11 91 distance_mat["plot_type"] = "heatmap"
kkonganti@11 92 distance_mat["pconfig"]["id"] = "bettercallsal_containment_index_heatmap"
kkonganti@11 93 distance_mat["pconfig"]["title"] = "Sourmash: containment index"
kkonganti@11 94 distance_mat["pconfig"]["xTitle"] = "Samples"
kkonganti@11 95 distance_mat["pconfig"]["yTitle"] = "Isolates (Genome assemblies)"
kkonganti@11 96 distance_mat["pconfig"]["ycats_samples"] = "False"
kkonganti@11 97 distance_mat["pconfig"]["xcats_samples"] = "False"
kkonganti@11 98 distance_mat["pconfig"]["square"] = "False"
kkonganti@11 99 distance_mat["pconfig"]["min"] = "0.0"
kkonganti@11 100 distance_mat["pconfig"]["max"] = "1.0"
kkonganti@11 101 distance_mat["data"]["data"] = list()
kkonganti@11 102
kkonganti@11 103 if pickled_sero and (not os.path.exists(pickled_sero) or not os.path.getsize(pickled_sero)):
kkonganti@11 104 logging.error(
kkonganti@11 105 "The pickle file,\n" + f"{os.path.basename(pickled_sero)} does not exist or is empty!"
kkonganti@11 106 )
kkonganti@11 107 exit(1)
kkonganti@11 108 else:
kkonganti@11 109 acc2sero = pickle.load(file=open(pickled_sero, "rb"))
kkonganti@11 110
kkonganti@11 111 if csv and (not os.path.exists(csv) or not os.path.getsize(csv) > 0):
kkonganti@11 112 logging.error("File,\n" + f"{csv}\ndoes not exist " + "or is empty!")
kkonganti@11 113 exit(0)
kkonganti@11 114
kkonganti@11 115 if labels and (not os.path.exists(labels) or not os.path.getsize(labels) > 0):
kkonganti@11 116 logging.error("File,\n" + f"{labels}\ndoes not exist " + "or is empty!")
kkonganti@11 117 exit(0)
kkonganti@11 118
kkonganti@11 119 # with open(out_labels, "w") as out_labels_fh:
kkonganti@11 120 with open(labels, "r") as labels_fh:
kkonganti@11 121 for line in labels_fh:
kkonganti@11 122 line = line.strip()
kkonganti@11 123 if line not in acc2sero.keys():
kkonganti@11 124 row_names.append(line)
kkonganti@11 125
kkonganti@11 126 labels_fh.close()
kkonganti@11 127
kkonganti@11 128 with open(out_csv, "w") as csv_out_fh:
kkonganti@11 129 with open(csv, "r") as csv_in_fh:
kkonganti@11 130 header = csv_in_fh.readline().strip().split(",")
kkonganti@11 131 acc_cols = [idx for idx, col in enumerate(header) if col in acc2sero.keys()]
kkonganti@11 132 sample_cols = [idx for idx, col in enumerate(header) if col not in acc2sero.keys()]
kkonganti@11 133
kkonganti@11 134 col_names = [
kkonganti@11 135 re.sub(r"serotype=|\,antigen_formula=.*?\|", "", s)
kkonganti@11 136 for s in [acc2sero[col] + f"| | {col}" for col in header if col in acc2sero.keys()]
kkonganti@11 137 ]
kkonganti@11 138
kkonganti@11 139 distance_mat["xcats"] = col_names
kkonganti@11 140 csv_out_fh.write("\t".join(["Sample"] + col_names) + "\n")
kkonganti@11 141 line_num = 0
kkonganti@11 142
kkonganti@11 143 for line in csv_in_fh:
kkonganti@11 144 if line_num not in sample_cols:
kkonganti@11 145 continue
kkonganti@11 146 else:
kkonganti@11 147
kkonganti@11 148 heatmap_rows = [
kkonganti@11 149 str(round(float(line.strip().split(",")[col]), 5)) for col in acc_cols
kkonganti@11 150 ]
kkonganti@11 151 # distance_mat["data"]["hmdata"].append(heatmap_rows)
kkonganti@11 152 # distance_mat["data"][row_names[line_num]] = heatmap_rows
kkonganti@11 153 distance_mat["data"]["data"].append(heatmap_rows)
kkonganti@11 154 # distance_mat["data"][row_names[line_num]] = dict(
kkonganti@11 155 # [(col_names[idx], val) for idx, val in enumerate(heatmap_rows)]
kkonganti@11 156 # )
kkonganti@11 157 csv_out_fh.write("\t".join([row_names[line_num]] + heatmap_rows) + "\n")
kkonganti@11 158 line_num += 1
kkonganti@11 159 csv_in_fh.close()
kkonganti@11 160 csv_out_fh.close()
kkonganti@11 161
kkonganti@11 162 distance_mat["ycats"] = row_names
kkonganti@11 163 json.dump(distance_mat, open(out_json, "w"))
kkonganti@11 164
kkonganti@11 165
kkonganti@11 166 if __name__ == "__main__":
kkonganti@11 167 main()