Coverage for barbet/markers.py: 59.12%
159 statements
« prev ^ index » next coverage.py v7.9.1, created at 2026-08-17 06:22 +0000
1import os
2import subprocess
3from typing import Dict, List, Tuple, Optional
4import gzip
5import shutil
6import tempfile
7from multiprocessing import Manager
8from concurrent.futures import ProcessPoolExecutor
10from rich import progress
12# Marker information
13BAC120_MARKERS = [
14 "PF00380.20",
15 "PF00410.20",
16 "PF00466.21",
17 "PF01025.20",
18 "PF02576.18",
19 "PF03726.15",
20 "TIGR00006",
21 "TIGR00019",
22 "TIGR00020",
23 "TIGR00029",
24 "TIGR00043",
25 "TIGR00054",
26 "TIGR00059",
27 "TIGR00061",
28 "TIGR00064",
29 "TIGR00065",
30 "TIGR00082",
31 "TIGR00083",
32 "TIGR00084",
33 "TIGR00086",
34 "TIGR00088",
35 "TIGR00090",
36 "TIGR00092",
37 "TIGR00095",
38 "TIGR00115",
39 "TIGR00116",
40 "TIGR00138",
41 "TIGR00158",
42 "TIGR00166",
43 "TIGR00168",
44 "TIGR00186",
45 "TIGR00194",
46 "TIGR00250",
47 "TIGR00337",
48 "TIGR00344",
49 "TIGR00362",
50 "TIGR00382",
51 "TIGR00392",
52 "TIGR00396",
53 "TIGR00398",
54 "TIGR00414",
55 "TIGR00416",
56 "TIGR00420",
57 "TIGR00431",
58 "TIGR00435",
59 "TIGR00436",
60 "TIGR00442",
61 "TIGR00445",
62 "TIGR00456",
63 "TIGR00459",
64 "TIGR00460",
65 "TIGR00468",
66 "TIGR00472",
67 "TIGR00487",
68 "TIGR00496",
69 "TIGR00539",
70 "TIGR00580",
71 "TIGR00593",
72 "TIGR00615",
73 "TIGR00631",
74 "TIGR00634",
75 "TIGR00635",
76 "TIGR00643",
77 "TIGR00663",
78 "TIGR00717",
79 "TIGR00755",
80 "TIGR00810",
81 "TIGR00922",
82 "TIGR00928",
83 "TIGR00959",
84 "TIGR00963",
85 "TIGR00964",
86 "TIGR00967",
87 "TIGR01009",
88 "TIGR01011",
89 "TIGR01017",
90 "TIGR01021",
91 "TIGR01029",
92 "TIGR01032",
93 "TIGR01039",
94 "TIGR01044",
95 "TIGR01059",
96 "TIGR01063",
97 "TIGR01066",
98 "TIGR01071",
99 "TIGR01079",
100 "TIGR01082",
101 "TIGR01087",
102 "TIGR01128",
103 "TIGR01146",
104 "TIGR01164",
105 "TIGR01169",
106 "TIGR01171",
107 "TIGR01302",
108 "TIGR01391",
109 "TIGR01393",
110 "TIGR01394",
111 "TIGR01510",
112 "TIGR01632",
113 "TIGR01951",
114 "TIGR01953",
115 "TIGR02012",
116 "TIGR02013",
117 "TIGR02027",
118 "TIGR02075",
119 "TIGR02191",
120 "TIGR02273",
121 "TIGR02350",
122 "TIGR02386",
123 "TIGR02397",
124 "TIGR02432",
125 "TIGR02729",
126 "TIGR03263",
127 "TIGR03594",
128 "TIGR03625",
129 "TIGR03632",
130 "TIGR03654",
131 "TIGR03723",
132 "TIGR03725",
133 "TIGR03953",
134]
136AR53_MARKERS = [
137 "PF04919.13",
138 "PF07541.13",
139 "PF01000.27",
140 "PF00687.22",
141 "PF00466.21",
142 "PF00827.18",
143 "PF01280.21",
144 "PF01090.20",
145 "PF01200.19",
146 "PF01015.19",
147 "PF00900.21",
148 "PF00410.20",
149 "TIGR00037",
150 "TIGR00064",
151 "TIGR00111",
152 "TIGR00134",
153 "TIGR00279",
154 "TIGR00291",
155 "TIGR00323",
156 "TIGR00335",
157 "TIGR00373",
158 "TIGR00405",
159 "TIGR00448",
160 "TIGR00483",
161 "TIGR00491",
162 "TIGR00522",
163 "TIGR00967",
164 "TIGR00982",
165 "TIGR01008",
166 "TIGR01012",
167 "TIGR01018",
168 "TIGR01020",
169 "TIGR01028",
170 "TIGR01046",
171 "TIGR01052",
172 "TIGR01171",
173 "TIGR01213",
174 "TIGR01952",
175 "TIGR02236",
176 "TIGR02338",
177 "TIGR02389",
178 "TIGR02390",
179 "TIGR03626",
180 "TIGR03627",
181 "TIGR03628",
182 "TIGR03629",
183 "TIGR03670",
184 "TIGR03671",
185 "TIGR03672",
186 "TIGR03673",
187 "TIGR03674",
188 "TIGR03676",
189 "TIGR03680",
190]
193def read_fasta(path: str) -> Dict[str, str]:
194 """
195 Read a FASTA file into a dictionary of sequences.
197 Parameters
198 ----------
199 path : str
200 Filesystem path to the input FASTA file.
202 Returns
203 -------
204 Dict[str, str]
205 Mapping from sequence ID (the first token after '>' in the header)
206 to the full sequence string, with any terminal “*” characters stripped.
208 Raises
209 ------
210 IOError
211 If the file cannot be opened for reading.
212 """
213 seqs: Dict[str, str] = {}
214 with open(path) as fh:
215 header, buffer = None, []
216 for line in fh:
217 line = line.rstrip()
218 if not line:
219 continue
220 if line.startswith(">"):
221 if header:
222 seqs[header] = "".join(buffer).strip("*")
223 header = line[1:].split()[0]
224 buffer = []
225 else:
226 buffer.append(line)
227 if header:
228 seqs[header] = "".join(buffer).strip("*")
229 return seqs
232def run_prodigal(
233 genome_id: str,
234 fasta_path: str,
235 out_dir: str,
236 force: bool,
237 translation_table: Optional[int] = None,
238) -> str:
239 """
240 Run Prodigal to predict protein-coding genes from a FASTA file.
241 Parameters
242 ----------
243 genome_id : str
244 Unique identifier for the genome (used for output directory and file names).
245 fasta_path : str
246 Path to the input FASTA file containing genomic sequences.
247 out_dir : str
248 Directory where the output protein FASTA file will be saved.
249 force : bool
250 If True, overwrite existing output files.
251 translation_table : Optional[int]
252 Translation table to pass to Prodigal via -g flag (e.g., 4, 11, 25).
253 progress_dict : Dict
254 Shared dictionary to track progress across multiple processes.
255 task_id : int
256 Unique task ID for this genome, used to update progress in the shared dict.
257 Returns
258 -------
259 str
260 Path to the output protein FASTA file generated by Prodigal.
262 """
263 prot_dir = os.path.join(out_dir, genome_id)
264 os.makedirs(prot_dir, exist_ok=True)
265 prot_fa = os.path.join(prot_dir, f"{genome_id}.faa")
267 if force and os.path.exists(prot_fa):
268 os.remove(prot_fa)
270 # If input is gzipped, decompress to a temporary file
271 if fasta_path.endswith(".gz"):
272 with gzip.open(fasta_path, "rt") as gz_in, \
273 tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".fasta") as tmp_fa:
274 shutil.copyfileobj(gz_in, tmp_fa)
275 input_path = tmp_fa.name
276 else:
277 input_path = fasta_path
279 # Run Prodigal
280 cmd = ["prodigal", "-a", prot_fa, "-p", "meta", "-i", input_path]
281 if translation_table is not None:
282 cmd.extend(["-g", str(translation_table)])
284 subprocess.run(
285 cmd,
286 check=True,
287 stdout=subprocess.DEVNULL,
288 stderr=subprocess.DEVNULL,
289 )
290 return prot_fa
293def parse_domtblout_top_hits(domtbl_path: str) -> Dict[str, List[str]]:
294 """
295 Parse a HMMER --domtblout file and select the top hit per sequence.
297 Implements the GTDB-Tk comparator logic: for each query sequence,
298 keeps the hit with highest bitscore; ties broken by lower e-value,
299 then by lexicographically smaller HMM ID.
301 Parameters
302 ----------
303 domtbl_path : str
304 Path to the HMMER --domtblout output file.
306 Returns
307 -------
308 Dict[str, List[str]]
309 Mapping from HMM ID to a list of sequence IDs that were chosen
310 as top hit(s) for that HMM.
312 Raises
313 ------
314 IOError
315 If the domtblout file cannot be opened.
316 ValueError
317 If a non-numeric e-value or bitscore is encountered.
318 """
319 seq_matches: Dict[str, Tuple[str, float, float]] = {}
320 with open(domtbl_path) as fh:
321 for line in fh:
322 if line.startswith("#"):
323 continue
324 parts = line.split()
325 seq_id = parts[0]
326 hmm_id = parts[3]
327 evalue = float(parts[4])
328 bitscore = float(parts[5])
329 # look for the best hit for this sequence
330 prev = seq_matches.get(seq_id)
331 if prev is None:
332 seq_matches[seq_id] = (hmm_id, bitscore, evalue)
333 else:
334 # only keep the best hit
335 prev_hmm_id, prev_b, prev_e = prev
336 if (
337 bitscore > prev_b
338 or (bitscore == prev_b and evalue < prev_e)
339 or (
340 bitscore == prev_b and evalue == prev_e and hmm_id < prev_hmm_id
341 )
342 ):
343 seq_matches[seq_id] = (hmm_id, bitscore, evalue)
345 # now, invert the mapping to get HMM IDs to sequences
346 hits: Dict[str, List[str]] = {}
347 for seq_id, (hmm_id, _, _) in seq_matches.items():
348 hits.setdefault(hmm_id, []).append(seq_id)
349 return {hmm: sorted(seq_list) for hmm, seq_list in hits.items()}
352def _process_single_genome(
353 args: Tuple[
354 str, # genome_id
355 str, # fasta_path
356 str, # out_dir
357 int, # cpus_per_proc
358 str, # pfam_db
359 str, # tigr_db
360 bool, # force
361 bool, # skip_multiple_hits
362 int, # number_of_hits_to_keep
363 Dict, # shared progress dict (Manager().dict())
364 int, # task_id in Rich.Progress
365 ],
366) -> Tuple[str, Dict[str, List[str]]]:
367 """
368 Worker-function for one genome, updated to report progress after each subtask.
369 Subtasks:
370 1) Prodigal
371 2) Pfam HMM search
372 3) TIGRFAM HMM search
373 4) Parse + write FASTAs
375 Returns: ( "path/to/genome.fasta", { "bac120": [...], "ar53": [...] } )
376 """
377 (
378 gid,
379 path,
380 out_dir,
381 cpus_per_proc,
382 pfam_db,
383 tigr_db,
384 force,
385 skip_multiple_hits,
386 number_of_hits_to_keep,
387 progress_dict,
388 task_id,
389 *rest,
390 ) = args
392 translation_table = rest[0] if rest else None
394 genome_fastas: Dict[str, List[str]] = {"bac120": [], "ar53": []}
396 # Signal that we are starting the job
397 progress_dict[task_id] = {"progress": 0}
399 # Prodigal
400 prot_fa = run_prodigal(gid, path, out_dir, force, translation_table=translation_table)
401 prot_seqs = read_fasta(prot_fa)
402 # Signal to the shared dict that we've completed step 1 (Prodigal)
403 progress_dict[task_id] = {"progress": 1}
405 # Pfam HMM search
406 pf_out = os.path.join(out_dir, gid, "pfam.tblout")
407 if force and os.path.exists(pf_out):
408 os.remove(pf_out)
410 subprocess.run(
411 [
412 "hmmsearch",
413 "--cpu",
414 str(cpus_per_proc),
415 "--notextw",
416 "-E",
417 "0.001",
418 "--domE",
419 "0.001",
420 "--tblout",
421 pf_out,
422 pfam_db,
423 prot_fa,
424 ],
425 check=True,
426 stdout=subprocess.DEVNULL,
427 stderr=subprocess.DEVNULL,
428 )
429 progress_dict[task_id] = {"progress": 2}
431 # TIGRFAM HMM search
432 tg_out = os.path.join(out_dir, gid, "tigrfam.tblout")
433 if force and os.path.exists(tg_out):
434 os.remove(tg_out)
436 subprocess.run(
437 [
438 "hmmsearch",
439 "--cpu",
440 str(cpus_per_proc),
441 "--noali",
442 "--notextw",
443 "--cut_nc",
444 "--tblout",
445 tg_out,
446 tigr_db,
447 prot_fa,
448 ],
449 check=True,
450 stdout=subprocess.DEVNULL,
451 stderr=subprocess.DEVNULL,
452 )
453 progress_dict[task_id] = {"progress": 3}
455 # Write FASTAs for each marker
456 pf_hits = parse_domtblout_top_hits(pf_out)
457 tg_hits = parse_domtblout_top_hits(tg_out)
458 combined_hits = {**pf_hits, **tg_hits}
460 for marker, seq_ids in sorted(combined_hits.items()):
461 if not seq_ids:
462 continue
463 elif len(seq_ids) == 1:
464 seqs = [prot_seqs[seq_ids[0]]]
465 else:
466 unique_seqs = set(prot_seqs[s] for s in seq_ids)
467 if len(unique_seqs) != 1 and skip_multiple_hits:
468 # faster but can miss some markers
469 continue
470 seqs = sorted(unique_seqs)[:number_of_hits_to_keep]
472 for dom in ("bac120", "ar53"):
473 if (dom == "bac120" and marker in BAC120_MARKERS) or (
474 dom == "ar53" and marker in AR53_MARKERS
475 ):
476 genome_dir = os.path.join(out_dir, gid, dom)
477 os.makedirs(genome_dir, exist_ok=True)
478 for i, seq in enumerate(seqs, start=1):
479 fa_path = os.path.join(genome_dir, f"{marker}.{i}.fa")
480 with open(fa_path, "w") as fh:
481 fh.write(f">{gid}\n{seq}\n")
482 genome_fastas[dom].append(fa_path)
484 # Signal completion of subtask 4 (and thus the entire genome job)
485 progress_dict[task_id] = {"progress": 4}
486 return (path, genome_fastas)
489def extract_markers_genes(
490 genomes: Dict[str, str],
491 out_dir: str,
492 cpus: int = 1,
493 pfam_db: str = os.environ.get("PFAM_HMMDB"),
494 tigr_db: str = os.environ.get("TIGR_HMMDB"),
495 force: bool = False,
496 skip_multiple_hits: bool = False,
497 number_of_hits_to_keep: int = 1,
498 translation_tables: Optional[Dict[str, int]] = None,
499) -> Dict[str, Dict[str, List[str]]]:
500 """
501 Extract marker genes from multiple genomes in parallel using Prodigal and HMMER.
502 """
504 # Early exit if no genomes
505 if not genomes:
506 return {}
508 n_genomes = len(genomes)
509 n_workers = min(n_genomes, cpus)
510 cpus_per_proc = max(1, cpus // n_workers)
512 # This shared dict allows us to track progress
513 # across multiple processes
514 manager = Manager()
515 progress_dict = manager.dict()
517 genome_to_task: Dict[str, int] = {}
518 futures_to_task: Dict = {}
520 results: Dict[str, Dict[str, List[str]]] = {}
522 with progress.Progress(
523 progress.TextColumn("[progress.description]{task.description}", justify="right"),
524 progress.BarColumn(),
525 progress.TaskProgressColumn(),
526 progress.TimeElapsedColumn(),
527 ) as rich_progress:
529 # Over progress bar that counts how many genomes finished all 4 steps
530 overall_task = rich_progress.add_task("[cyan]Extracting Markers", total=n_genomes)
532 # Create one sub‐task per genome (each with total=4)
533 for idx, (gid, fasta_path) in enumerate(genomes.items()):
534 # 4 total steps per genome: Prodigal, Pfam, TIGRFAM, and writing FASTAs
535 task_id = rich_progress.add_task(gid, total=4, visible=False, start=False)
536 genome_to_task[gid] = task_id
537 progress_dict[task_id] = {}
539 # Submit all genome jobs to a ProcessPoolExecutor
540 with ProcessPoolExecutor(max_workers=n_workers) as executor:
541 for gid, fasta_path in genomes.items():
542 task_id = genome_to_task[gid]
543 trans_table = translation_tables.get(gid) if translation_tables else None
544 args = (
545 gid,
546 fasta_path,
547 out_dir,
548 cpus_per_proc,
549 pfam_db,
550 tigr_db,
551 force,
552 skip_multiple_hits,
553 number_of_hits_to_keep,
554 progress_dict,
555 task_id,
556 trans_table,
557 )
558 future = executor.submit(_process_single_genome, args)
559 futures_to_task[future] = (task_id, fasta_path)
561 started_tasks = set() # Track which tasks have started
562 # Monitor progress until all futures complete
563 while futures_to_task:
564 for task_id, status in progress_dict.items():
565 started = "progress" in status # only started tasks have a "progress" key
566 if started and task_id not in started_tasks:
567 # If the task has started or progressed, ensure the progress bar is visible
568 rich_progress.update(task_id, visible=True)
569 rich_progress.start_task(task_id)
570 started_tasks.add(task_id)
571 elif started:
572 # Update the individual genome progress bar
573 rich_progress.update(task_id, completed=status["progress"])
575 # Check for any finished futures
576 done_now = []
577 for fut, (task_id, fasta_path) in futures_to_task.items():
578 if fut.done():
579 done_now.append(fut)
580 rich_progress.update(task_id, visible=False) # Hide the task bar
582 for fut in done_now:
583 task_id, fasta_path = futures_to_task.pop(fut)
584 try:
585 path_key, genome_dict = fut.result()
586 except Exception as e:
587 # If a worker failed, re‐raise
588 # TODO: consider logging the error instead
589 raise RuntimeError(f"Error processing {fasta_path}: {e}") from e
591 # Store the results
592 results[path_key] = genome_dict
594 status = progress_dict.get(task_id, {})
595 if status.get("progress", 0) >= 4:
596 # advance overall by one
597 rich_progress.update(overall_task, advance=1)
599 # Small sleep to avoid busy‐waiting too tightly
600 import time
601 time.sleep(0.2)
602 # Ensure we update the overall task to the final count
603 rich_progress.update(overall_task, completed=n_genomes)
605 return {gid: results[fasta_path] for gid, fasta_path in genomes.items() if fasta_path in results}