Mercurial > repos > kkonganti > cfsan_bettercallsal

comparison 0.7.0/bin/gen_otf_genome.py @ 17:0e7a0053e4a6

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload
author kkonganti
date Mon, 15 Jul 2024 10:42:02 -0400
parents
children
comparison
equal deleted inserted replaced
16:b90e5a7a3d4f 17:0e7a0053e4a6
1 #!/usr/bin/env python3
2
3 # Kranti Konganti
4
5 import argparse
6 import glob
7 import gzip
8 import inspect
9 import logging
10 import os
11 import pprint
12 import re
13
14 # Set logging.
15 logging.basicConfig(
16 format="\n"
17 + "=" * 55
18 + "\n%(asctime)s - %(levelname)s\n"
19 + "=" * 55
20 + "\n%(message)s\n\n",
21 level=logging.DEBUG,
22 )
23
24 # Debug print.
25 ppp = pprint.PrettyPrinter(width=50, indent=4)
26
27
28 # Multiple inheritence for pretty printing of help text.
29 class MultiArgFormatClasses(
30 argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter
31 ):
32 pass
33
34
35 def main() -> None:
36 """
37 This script works only in the context of `bettercallsal` Nextflow workflow.
38 It takes:
39 1. A text file containing accessions or FASTA IDs, one per line and
40 then,
41 2. Searches for a genome FASTA file in gzipped format in specified
42 search path, where the prefix of the filename is the accession or
43 FASTA ID from 1. and then,
44 creates a new concatenated gzipped genome FASTA file with all the genomes
45 in the text file from 1.
46 """
47
48 prog_name = os.path.basename(inspect.stack()[0].filename)
49
50 parser = argparse.ArgumentParser(
51 prog=prog_name, description=main.__doc__, formatter_class=MultiArgFormatClasses
52 )
53
54 required = parser.add_argument_group("required arguments")
55
56 required.add_argument(
57 "-txt",
58 dest="accs_txt",
59 default=False,
60 required=True,
61 help="Absolute UNIX path to .txt file containing accessions\n"
62 + "FASTA IDs, one per line.",
63 )
64 required.add_argument(
65 "-gd",
66 dest="genomes_dir",
67 default=False,
68 required=True,
69 help="Absolute UNIX path to a directory containing\n"
70 + "gzipped genome FASTA files.\n"
71 + "Required if -m is on.",
72 )
73 parser.add_argument(
74 "-gds",
75 dest="genomes_dir_suffix",
76 default="_scaffolded_genomic.fna.gz",
77 required=False,
78 help="Genome FASTA file suffix to search for\nin the directory mentioned using\n-gd.",
79 )
80 parser.add_argument(
81 "-op",
82 dest="out_prefix",
83 default="CATTED_GENOMES",
84 help="Set the output file prefix for .fna.gz and .txt\n" + "files.",
85 )
86 parser.add_argument(
87 "-txts",
88 dest="accs_suffix",
89 default="_template_hits.txt",
90 required=False,
91 help="The suffix of the file supplied with -txt option. It is assumed that the\n"
92 + "sample name is present in the file supplied with -txt option and the suffix\n"
93 + "will be stripped and stored in a file that logs samples which have no hits.",
94 )
95 parser.add_argument(
96 "-frag_delim",
97 dest="frag_delim",
98 default="\t",
99 required=False,
100 help="The delimitor by which the fields are separated in *_frag.gz file.",
101 )
102
103 args = parser.parse_args()
104 accs_txt = args.accs_txt
105 genomes_dir = args.genomes_dir
106 genomes_dir_suffix = args.genomes_dir_suffix
107 out_prefix = args.out_prefix
108 accs_suffix = args.accs_suffix
109 frag_delim = args.frag_delim
110 accs_seen = dict()
111 cat_genomes_gz = os.path.join(os.getcwd(), out_prefix + "_" + genomes_dir_suffix)
112 cat_genomes_gz = re.sub("__", "_", str(cat_genomes_gz))
113 frags_gz = os.path.join(os.getcwd(), out_prefix + ".frag.gz")
114 cat_reads_gz = os.path.join(os.getcwd(), out_prefix + "_aln_reads.fna.gz")
115 cat_reads_gz = re.sub("__", "_", cat_reads_gz)
116
117 if (
118 accs_txt
119 and os.path.exists(cat_genomes_gz)
120 and os.path.getsize(cat_genomes_gz) > 0
121 ):
122 logging.error(
123 "A concatenated genome FASTA file,\n"
124 + f"{os.path.basename(cat_genomes_gz)} already exists in:\n"
125 + f"{os.getcwd()}\n"
126 + "Please remove or move it as we will not "
127 + "overwrite it."
128 )
129 exit(1)
130
131 if accs_txt and (not os.path.exists(accs_txt) or not os.path.getsize(accs_txt) > 0):
132 logging.error("File,\n" + f"{accs_txt}\ndoes not exist " + "or is empty!")
133 failed_sample_name = re.sub(accs_suffix, "", os.path.basename(accs_txt))
134 with open(
135 os.path.join(os.getcwd(), "_".join([out_prefix, "FAILED.txt"])), "w"
136 ) as failed_sample_fh:
137 failed_sample_fh.write(f"{failed_sample_name}\n")
138 failed_sample_fh.close()
139 exit(0)
140
141 if genomes_dir:
142 if not os.path.isdir(genomes_dir):
143 logging.error("UNIX path\n" + f"{genomes_dir}\n" + "does not exist!")
144 exit(1)
145 if len(glob.glob(os.path.join(genomes_dir, "*" + genomes_dir_suffix))) <= 0:
146 logging.error(
147 "Genomes directory"
148 + f"{genomes_dir}"
149 + "\ndoes not seem to have any\n"
150 + f"files ending with suffix: {genomes_dir_suffix}"
151 )
152 exit(1)
153
154 # ppp.pprint(mash_hits)
155 empty_lines = 0
156 empty_lines_msg = ""
157 with open(cat_genomes_gz, "wb") as genomes_out_gz:
158 with open(accs_txt, "r") as accs_txt_fh:
159 for line in accs_txt_fh:
160 if line in ["\n", "\n\r"]:
161 empty_lines += 1
162 continue
163 else:
164 line = line.strip()
165
166 if line in accs_seen.keys():
167 continue
168 else:
169 accs_seen[line] = 1
170
171 genome_file = os.path.join(genomes_dir, line + genomes_dir_suffix)
172
173 if (
174 not os.path.exists(genome_file)
175 or os.path.getsize(genome_file) <= 0
176 ):
177 logging.error(
178 f"Genome file {os.path.basename(genome_file)} does not\n"
179 + "exits or is empty!"
180 )
181 exit(1)
182 else:
183 with open(genome_file, "rb") as genome_file_h:
184 genomes_out_gz.writelines(genome_file_h.readlines())
185 genome_file_h.close()
186 accs_txt_fh.close()
187 genomes_out_gz.close()
188
189 if (
190 len(accs_seen.keys()) > 0
191 and os.path.exists(frags_gz)
192 and os.path.getsize(frags_gz) > 0
193 ):
194 with gzip.open(
195 cat_reads_gz, "wt", encoding="utf-8", compresslevel=6
196 ) as cat_reads_gz_fh:
197 with gzip.open(frags_gz, "rb", compresslevel=6) as fragz_gz_fh:
198 for frag_line in fragz_gz_fh:
199 frag_lines = frag_line.decode("utf-8").strip().split(frag_delim)
200 # Per KMA specification, 6=template, 7=query, 1=read
201 cat_reads_gz_fh.write(f">{frag_lines[6]}\n{frag_lines[0]}\n")
202 fragz_gz_fh.close()
203 cat_reads_gz_fh.close()
204
205 if empty_lines > 0:
206 empty_lines_msg = f"Skipped {empty_lines} empty line(s).\n"
207
208 logging.info(
209 empty_lines_msg
210 + f"File {os.path.basename(cat_genomes_gz)}\n"
211 + f"written in:\n{os.getcwd()}\nDone! Bye!"
212 )
213 exit(0)
214
215
216 if __name__ == "__main__":
217 main()