kkonganti@1: #!/usr/bin/env python3 kkonganti@1: kkonganti@1: # Kranti Konganti kkonganti@1: kkonganti@1: import os kkonganti@1: import glob kkonganti@1: import pickle kkonganti@1: import argparse kkonganti@1: import inspect kkonganti@1: import logging kkonganti@1: import re kkonganti@1: import pprint kkonganti@1: from collections import defaultdict 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: This script works only in the context of `bettercallsal` Nextflow workflow. kkonganti@1: It takes: kkonganti@1: 1. A pickle file containing a dictionary object where genome accession kkonganti@1: is the key and the computed serotype is the value. kkonganti@1: 2. A file with `mash screen` results run against the Salmonella SNP kkonganti@1: Cluster genomes' sketch. kkonganti@1: 3. A directory containing genomes' FASTA in gzipped format where the kkonganti@1: FASTA file contains 2 lines: one FASTA header followed by kkonganti@1: genome Sequence. kkonganti@1: and then generates a concatenated FASTA file of top N unique `mash screen` kkonganti@1: genome hits as requested. kkonganti@1: """ 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 = os.path.basename(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: parser.add_argument( kkonganti@1: "-s", kkonganti@1: dest="sero_snp_metadata", kkonganti@1: default=False, kkonganti@1: required=False, kkonganti@1: help="Absolute UNIX path to metadata text file with the field separator, | " kkonganti@1: + "\nand 5 fields: serotype|asm_lvl|asm_url|snp_cluster_id" kkonganti@1: + "\nEx: serotype=Derby,antigen_formula=4:f,g:-|Scaffold|402440|ftp://...\n|PDS000096654.2\n" kkonganti@1: + "Mentioning this option will create a pickle file for the\nprovided metadata and exits.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-fs", kkonganti@1: dest="force_write_pick", kkonganti@1: action="store_true", kkonganti@1: required=False, kkonganti@1: help="By default, when -s flag is on, the pickle file named *.ACC2SERO.pickle\n" kkonganti@1: + "is written to CWD. If the file exists, the program will not overwrite\n" kkonganti@1: + "and exit. Use -fs option to overwrite.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-m", kkonganti@1: dest="mash_screen_res", kkonganti@1: default=False, kkonganti@1: required=False, kkonganti@1: help="Absolute UNIX path to `mash screen` results file.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-ms", kkonganti@1: dest="mash_screen_res_suffix", kkonganti@1: default=".screened", kkonganti@1: required=False, kkonganti@1: help="Suffix of the `mash screen` result file.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-ps", kkonganti@1: dest="pickled_sero", kkonganti@1: default=False, kkonganti@1: required=False, kkonganti@1: help="Absolute UNIX Path to serialized metadata object in a pickle file.\n" kkonganti@1: + "You can create the pickle file of the metadata using -s option.\n" kkonganti@1: + "Required if -m is on.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-gd", kkonganti@1: dest="genomes_dir", kkonganti@1: default=False, kkonganti@1: required=False, kkonganti@1: help="Absolute UNIX path to a directory containing\n" kkonganti@1: + "gzipped genome FASTA files.\n" kkonganti@1: + "Required if -m is on.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-gds", kkonganti@1: dest="genomes_dir_suffix", kkonganti@1: default="_scaffolded_genomic.fna.gz", kkonganti@1: required=False, kkonganti@1: help="Genome FASTA file suffix to search for\nin the directory mentioned using\n-gd.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-n", kkonganti@1: dest="num_uniq_hits", kkonganti@1: default=10, kkonganti@1: help="This many number of serotype genomes' accessions are " + "\nreturned.", kkonganti@1: ) kkonganti@1: parser.add_argument( kkonganti@1: "-op", kkonganti@1: dest="out_prefix", kkonganti@1: default="MASH_SCREEN", kkonganti@1: help="Set the output file prefix for .fna.gz and .txt files.", kkonganti@1: ) kkonganti@1: # required = parser.add_argument_group('required arguments') kkonganti@1: kkonganti@1: args = parser.parse_args() kkonganti@1: num_uniq_hits = int(args.num_uniq_hits) kkonganti@1: mash_screen_res = args.mash_screen_res kkonganti@1: mash_screen_res_suffix = args.mash_screen_res_suffix kkonganti@1: pickle_sero = args.sero_snp_metadata kkonganti@1: pickled_sero = args.pickled_sero kkonganti@1: f_write_pick = args.force_write_pick kkonganti@1: genomes_dir = args.genomes_dir kkonganti@1: genomes_dir_suffix = args.genomes_dir_suffix kkonganti@1: out_prefix = args.out_prefix kkonganti@1: mash_genomes_gz = os.path.join( kkonganti@1: os.getcwd(), out_prefix + "_TOP_" + str(num_uniq_hits) + "_UNIQUE_HITS.fna.gz" kkonganti@1: ) kkonganti@1: mash_uniq_hits_txt = os.path.join( kkonganti@1: os.getcwd(), re.sub(".fna.gz", ".txt", os.path.basename(mash_genomes_gz)) kkonganti@1: ) kkonganti@1: kkonganti@1: if mash_screen_res and os.path.exists(mash_genomes_gz): kkonganti@1: logging.error( kkonganti@1: "A concatenated genome FASTA file,\n" kkonganti@1: + f"{os.path.basename(mash_genomes_gz)} already exists in:\n" kkonganti@1: + f"{os.getcwd()}\n" kkonganti@1: + "Please remove or move it as we will not " kkonganti@1: + "overwrite it." kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if os.path.exists(mash_uniq_hits_txt) and os.path.getsize(mash_uniq_hits_txt) > 0: kkonganti@1: os.remove(mash_uniq_hits_txt) kkonganti@1: kkonganti@1: if mash_screen_res and (not genomes_dir or not pickled_sero): kkonganti@1: logging.error("When -m is on, -ps and -gd are also required.") kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if genomes_dir: kkonganti@1: if not os.path.isdir(genomes_dir): kkonganti@1: logging.error("UNIX path\n" + f"{genomes_dir}\n" + "does not exist!") kkonganti@1: exit(1) kkonganti@1: if len(glob.glob(os.path.join(genomes_dir, "*" + genomes_dir_suffix))) <= 0: kkonganti@1: logging.error( kkonganti@1: "Genomes directory" kkonganti@1: + f"{genomes_dir}" kkonganti@1: + "\ndoes not seem to have any\n" kkonganti@1: + f"files ending with suffix: {genomes_dir_suffix}" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if pickle_sero and os.path.exists(pickle_sero) and os.path.getsize(pickle_sero) > 0: kkonganti@1: acc2serotype = defaultdict() kkonganti@1: init_pickled_sero = os.path.join(os.getcwd(), out_prefix + ".ACC2SERO.pickle") kkonganti@1: kkonganti@1: if ( kkonganti@1: os.path.exists(init_pickled_sero) kkonganti@1: and os.path.getsize(init_pickled_sero) kkonganti@1: and not f_write_pick kkonganti@1: ): kkonganti@1: logging.error( kkonganti@1: f"File {os.path.basename(init_pickled_sero)} already exists in\n{os.getcwd()}\n" kkonganti@1: + "Use -fs to force overwrite it." kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: with open(pickle_sero, "r") as sero_snp_meta: kkonganti@1: for line in sero_snp_meta: kkonganti@1: cols = line.strip().split("|") kkonganti@1: url_cols = cols[3].split("/") kkonganti@1: kkonganti@1: if not 4 <= len(cols) <= 5: kkonganti@1: logging.error( kkonganti@1: f"The metadata file {pickle_sero} is malformed.\n" kkonganti@1: + f"Expected 4-5 columns. Got {len(cols)} columns.\n" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if not len(url_cols) > 5: kkonganti@1: acc = url_cols[3] kkonganti@1: else: kkonganti@1: acc = url_cols[9] kkonganti@1: kkonganti@1: if not re.match(r"^GC[AF]\_\d+\.\d+$", acc): kkonganti@1: logging.error( kkonganti@1: f"Did not find accession in either field number 4\n" kkonganti@1: + "or field number 10 of column 4." kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: acc2serotype[acc] = cols[0] kkonganti@1: kkonganti@1: with open(init_pickled_sero, "wb") as write_pickled_sero: kkonganti@1: pickle.dump(file=write_pickled_sero, obj=acc2serotype) kkonganti@1: kkonganti@1: logging.info( kkonganti@1: f"Created the pickle file for\n{os.path.basename(pickle_sero)}.\n" kkonganti@1: + "This was the only requested function." kkonganti@1: ) kkonganti@1: sero_snp_meta.close() kkonganti@1: write_pickled_sero.close() kkonganti@1: exit(0) kkonganti@1: elif pickle_sero and not (os.path.exists(pickle_sero) and os.path.getsize(pickle_sero) > 0): kkonganti@1: kkonganti@1: logging.error( kkonganti@1: "Requested to create pickle from metadata, but\n" kkonganti@1: + f"the file, {os.path.basename(pickle_sero)} is empty or\ndoes not exist!" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: if mash_screen_res and os.path.exists(mash_screen_res): kkonganti@1: if os.path.getsize(mash_screen_res) > 0: kkonganti@1: kkonganti@1: seen_uniq_hits = 0 kkonganti@1: unpickled_acc2serotype = pickle.load(file=open(pickled_sero, "rb")) kkonganti@1: kkonganti@1: with open(mash_screen_res, "r") as msh_res: kkonganti@1: mash_hits = defaultdict() kkonganti@1: seen_mash_sero = defaultdict() kkonganti@1: kkonganti@1: for line in msh_res: kkonganti@1: cols = line.strip().split("\t") kkonganti@1: kkonganti@1: if len(cols) < 5: kkonganti@1: logging.error( kkonganti@1: f"The file {os.path.basename(mash_screen_res)} seems to\n" kkonganti@1: + "be malformed. It contains less than required 5-6 columns." kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: kkonganti@1: mash_hit_acc = re.sub( kkonganti@1: genomes_dir_suffix, kkonganti@1: "", kkonganti@1: str((re.search(r"GC[AF].*?" + genomes_dir_suffix, cols[4])).group()), kkonganti@1: ) kkonganti@1: kkonganti@1: if mash_hit_acc: kkonganti@1: mash_hits.setdefault(cols[0], []).append(mash_hit_acc) kkonganti@1: else: kkonganti@1: logging.error( kkonganti@1: "Did not find an assembly accession in column\n" kkonganti@1: + f"number 5. Found {cols[4]} instead. Cannot proceed!" kkonganti@1: ) kkonganti@1: exit(1) kkonganti@1: msh_res.close() kkonganti@1: elif os.path.getsize(mash_screen_res) == 0: kkonganti@1: failed_sample_name = os.path.basename(mash_screen_res).rstrip(mash_screen_res_suffix) kkonganti@1: with open( kkonganti@1: os.path.join(os.getcwd(), "_".join([out_prefix, "FAILED.txt"])), "w" kkonganti@1: ) as failed_sample_fh: kkonganti@1: failed_sample_fh.write(f"{failed_sample_name}\n") kkonganti@1: failed_sample_fh.close() kkonganti@1: exit(0) kkonganti@1: kkonganti@1: # ppp.pprint(mash_hits) kkonganti@1: msh_out_txt = open(mash_uniq_hits_txt, "w") kkonganti@1: with open(mash_genomes_gz, "wb") as msh_out_gz: kkonganti@1: for _, (ident, acc_list) in enumerate(sorted(mash_hits.items(), reverse=True)): kkonganti@1: for acc in acc_list: kkonganti@1: if seen_uniq_hits >= num_uniq_hits: kkonganti@1: break kkonganti@1: if unpickled_acc2serotype[acc] not in seen_mash_sero.keys(): kkonganti@1: seen_mash_sero[unpickled_acc2serotype[acc]] = 1 kkonganti@1: seen_uniq_hits += 1 kkonganti@1: # print(acc.strip() + '\t' + ident + '\t' + unpickled_acc2serotype[acc], file=sys.stdout) kkonganti@1: msh_out_txt.write( kkonganti@1: f"{acc.strip()}\t{unpickled_acc2serotype[acc]}\t{ident}\n" kkonganti@1: ) kkonganti@1: with open( kkonganti@1: os.path.join(genomes_dir, acc + genomes_dir_suffix), "rb" kkonganti@1: ) as msh_in_gz: kkonganti@1: msh_out_gz.writelines(msh_in_gz.readlines()) kkonganti@1: msh_in_gz.close() kkonganti@1: msh_out_gz.close() kkonganti@1: msh_out_txt.close() kkonganti@1: logging.info( kkonganti@1: f"File {os.path.basename(mash_genomes_gz)}\n" kkonganti@1: + f"written in:\n{os.getcwd()}\nDone! Bye!" kkonganti@1: ) kkonganti@1: exit(0) kkonganti@1: kkonganti@1: kkonganti@1: if __name__ == "__main__": kkonganti@1: main()