Coverage for barbet/apps.py: 56.09%
353 statements
« prev ^ index » next coverage.py v7.9.1, created at 2026-08-17 06:22 +0000
1import sys
2import logging
3from typing import TYPE_CHECKING
4from pathlib import Path
5from enum import Enum
6from collections import defaultdict
7import time
8from rich.progress import (
9 Progress,
10 TextColumn,
11 BarColumn,
12 TaskProgressColumn,
13 TimeElapsedColumn,
14 TimeRemainingColumn,
15)
16from torchapp import TorchApp, Param, method, main, tool
18from .output import print_polars_df
19from .logging import setup_logger
21if TYPE_CHECKING:
22 from collections.abc import Iterable
23 from torchmetrics import Metric
24 from hierarchicalsoftmax import SoftmaxNode
25 from torch import nn
26 import lightning as L
27 # import pandas as pd
28 import polars as pl
32class ImageFormat(str, Enum):
33 """The image format to use for the output images."""
35 NONE = ""
36 PNG = "png"
37 JPG = "jpg"
38 SVG = "svg"
39 PDF = "pdf"
40 DOT = "dot"
42 def __str__(self):
43 return self.value
45 def __bool__(self) -> bool:
46 """Returns True if the image format is not empty."""
47 return self.value != ""
50class Barbet(TorchApp):
51 def __init__(self, **kwargs):
52 super().__init__(**kwargs)
53 if hasattr(self, "main_app") and hasattr(self.main_app, "info"):
54 self.main_app.info.context_settings = {"help_option_names": ["-h", "--help"]}
55 if hasattr(self, "tools_app") and hasattr(self.tools_app, "info"):
56 self.tools_app.info.context_settings = {"help_option_names": ["-h", "--help"]}
58 @property
59 def logger(self) -> logging.Logger:
60 if not hasattr(self, "_logger") or self._logger is None:
61 output_dir = getattr(self, "output_dir", Path("output"))
62 self._logger = setup_logger(output_dir)
63 return self._logger
65 @logger.setter
66 def logger(self, logger: logging.Logger) -> None:
67 self._logger = logger
69 @method
70 def setup(
71 self,
72 memmap: str = None,
73 memmap_index: str = None,
74 treedict: str = None,
75 stack_size: int = 32,
76 in_memory: bool = False,
77 tip_alpha: float = None,
78 ) -> None:
79 if not treedict:
80 raise ValueError("treedict is required")
81 if not memmap:
82 raise ValueError("memmap is required")
83 if not memmap_index:
84 raise ValueError("memmap_index is required")
86 from hierarchicalsoftmax import TreeDict
87 import numpy as np
88 from barbet.data import read_memmap
90 self.stack_size = stack_size
92 self.logger.info(f"Loading treedict {treedict}")
93 individual_treedict = TreeDict.load(treedict)
94 self.treedict = TreeDict(
95 classification_tree=individual_treedict.classification_tree
96 )
98 # Sets the loss weighting for the tips
99 if tip_alpha:
100 for tip in self.treedict.classification_tree.leaves:
101 tip.parent.alpha = tip_alpha
103 self.logger.info("Loading memmap")
104 self.accession_to_array_index = defaultdict(list)
105 with open(memmap_index) as f:
106 for key_index, key in enumerate(f):
107 key = key.strip()
108 accession = key.strip().split("/")[0]
110 if len(self.accession_to_array_index[accession]) == 0:
111 self.treedict[accession] = individual_treedict[key]
113 self.accession_to_array_index[accession].append(key_index)
114 count = key_index + 1
115 self.array = read_memmap(memmap, count)
117 # If there's enough memory, then read into RAM
118 if in_memory:
119 self.array = np.array(self.array)
121 self.classification_tree = self.treedict.classification_tree
122 assert self.classification_tree is not None
124 # Get list of gene families
125 family_ids = set()
126 for accession in self.treedict:
127 gene_id = accession.split("/")[-1]
128 family_ids.add(gene_id)
130 @method
131 def model(
132 self,
133 features: int = 768,
134 intermediate_layers: int = 2,
135 growth_factor: float = 2.0,
136 attention_size: int = 512,
137 ) -> "nn.Module":
138 from barbet.models import BarbetModel
140 return BarbetModel(
141 classification_tree=self.classification_tree,
142 features=features,
143 intermediate_layers=intermediate_layers,
144 growth_factor=growth_factor,
145 attention_size=attention_size,
146 )
148 @method
149 def loss_function(self):
150 from hierarchicalsoftmax import HierarchicalSoftmaxLoss
152 return HierarchicalSoftmaxLoss(root=self.classification_tree)
154 @method
155 def metrics(self) -> "list[tuple[str,Metric]]":
156 from hierarchicalsoftmax.metrics import RankAccuracyTorchMetric
157 from barbet.data import RANKS
159 rank_accuracy = RankAccuracyTorchMetric(
160 root=self.classification_tree,
161 ranks={1 + i: rank for i, rank in enumerate(RANKS)},
162 )
164 return [("rank_accuracy", rank_accuracy)]
166 @method
167 def data(
168 self,
169 max_items: int = 0,
170 num_workers: int = 4,
171 validation_partition: int = 0,
172 batch_size: int = 4,
173 test_partition: int = -1,
174 train_all: bool = False,
175 ) -> "Iterable|L.LightningDataModule":
176 from barbet.data import BarbetDataModule
178 return BarbetDataModule(
179 array=self.array,
180 accession_to_array_index=self.accession_to_array_index,
181 treedict=self.treedict,
182 max_items=max_items,
183 batch_size=batch_size,
184 num_workers=num_workers,
185 validation_partition=validation_partition,
186 test_partition=test_partition,
187 stack_size=self.stack_size,
188 train_all=train_all,
189 )
191 @method
192 def module_class(self) :
193 from .modules import BarbetLightningModule
194 return BarbetLightningModule
196 @method
197 def extra_hyperparameters(self, embedding_model: str = "") -> dict:
198 """Extra hyperparameters to save with the module."""
199 assert embedding_model, "Please provide an embedding model."
200 from barbet.embeddings.esm import ESMEmbedding
202 embedding_model = embedding_model.lower()
203 if embedding_model.startswith("esm"):
204 layers = embedding_model[3:].strip()
205 embedding_model = ESMEmbedding()
206 embedding_model.setup(layers=layers)
207 else:
208 raise ValueError(f"Cannot understand embedding model: {embedding_model}")
210 return dict(
211 embedding_model=embedding_model,
212 classification_tree=self.treedict.classification_tree,
213 stack_size=self.stack_size,
214 )
216 @method
217 def prediction_dataloader(
218 self,
219 module,
220 genome_path: Path,
221 markers: dict[str, str],
222 batch_size: int = Param(
223 64, help="The batch size for the prediction dataloader."
224 ),
225 cpus: int = Param(
226 1, param_decls=["--cpus", "-c"], help="The number of CPUs to use for the prediction dataloader."
227 ),
228 dataloader_workers: int = Param(
229 4, help="The number of workers to use for the dataloader."
230 ),
231 repeats: int = Param(
232 2,
233 help="The minimum number of times to use each protein embedding in the prediction.",
234 ),
235 genome_idx: int = Param(1, hidden=True),
236 total_genomes: int = Param(1, hidden=True),
237 **kwargs,
238 ) -> "Iterable":
239 import torch
240 import numpy as np
241 from torch.utils.data import DataLoader
242 from barbet.data import BarbetPredictionDataset
244 # Set PyTorch thread limits
245 torch.set_num_threads(cpus)
247 # Get hyperparameters from checkpoint
248 stack_size = module.hparams.get("stack_size", 32)
249 self.classification_tree = module.hparams.classification_tree
251 # extract domain from the model
252 domain = "ar53" if self.classification_tree.name == "d__Archaea" else "bac120"
254 #######################
255 # Create Embeddings
256 #######################
257 embeddings = []
258 accessions = []
260 fastas = sorted(markers[domain]) # sort for determinism independent of HMMER output order
261 pct = (genome_idx / total_genomes) * 100 if total_genomes else 100.0
262 description = f"[cyan]Embedding ({genome_idx:,}/{total_genomes:,} genomes, {pct:.1f}%)..."
264 embedding_model = module.hparams.embedding_model
266 with Progress(
267 TextColumn("{task.description}"),
268 BarColumn(),
269 TaskProgressColumn(),
270 TimeElapsedColumn(),
271 TimeRemainingColumn(),
272 ) as progress:
273 task = progress.add_task(description, total=len(fastas))
274 for fasta in fastas:
275 # read the fasta file sequence remove the header
276 fasta = Path(fasta)
277 seq = fasta.read_text().split("\n")[1]
278 vector = embedding_model(seq)
279 if vector is not None and not torch.isnan(vector).any():
280 vector = vector.cpu().detach().clone().numpy()
281 embeddings.append(vector)
283 gene_family_id = fasta.stem
284 accession = f"{genome_path.name}/{gene_family_id}"
285 accessions.append(accession)
287 del vector
288 progress.advance(task)
290 embeddings = np.asarray(embeddings).astype(np.float16)
292 self.prediction_dataset = BarbetPredictionDataset(
293 array=embeddings,
294 accessions=accessions,
295 stack_size=stack_size,
296 repeats=repeats,
297 seed=42,
298 )
299 dataloader = DataLoader(
300 self.prediction_dataset,
301 batch_size=batch_size,
302 num_workers=dataloader_workers,
303 shuffle=False,
304 )
306 return dataloader
308 def node_to_str(self, node: "SoftmaxNode") -> str:
309 """
310 Converts the node to a string
311 """
312 return str(node).split(",")[-1].strip()
314 @main(
315 "load_checkpoint",
316 "prediction_trainer",
317 "prediction_dataloader",
318 )
319 def predict(
320 self,
321 input: list[Path] = Param(
322 default=...,
323 param_decls=["--input", "-i"],
324 help="FASTA files, directories of FASTA files, or a TSV file (columns: fasta_path, [genome_id], [translation_table]). Requires genomes to be in individual FASTA file."
325 ),
326 output_dir: Path = Param(
327 default="output",
328 param_decls=["--output-dir", "-o"],
329 help="A path to the output directory."
330 ),
331 output_csv: Path = Param(
332 default=None,
333 help="A path to output the results as a CSV."
334 ),
335 rm_intermediate: bool = Param(
336 default=False,
337 help="If set, remove the intermediate results directory after processing."
338 ),
339 cpus: int = Param(
340 default=1,
341 param_decls=["--cpus", "-c"],
342 help="The number of CPUs to use."
343 ),
344 pfam_db: str = Param(
345 default="https://data.ace.uq.edu.au/public/gtdbtk/release95/markers/pfam/Pfam-A.hmm",
346 help="The Pfam database to use.",
347 ),
348 tigr_db: str = Param(
349 default="https://data.ace.uq.edu.au/public/gtdbtk/release95/markers/tigrfam/tigrfam.hmm",
350 help="The TIGRFAM database to use.",
351 ),
352 **kwargs,
353 ):
354 """Barbet is a tool for assigning taxonomic labels to genomes using Machine Learning."""
355 start_time = time.perf_counter()
356 self.start_time = start_time
357 self.output_dir = Path(output_dir)
358 self.logger = setup_logger(self.output_dir)
360 # import pandas as pd
361 import polars as pl
362 from itertools import chain
363 from barbet.markers import extract_markers_genes
365 # Get list of files & metadata
366 files = []
367 genome_id_map = {} # str(file_path) -> genome_id
368 translation_tables = {} # genome_id -> translation_table (int)
370 if isinstance(input, (str, Path)):
371 input = [Path(input)]
372 else:
373 input = [Path(p) for p in input]
375 assert len(input) > 0, "No input files provided."
377 def _is_tsv_file(p: Path) -> bool:
378 if not p.is_file():
379 return False
380 if p.suffix.lower() in (".tsv", ".tab"):
381 return True
382 if p.suffix.lower() in (".fa", ".fasta", ".fna", ".gz"):
383 return False
384 try:
385 with open(p, "r", encoding="utf-8") as f:
386 first_line = f.readline()
387 if first_line and not first_line.startswith(">") and "\t" in first_line:
388 return True
389 except Exception:
390 pass
391 return False
393 if len(input) == 1 and _is_tsv_file(input[0]):
394 tsv_path = input[0]
395 with open(tsv_path, "r", encoding="utf-8") as f:
396 for line_idx, line in enumerate(f, start=1):
397 line_str = line.strip()
398 if not line_str or line_str.startswith("#"):
399 continue
400 parts = [p.strip() for p in line_str.split("\t")]
401 if not parts or not parts[0]:
402 continue
404 fasta_str = parts[0]
405 fasta_path = Path(fasta_str)
407 # Skip header line if col1 looks like header and file doesn't exist
408 if line_idx == 1 and not fasta_path.exists():
409 col1_lower = fasta_str.lower()
410 if col1_lower in ("path", "fasta", "fasta_path", "file", "genome_path", "genome") or "fasta" in col1_lower or "path" in col1_lower:
411 continue
413 if not fasta_path.exists():
414 raise FileNotFoundError(f"FASTA file specified in TSV does not exist: {fasta_str}")
416 files.append(fasta_path)
418 col2 = parts[1] if len(parts) > 1 and parts[1] else None
419 gid = col2 if col2 else fasta_path.name
420 genome_id_map[str(fasta_path)] = gid
422 if len(parts) > 2 and parts[2]:
423 raw_tt = parts[2]
424 try:
425 tt = int(raw_tt)
426 except ValueError:
427 raise ValueError(
428 f"Invalid translation table '{raw_tt}' at line {line_idx} in {tsv_path}. Must be 4, 11, or 25."
429 )
430 if tt not in (4, 11, 25):
431 raise ValueError(
432 f"Invalid translation table '{tt}' at line {line_idx} in {tsv_path}. Must be either 4, 11, or 25."
433 )
434 translation_tables[gid] = tt
435 else:
436 for path in input:
437 if path.is_dir():
438 for file in chain(
439 path.rglob("*.fa"),
440 path.rglob("*.fasta"),
441 path.rglob("*.fna"),
442 path.rglob("*.fa.gz"),
443 path.rglob("*.fasta.gz"),
444 path.rglob("*.fna.gz"),
445 ):
446 files.append(file)
447 genome_id_map[str(file)] = file.name
448 elif path.is_file():
449 files.append(path)
450 genome_id_map[str(path)] = path.name
452 # Check if any files were found
453 if len(files) == 0:
454 raise ValueError(
455 f"No files found in {input}. Please provide a directory, a list of files, or a TSV file."
456 )
458 # Check if output directory exists
459 output_csv = output_csv or self.output_dir / "barbet-predictions.csv"
460 output_csv = Path(output_csv)
461 output_csv.parent.mkdir(exist_ok=True, parents=True)
462 self.logger.info(
463 f"Writing results for {len(files)} genome{'s' if len(files) > 1 else ''} to '{output_csv}'"
464 )
466 results_dir = self.output_dir / "results"
467 results_dir.mkdir(exist_ok=True, parents=True)
469 ####################
470 # Extract single copy marker genes
471 ####################
472 start_time_markers = time.perf_counter()
473 markers_gene_map = extract_markers_genes(
474 genomes={genome_id_map[str(file)]: str(file) for file in files},
475 out_dir=str(results_dir),
476 cpus=cpus,
477 force=True,
478 pfam_db=self.process_location(pfam_db),
479 tigr_db=self.process_location(tigr_db),
480 translation_tables=translation_tables if translation_tables else None,
481 )
482 extract_markers_time = time.perf_counter() - start_time_markers
483 self.logger.info(
484 f"Total time to extract and identify marker genes: {extract_markers_time:.2f} seconds"
485 )
487 # Load the model
488 module = self.load_checkpoint(**kwargs)
489 trainer = self.prediction_trainer(module, **kwargs)
491 # Make predictions for each file
492 total_df = None
493 total_genomes = len(markers_gene_map)
494 kwargs_dataloader = dict(kwargs)
495 kwargs_dataloader.pop("genome_idx", None)
496 kwargs_dataloader.pop("total_genomes", None)
498 stack_size = module.hparams.get("stack_size", 32)
499 repeats = kwargs.get("repeats", 2)
500 batch_size = kwargs.get("batch_size", 64)
501 dataloader_workers = kwargs.get("dataloader_workers", 4)
503 start_time_embed_classify = time.perf_counter()
505 all_embeddings_list = []
506 all_accessions = []
508 for idx, (gid, maker_genes) in enumerate(markers_gene_map.items(), start=1):
509 genome_path = Path(gid)
510 self.prediction_dataloader(
511 module,
512 genome_path,
513 maker_genes,
514 cpus=cpus,
515 genome_idx=idx,
516 total_genomes=total_genomes,
517 **kwargs_dataloader,
518 )
519 if len(self.prediction_dataset.array) > 0:
520 all_embeddings_list.append(self.prediction_dataset.array)
521 all_accessions.extend(self.prediction_dataset.accessions)
523 if all_embeddings_list:
524 import numpy as np
525 from torch.utils.data import DataLoader
526 from barbet.data import BarbetPredictionDataset
528 all_embeddings_arr = np.concatenate(all_embeddings_list, axis=0)
529 self.prediction_dataset = BarbetPredictionDataset(
530 array=all_embeddings_arr,
531 accessions=all_accessions,
532 stack_size=stack_size,
533 repeats=repeats,
534 seed=42,
535 )
536 combined_dataloader = DataLoader(
537 self.prediction_dataset,
538 batch_size=batch_size,
539 num_workers=dataloader_workers,
540 shuffle=False,
541 )
542 names = [stack.genome for stack in self.prediction_dataset.stacks]
543 module.setup_prediction(self, names)
544 trainer.predict(module, dataloaders=combined_dataloader)
545 total_df = module.results_df
547 embed_classify_time = time.perf_counter() - start_time_embed_classify
548 self.logger.info(
549 f"Total time to embed and classify: {embed_classify_time:.2f} seconds"
550 )
552 if total_df is not None:
553 if output_csv:
554 total_df.write_csv(output_csv)
556 print_polars_df(
557 total_df[["name", "species_prediction", "species_probability", ]],
558 column_names=["Genome", "Species", "Probability"],
559 )
560 self.logger.info(f"Saved to: '{output_csv}'")
561 else:
562 self.logger.warning("No predictions were generated.")
564 if rm_intermediate and results_dir.exists():
565 import shutil
566 shutil.rmtree(results_dir)
568 total_time = time.perf_counter() - start_time
569 self.logger.info(f"Total time: {total_time:.2f} seconds")
570 self.logger.info("Done")
571 return total_df
573 @tool(
574 "load_checkpoint",
575 "prediction_trainer",
576 "prediction_dataloader_memmap",
577 )
578 def predict_memmap(
579 self,
580 output_csv: Path = Param(
581 default=None, help="A path to output the results as a CSV."
582 ),
583 treedict:Path = Param(None, help="A path to a TreeDict with the ground truth lineage."),
584 probabilities: bool = Param(
585 default=False, help="If True, include probabilities for all the nodes in the taxonomic tree."
586 ),
587 **kwargs,
588 ):
589 """Barbet is a tool for assigning taxonomic labels to genomes using Machine Learning."""
590 start_time = time.perf_counter()
591 self.start_time = start_time
592 module = self.load_checkpoint(**kwargs)
593 trainer = self.prediction_trainer(module, **kwargs)
594 prediction_dataloader = self.prediction_dataloader_memmap(module, **kwargs)
596 module.setup_prediction(self, [stack.genome for stack in self.prediction_dataset.stacks], save_probabilities=probabilities)
597 trainer.predict(module, dataloaders=prediction_dataloader, return_predictions=False)
598 results_df = module.results_df
600 genome_name_set = set(results_df['name'].unique())
602 if treedict is not None:
603 from hierarchicalsoftmax import TreeDict
604 from barbet.data import RANKS
605 import polars as pl
607 true_values = defaultdict(dict)
609 self.logger.info(f"Adding true values from TreeDict '{treedict}'")
610 treedict = TreeDict.load(treedict)
612 # Get lineage to map
613 for accession in track(treedict.keys()):
614 genome_name = accession.split("/")[0]
615 if genome_name in genome_name_set:
616 node = treedict.node(accession)
617 lineage = node.ancestors[1:] + (node,)
618 for rank, lineage_node in zip(RANKS, lineage):
619 true_values[rank][genome_name] = lineage_node.name.strip()
622 for rank in RANKS:
623 results_df = results_df.with_columns(
624 pl.col("name").map_elements(true_values[rank].get, return_dtype=pl.Utf8).alias(f"{rank}_true")
625 )
627 self.logger.info(f"Writing to '{output_csv}'")
628 output_csv = Path(output_csv)
629 output_csv.parent.mkdir(exist_ok=True, parents=True)
630 results_df.write_csv(output_csv)
632 total_time = time.perf_counter() - start_time
633 self.logger.info(f"Total time: {total_time:.2f} seconds")
634 self.logger.info("Done")
636 return results_df
638 @method
639 def prediction_dataloader_memmap(
640 self,
641 module,
642 memmap:Path = Param(None, help="A path to the memmap file containing the protein embeddings."),
643 memmap_index:Path = Param(None, help="A path to the memmap index file containing the accessions."),
644 batch_size: int = Param(
645 64, help="The batch size for the prediction dataloader."
646 ),
647 num_workers: int = 4,
648 repeats: int = Param(
649 2,
650 help="The minimum number of times to use each protein embedding in the prediction.",
651 ),
652 genomes:Path=Param(None, help="A path to a text file with the accessions for the genome to use."),
653 **kwargs,
654 ) -> "Iterable":
655 from barbet.data import read_memmap
656 from torch.utils.data import DataLoader
657 from barbet.data import BarbetPredictionDataset
659 assert memmap is not None, "Please provide a path to the memmap file."
660 assert memmap.exists(), f"Memmap file does not exist: {memmap}"
661 assert memmap_index is not None, "Please provide a path to the memmap index file."
662 assert memmap_index.exists(), f"Memmap index file does not exist: {memmap_index}"
664 # Read the memmap array index
665 self.logger.info(f"Reading memmap array index '{memmap_index}'")
666 accessions = memmap_index.read_text().strip().split("\n")
667 count = len(accessions)
668 self.logger.info(f"Found {count} accessions")
670 # Load the memmap array itself
671 self.logger.info(f"Loading memmap array '{memmap}'")
672 array = read_memmap(memmap, count)
674 # Get hyperparameters from checkpoint
675 self.classification_tree = module.hparams.classification_tree
676 stack_size = module.hparams.get("stack_size", 32)
678 # If treedict is provided, then we filter the accessions to only those that are in the treedict
679 genome_filter = None
680 if genomes:
681 assert genomes.exists(), f"Genomes file does not exist: {genomes}"
682 genome_filter = set(Path(genomes).read_text().strip().split("\n"))
684 self.prediction_dataset = BarbetPredictionDataset(
685 array=array,
686 accessions=accessions,
687 stack_size=stack_size,
688 repeats=repeats,
689 genome_filter=genome_filter,
690 seed=42,
691 )
692 dataloader = DataLoader(
693 self.prediction_dataset,
694 batch_size=batch_size,
695 num_workers=num_workers,
696 shuffle=False,
697 )
699 return dataloader
701 @method
702 def monitor(
703 self,
704 train_all: bool = False,
705 **kwargs,
706 ) -> str:
707 if train_all:
708 return "valid_loss"
709 return "genus"
711 def checkpoint(
712 self,
713 checkpoint:Path=Param(None, help="The path to a checkpoint file for the Barbet parameters. If not provided, then it will use a standard checkpoint."),
714 large:bool=Param(False, help="Whether or not to use the large standard checkpoint of the Barbet parameters."),
715 archaea:bool=Param(False, help="Whether or not to use the standard model for archaea. If not, then it uses the default model for bacteria."),
716 ) -> str:
717 if checkpoint:
718 return checkpoint
720 # Weights are here: https://figshare.unimelb.edu.au/articles/dataset/Trained_weights_for_Barbet/
721 # DOI: https://doi.org/10.26188/29578964
723 if archaea:
724 if large:
725 # barbet-ar53-ESM12-large.ckpt
726 return "https://figshare.unimelb.edu.au/ndownloader/files/56332160"
728 # barbet-ar53-ESM12-base.ckpt
729 return "https://figshare.unimelb.edu.au/ndownloader/files/56332157"
731 if large:
732 # barbet-bac120-ESM6-large.ckpt
733 return "https://figshare.unimelb.edu.au/ndownloader/files/56307647"
735 # barbet-bac120-ESM6-base.ckpt
736 return "https://figshare.unimelb.edu.au/ndownloader/files/56307671"