diff --git a/gtranslate/tools.py b/gtranslate/tools.py index 0f78cf8..268f83b 100644 --- a/gtranslate/tools.py +++ b/gtranslate/tools.py @@ -1,15 +1,12 @@ import hashlib -import json import logging import shutil -from typing import Optional import math import os import random import re import time -import urllib.request from itertools import islice from tqdm import tqdm @@ -23,6 +20,29 @@ ############MISC UTILITIES######################## ################################################## +RE_CANONICAL = re.compile(r'^(?:GB_)?(?:RS_)?(?:GCF_)?(?:GCA_)?(\d{9})\.\d') + + +def canonical_gid(gid: str) -> str: + """Get canonical form of NCBI genome accession. + + Example: + G005435135 -> G005435135 + GCF_005435135.1 -> G005435135 + GCF_005435135.1_ASM543513v1_genomic -> G005435135 + RS_GCF_005435135.1 -> G005435135 + GB_GCA_005435135.1 -> + + :param gid: Genome accesion to conver to canonical form. + :return: Canonical form of accession. + """ + + match = RE_CANONICAL.match(gid) + if match: + return f'G{match[1]}' + else: + return gid + def get_genomes_size(genome_path): """Returns the size of a specific genome file.""" diff --git a/gtranslate/training/ground_truth_by_taxonomy.py b/gtranslate/training/ground_truth_by_taxonomy.py deleted file mode 100755 index e211eb2..0000000 --- a/gtranslate/training/ground_truth_by_taxonomy.py +++ /dev/null @@ -1,218 +0,0 @@ -#! /usr/bin/env python3 - -# Determine the ground truth for genomes based on their taxonomic classification. -# -# Genomes are assumed to be classified according to the GTDB taxonomic framework with the -# exception of Hodgkinia cicadicola, Nasuia deltocephalinicola, and Stammera capleta which -# have highly reduced genomes and are not included in GTDB as they lack sufficient phylogenetically -# informative genes to be robustly placed in the reference trees used by this taxonomic resource. -# -# This function requires a taxonomy file that indicates the taxonomic assignment of each genome. This -# can either be a 2 column TSV file with the headers "Genome ID" and "Taxonomy", or a 3 column TSV -# file with the headers "Genome ID", "GTDB taxonomy", and "NCBI taxonomy". Taxonomy strings must be -# in Greengenes-style and indicate all ranks, e.g. -# d__Bacteria;p__Bacillota;c__Bacilli;o__Bacillales;f__Bacillaceae;g__Bacillus;s__Bacillus subtilis - -__prog_name__ = 'ground_truth_by_taxonomy.py' -__prog_desc__ = 'Determine the ground truth for genomes based on their taxonomic classification.' - -__author__ = 'Donovan Parks' -__copyright__ = 'Copyright 2026' -__credits__ = ['Donovan Parks'] -__license__ = 'GPL3' -__version__ = '0.1.1' -__maintainer__ = 'Donovan Parks' -__email__ = 'donovan.parks@gmail.com' -__status__ = 'Development' - - -import logging -import argparse -import gzip -from collections import defaultdict - -from gtranslate.biolib_lite.logger import logger_setup - - -class GroundTruthByTaxonomy(object): - """Determine the ground truth for genomes based on their taxonomic classification.""" - - def __init__(self): - """Initialize.""" - - # Ground truth from GTDB classifications - self.GTDB_TT25 = set(['c__JAEDAM01']) - self.GTDB_TT4 = set(['o__Mycoplasmatales', 's__Zinderia insecticola']) - - # Eggerthellacea genera using table 4; will need to be updated to names in Parks et al., 2026 - # once these appear in GTDB - self.GTDB_TT4.update(set(['g__CAVGFB01', 'g__JAUNQF01'])) - - # Minisyncoccia family identified in gTranslate manuscript that uses table 4. The majority, but not - # all genomes in g__GCA-2747955 were also identified as using table 4. Currently, this is handled by - # explicitly indicating the species in this genus identified as using table 4. - self.GTDB_TT4.update(set(['f__JAKLIH01', 's__GCA-2747955 sp027024305', 's__GCA-2747955 sp027039745', 's__GCA-2747955 sp947311625'])) - - # Must include the Fastidiosibacteraceae XS4 species cluster once (if) this genome appears in GTDB: - # - https://www.ncbi.nlm.nih.gov/nuccore/AP038919.1 - # - https://pmc.ncbi.nlm.nih.gov/articles/PMC12213064 - - # Ground truth from NCBI classifications - self.NCBI_TT4 = set(['s__Candidatus Hodgkinia cicadicola', 's__Candidatus Nasuia deltocephalinicola', 's__Candidatus Stammera capleta']) - self.NCBI_TT4.update(set(['s__Candidatus Organicella extenuata', 's__Candidatus Pinguicoccus supinus'])) - self.NCBI_TT4.update(set(['s__Hodgkinia cicadicola', 's__Nasuia deltocephalinicola', 's__Stammera capleta'])) - self.NCBI_TT4.update(set(['s__Organicella extenuata', 's__Pinguicoccus supinus'])) - - # These species clusters have an unclear ground truth, see https://doi.org/10.1093/gbe/evad164 - self.GTDB_UNRESOLVED = set(['s__Providencia_A siddallii', 's__Providencia_A siddallii_A']) - - self.log = logging.getLogger('timestamp') - - def parse_manual_ground_truth_file(self, manual_gt_file: str) -> dict: - """Parse manual ground truth file.""" - - manual_ground_truth = {} - open_file = gzip.open if manual_gt_file.endswith('.gz') else open - with open_file(manual_gt_file, 'rt') as f: - header = f.readline().strip().split('\t') - gid_idx = header.index("Genome ID") - gt_idx = header.index("Translation table") - - for line in f: - tokens = line.strip().split('\t') - manual_ground_truth[tokens[gid_idx]] = tokens[gt_idx] - - return manual_ground_truth - - def run(self, taxonomy_file: str, manual_gt_file: str, out_file: str) -> None: - """Determine the ground truth for genomes based on their taxonomic classification.""" - - # read files with manually specific ground truth - manual_ground_truth = {} - if manual_gt_file: - self.log.info('Parsing manual ground truth file:') - manual_ground_truth = self.parse_manual_ground_truth_file(manual_gt_file) - self.log.info(f' - identified manual ground truth for {len(manual_ground_truth):,} genomes') - - # determine ground truth for genomes based on their taxonomic classification - self.log.info('Determining ground truth for genomes:') - total_genomes = 0 - gt_table_count = defaultdict(int) - num_by_manual_gt = 0 - - open_file = gzip.open if taxonomy_file.endswith('.gz') else open - with open_file(taxonomy_file, 'rt') as f: - header = f.readline().strip().split('\t') - - if 'Genome ID' in header: - gid_idx = header.index("Genome ID") - else: - self.log.error("Taxonomy file must have a 'Genome ID' column.") - - fout = open(out_file, 'w') - fout.write('Genome ID\tGround truth table') - - taxonomy_idx = None - if 'Taxonomy' in header: - taxonomy_idx = header.index("Taxonomy") - fout.write('\tTaxonomy') - - gtdb_taxonomy_idx = None - if 'GTDB taxonomy' in header: - gtdb_taxonomy_idx = header.index('GTDB taxonomy') - fout.write('\tGTDB taxonomy') - - ncbi_taxonomy_idx = None - if 'NCBI taxonomy' in header: - ncbi_taxonomy_idx = header.index('NCBI taxonomy') - fout.write('\tNCBI taxonomy') - - if taxonomy_idx is None and gtdb_taxonomy_idx is None and ncbi_taxonomy_idx is None: - self.log.error("Taxonomy file must have a 'Taxonomy' column or a 'GTDB taxonomy' and 'NCBI taxonomy' columns.") - - fout.write('\n') - - for line in f: - tokens = line.strip().split('\t') - - total_genomes += 1 - - gid = tokens[gid_idx] - - if gid in manual_ground_truth: - ground_truth_tt = manual_ground_truth[gid] - num_by_manual_gt += 1 - else: - # determine ground truth translation table based on GTDB - # or NCBI taxonomic classification of genome - taxa = set() - if taxonomy_idx: - taxa.update(set(tokens[taxonomy_idx].split(';'))) - - gtdb_taxa = set() - if gtdb_taxonomy_idx: - gtdb_taxa.update(set(tokens[gtdb_taxonomy_idx].split(';'))) - - ncbi_taxa = set() - if ncbi_taxonomy_idx: - ncbi_taxa.update(set(tokens[ncbi_taxonomy_idx].split(';'))) - - if taxa.intersection(self.GTDB_TT25) or gtdb_taxa.intersection(self.GTDB_TT25): - ground_truth_tt = '25' - elif taxa.intersection(self.GTDB_TT4) or gtdb_taxa.intersection(self.GTDB_TT4): - ground_truth_tt = '4' - elif taxa.intersection(self.NCBI_TT4) or ncbi_taxa.intersection(self.NCBI_TT4): - ground_truth_tt = '4' - elif taxa.intersection(self.GTDB_UNRESOLVED) or gtdb_taxa.intersection(self.GTDB_UNRESOLVED): - ground_truth_tt = 'UNRESOLVED' - else: - ground_truth_tt = '11' - - gt_table_count[ground_truth_tt] += 1 - - # write out ground truth results - fout.write(f'{gid}\t{ground_truth_tt}') - - if taxonomy_idx: - fout.write(f'\t{tokens[taxonomy_idx]}') - - if gtdb_taxonomy_idx: - fout.write(f'\t{tokens[gtdb_taxonomy_idx]}') - - if ncbi_taxonomy_idx: - fout.write(f'\t{tokens[ncbi_taxonomy_idx]}') - - fout.write('\n') - - fout.close() - - # write out number of genomes assigned to each translation table - self.log.info(f' - determined ground truth for {total_genomes:,} genomes') - if manual_gt_file: - self.log.info(f' - ground truth set manually for {num_by_manual_gt:,} genomes') - - for tran_table, genome_count in sorted(gt_table_count.items()): - self.log.info(f'Table {tran_table}: {genome_count:,} ({100*genome_count/total_genomes:.2f}%)') - - -def main(): - print(__prog_name__ + ' v' + __version__ + ': ' + __prog_desc__) - print(' by ' + __author__ + ' (' + __email__ + ')' + '\n') - - parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument('--taxonomy_file', required=True, help='File indicating taxonomic classification of each genome.') - parser.add_argument('--out_file', required=True, help='Output file to write ground truth translation table.') - parser.add_argument('--manual_gt_file', help='File indicating manually specific ground truth for select genomes.') - - args = parser.parse_args() - - # setup logger - logger_setup(None, "ground_truth_by_taxonomy.log", "gTranslate", __version__, False, False) - - # run program - p = GroundTruthByTaxonomy() - p.run(args.taxonomy_file, args.manual_gt_file, args.out_file) - - -if __name__ == '__main__': - main() diff --git a/gtranslate/training_manager.py b/gtranslate/training_manager.py index 869635b..deda1ff 100644 --- a/gtranslate/training_manager.py +++ b/gtranslate/training_manager.py @@ -56,7 +56,7 @@ def __init__(self, cpus: int = 1 , seed: int | None = None) -> None: self.seed = seed if seed is not None else np.random.randint(0, 1000) # Ground truth from GTDB classifications - self.GTDB_TT25 = set(['c__JAEDAM01']) + self.GTDB_TT25 = set(['o__Absconditabacterales', 'o__BD1-5']) self.GTDB_TT4 = set(['o__Mycoplasmatales', 's__Zinderia insecticola']) # Eggerthellacea genera using table 4; will need to be updated to names in Parks et al., 2026 @@ -74,8 +74,11 @@ def __init__(self, cpus: int = 1 , seed: int | None = None) -> None: # - https://pmc.ncbi.nlm.nih.gov/articles/PMC12213064 # Ground truth from NCBI classifications - self.NCBI_TT4 = set(['s__Candidatus Hodgkinia cicadicola', 's__Candidatus Nasuia deltocephalinicola', 's__Candidatus Stammera capleta']) - self.NCBI_TT4.update(set(['s__Hodgkinia cicadicola', 's__Nasuia deltocephalinicola', 's__Stammera capleta'])) + self.NCBI_TT4 = set(['s__Candidatus Hodgkinia cicadicola', 's__Candidatus Nasuia deltocephalincola', 's__Candidatus Stammera capleta']) + self.NCBI_TT4.update(set(['s__Hodgkinia cicadicola', 's__Nasuia deltocephalincola', 's__Stammera capleta'])) + + # species appears to have incorrect spelling at NCBI + self.NCBI_TT4.update(set(['s__Candidatus Nasuia deltocephalinicola', 's__Nasuia deltocephalinicola'])) # These species clusters have an unclear ground truth, see https://doi.org/10.1093/gbe/evad164 self.GTDB_UNRESOLVED = set(['s__Providencia_A siddallii', 's__Providencia_A siddallii_A']) diff --git a/scripts/gtdb_training_files/gtdb_r232_manual_ground_truth.tsv b/scripts/gtdb_training_files/gtdb_r232_manual_ground_truth.tsv new file mode 100644 index 0000000..bb56ac8 --- /dev/null +++ b/scripts/gtdb_training_files/gtdb_r232_manual_ground_truth.tsv @@ -0,0 +1,6 @@ +Genome ID Translation table Taxonomy +G036689325 11 d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales_A;f__Enterobacteriaceae_A;g__Stammera;s__Stammera capleta_Y d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales;f__Enterobacteriaceae;g__Candidatus Stammera;s__Candidatus Stammera capleta +G036689305 11 d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales_A;f__Enterobacteriaceae_A;g__Stammera;s__Stammera capleta_AM d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales;f__Enterobacteriaceae;g__Candidatus Stammera;s__Candidatus Stammera capleta +G036689245 11 d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales_A;f__Enterobacteriaceae_A;g__Stammera;s__Stammera capleta_AD d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales;f__Enterobacteriaceae;g__Candidatus Stammera;s__Candidatus Stammera capleta +G036689255 11 d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales_A;f__Enterobacteriaceae_A;g__Stammera;s__Stammera capleta_T d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales;f__Enterobacteriaceae;g__Candidatus Stammera;s__Candidatus Stammera capleta +G036689265 11 d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales_A;f__Enterobacteriaceae_A;g__Stammera;s__Stammera capleta_AF d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales;f__Enterobacteriaceae;g__Candidatus Stammera;s__Candidatus Stammera capleta diff --git a/scripts/gtdb_training_set.py b/scripts/gtdb_training_set.py new file mode 100755 index 0000000..e1b3aa2 --- /dev/null +++ b/scripts/gtdb_training_set.py @@ -0,0 +1,278 @@ +#! /usr/bin/env python3 + +# Creating files required to train gTranslate on a new GTDB release. +# +# Training should be performed on: +# - all GTDB genomes passing standard QC +# - additional genomes with reassigned stop codons that pass a less stringent set of QC criteria + +__prog_name__ = 'gtdb_training_set.py' +__prog_desc__ = 'Creating files required to train gTranslate on a new GTDB release.' + +__author__ = 'Donovan Parks' +__copyright__ = 'Copyright 2026' +__credits__ = ['Donovan Parks'] +__license__ = 'GPL3' +__version__ = '0.1.0' +__maintainer__ = 'Donovan Parks' +__email__ = 'donovan.parks@gmail.com' +__status__ = 'Development' + + +import os +import logging +import argparse +import gzip +from collections import defaultdict +from dataclasses import dataclass + +from gtranslate.biolib_lite.logger import logger_setup +from gtranslate.tools import canonical_gid + + +@dataclass +class GenomeData: + gtdb_taxonomy: str + ncbi_taxonomy: str + + +class GtdbTrainingSet(object): + """Creating files required to train gTranslate on a new GTDB release.""" + + def __init__(self): + """Initialize.""" + + # list of species with reassigned stop codon that should be included in training + # set using more relaxed QC criteria + self.GTDB_REASSIGNED_SP = set(['s__Zinderia insecticola']) + self.NCBI_REASSIGNED_SP = set(['s__Candidatus Hodgkinia cicadicola', 's__Candidatus Nasuia deltocephalincola', 's__Candidatus Stammera capleta']) + self.NCBI_REASSIGNED_SP.update(set(['s__Candidatus Organicella extenuata', 's__Candidatus Pinguicoccus supinus'])) + self.NCBI_REASSIGNED_SP.update(set(['s__Hodgkinia cicadicola', 's__Nasuia deltocephalincola', 's__Stammera capleta'])) + self.NCBI_REASSIGNED_SP.update(set(['s__Organicella extenuata', 's__Pinguicoccus supinus'])) + + # species appears to have incorrect spelling at NCBI + self.NCBI_REASSIGNED_SP.update(set(['s__Candidatus Nasuia deltocephalinicola', 's__Nasuia deltocephalinicola'])) + + # Must include the Fastidiosibacteraceae XS4 species cluster once (if) this genome appears in NCBI Assembly DB: + # - https://www.ncbi.nlm.nih.gov/nuccore/AP038919.1 + # - https://pmc.ncbi.nlm.nih.gov/articles/PMC12213064 + + # QC criteria for above taxa + self.MIN_COMP = 50 + self.MAX_CONT = 10 + self.MIN_QUAL = 50 + + # genomes to retain without needing to pass QC + self.QC_EXCEPTIONS = set(['G015697925']) # s__Pinguicoccus supinus representative genome + + self.log = logging.getLogger('timestamp') + + def parse_metadata_file(self, gtdb_metadata_file: str) -> dict: + """Parse manual ground truth file.""" + + genome_data = {} + open_file = gzip.open if gtdb_metadata_file.endswith('.gz') else open + with open_file(gtdb_metadata_file, 'rt') as f: + header = f.readline().strip().split('\t') + gid_idx = header.index("accession") + gtdb_taxonomy_idx = header.index('gtdb_taxonomy') + ncbi_taxonomy_idx = header.index('ncbi_taxonomy') + + for line in f: + tokens = line.strip().split('\t') + gid = canonical_gid(tokens[gid_idx]) + genome_data[gid] = GenomeData(tokens[gtdb_taxonomy_idx], tokens[ncbi_taxonomy_idx]) + + return genome_data + + def parse_taxonomy_file(self, taxonomy_file: str) -> dict: + """Parse taxonomy file.""" + + taxonomy = {} + open_file = gzip.open if taxonomy_file.endswith('.gz') else open + with open_file(taxonomy_file, 'rt') as f: + for line in f: + tokens = line.strip().split('\t') + gid = canonical_gid(tokens[0]) + taxonomy[gid] = tokens[1] + + return taxonomy + + def parse_qc_failed_file(self, gtdb_qc_failed_file: str, ncbi_taxonomy: dict) -> tuple[dict, dict]: + """Parse genomes with reassigned stop codon that failed the GTDB QC, but should still be included.""" + + genome_data = {} + sp_counts = defaultdict(int) + open_file = gzip.open if gtdb_qc_failed_file.endswith('.gz') else open + with open_file(gtdb_qc_failed_file, 'rt') as f: + header = f.readline().strip().split('\t') + + gid_idx = header.index("Accession") + gtdb_taxonomy_idx = header.index('GTDB taxonomy') + ncbi_sp_idx = header.index('NCBI species') + + cm2_comp_idx = header.index("CM2 completeness (%)") + cm2_cont_idx = header.index("CM2 contamination (%)") + cm2_qual_idx = header.index("CM2 quality") + + for line in f: + tokens = line.strip().split('\t') + + gid = canonical_gid(tokens[gid_idx]) + + comp = float(tokens[cm2_comp_idx]) + cont = float(tokens[cm2_cont_idx]) + qual = float(tokens[cm2_qual_idx]) + + pass_qc = comp >= self.MIN_COMP and cont <= self.MAX_CONT and qual >= self.MIN_QUAL + + if gid not in self.QC_EXCEPTIONS and not pass_qc: + # genome failed even more lenient QC + continue + + gtdb_taxonomy = tokens[gtdb_taxonomy_idx] + gtdb_sp = [t.strip() for t in gtdb_taxonomy.split(';')][-1] + ncbi_sp = tokens[ncbi_sp_idx] + + if ncbi_sp in self.NCBI_REASSIGNED_SP: + assert ncbi_sp.replace('s__', '') in ncbi_taxonomy[gid] + genome_data[gid] = GenomeData(gtdb_taxonomy, ncbi_taxonomy[gid]) + sp_counts[ncbi_sp] += 1 + elif gtdb_sp in self.GTDB_REASSIGNED_SP: + genome_data[gid] = GenomeData(gtdb_taxonomy, ncbi_taxonomy[gid]) + sp_counts[gtdb_sp] += 1 + + return genome_data, sp_counts + + def parse_genome_dir_file(self, genome_dir_file: str, training_gids: dict) -> dict: + """Parse path to genomic FASTA files for training genomes.""" + + genome_paths = {} + open_file = gzip.open if genome_dir_file.endswith('.gz') else open + with open_file(genome_dir_file, 'rt') as f: + for line in f: + tokens = line.strip().split('\t') + + gid = canonical_gid(tokens[0]) + if gid not in training_gids: + continue + + genome_dir = tokens[1] + accn = genome_dir.split('/')[-1] + genome_path = os.path.join(genome_dir, f'{accn}_genomic.fna.gz') + + genome_paths[gid] = genome_path + + return genome_paths + + def run(self, gtdb_bac_metadata_file: str, + gtdb_ar_metadata_file: str, + gtdb_qc_failed: str, + ncbi_taxonomy_file: str, + genome_dir_file: str, + out_dir: str) -> None: + """Creating files required to train gTranslate on a new GTDB release.""" + + # parse genomes from GTDB metadata files + self.log.info('Parsing GTDB metadata files:') + gtdb_bac_gids = self.parse_metadata_file(gtdb_bac_metadata_file) + gtdb_ar_gids = self.parse_metadata_file(gtdb_ar_metadata_file) + gtdb_gids = {**gtdb_bac_gids, **gtdb_ar_gids} + self.log.info(f' - identified {len(gtdb_bac_gids):,} bacterial genomes') + self.log.info(f' - identified {len(gtdb_ar_gids):,} archaeal genomes') + self.log.info(f' - identified {len(gtdb_gids):,} total genomes') + + # read NCBI taxonomy for genomes + self.log.info('Parsing NCBI taxonomy file:') + ncbi_taxonomy = self.parse_taxonomy_file(ncbi_taxonomy_file) + self.log.info(f' - identified taxonomy for {len(ncbi_taxonomy):,} genomes') + + # parse genomes from taxa with reassigned stop codon that failed the GTDB QC, + # but should still be included + self.log.info('Parsing genomes that failed GTDB QC:') + gtdb_reassigned_gids, sp_counts = self.parse_qc_failed_file(gtdb_qc_failed, ncbi_taxonomy) + self.log.info(f' - identified {len(gtdb_reassigned_gids):,} genomes to retain for training') + + for sp, count in sorted(sp_counts.items()): + self.log.info(f' - {sp}: {count}') + + # combine all training genomes + self.log.info('Combining all genomes to use for training:') + training_gids = {**gtdb_gids, **gtdb_reassigned_gids} + self.log.info(f' - identified {len(training_gids):,} total genomes for training') + + # sanity check that all genomes that are a QC expection are accounted for + for gid in self.QC_EXCEPTIONS: + assert gid in training_gids + + # get final count of reassigned species + self.log.info('Tabulating number of genomes in reassigned species:') + reassigned_sp_count = defaultdict(int) + for gid, genome_data in training_gids.items(): + ncbi_sp = [t.strip() for t in genome_data.ncbi_taxonomy.split(';')][-1] + gtdb_sp = [t.strip() for t in genome_data.gtdb_taxonomy.split(';')][-1] + + if ncbi_sp in self.NCBI_REASSIGNED_SP: + reassigned_sp_count[ncbi_sp] += 1 + elif gtdb_sp in self.GTDB_REASSIGNED_SP: + reassigned_sp_count[gtdb_sp] += 1 + + for sp, count in sorted(reassigned_sp_count.items()): + self.log.info(f' - {sp}: {count}') + + # read path to genomic FASTA files + self.log.info('Parsing path to genomic FASTA files for training genomes:') + genome_paths = self.parse_genome_dir_file(genome_dir_file, training_gids) + self.log.info(f' - identified path for {len(training_gids):,} genomes') + assert len(genome_paths) == len(training_gids) + + # create taxonomy file for all training genomes + self.log.info('Creating taxonomy file.') + fout = open(os.path.join(out_dir, 'training_taxonomy.tsv'), 'wt') + fout.write('Genome ID\tGTDB taxonomy\tNCBI taxonomy\n') + for gid, genome_data in training_gids.items(): + fout.write('{}\t{}\t{}\n'.format( + gid, + genome_data.gtdb_taxonomy, + genome_data.ncbi_taxonomy)) + fout.close() + + # create file with path to genomic FASTA file for all training genomes + self.log.info('Creating genome path file.') + fout = open(os.path.join(out_dir, 'training_genome_path.tsv'), 'wt') + for gid in training_gids: + fout.write(f'{genome_paths[gid]}\t{gid}\n') + fout.close() + + self.log.info('Done.') + + +def main(): + print(__prog_name__ + ' v' + __version__ + ': ' + __prog_desc__) + print(' by ' + __author__ + ' (' + __email__ + ')' + '\n') + + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument('--gtdb_bac_metadata_file', required=True, help='GTDB metadata file for bacterial genomes.') + parser.add_argument('--gtdb_ar_metadata_file', required=True, help='GTDB metadata file for bacterial genomes.') + parser.add_argument('--gtdb_qc_failed', required=True, help='File indicating genomes that failed standard GTDB QC criteria.') + parser.add_argument('--ncbi_taxonomy_file', required=True, help='File indicating standardized NCBI taxonomic classification of each genome.') + parser.add_argument('--genome_dir_file', required=True, help='File indicating path to data files for genome assemblies.') + parser.add_argument('--out_dir', help='Output directory.') + + args = parser.parse_args() + + # setup logger + logger_setup(args.out_dir, "gtdb_training_set.log", "gTranslate", __version__, False, False) + + # run program + p = GtdbTrainingSet() + p.run(args.gtdb_bac_metadata_file, + args.gtdb_ar_metadata_file, + args.gtdb_qc_failed, + args.ncbi_taxonomy_file, + args.genome_dir_file, + args.out_dir) + + +if __name__ == '__main__': + main()