chklovski / CheckM2

Assessing the quality of metagenome-derived genome bins using machine learning
GNU General Public License v3.0
163 stars 19 forks source link

[Question] How to determine the number of splits used for Diamond annotation? #88

Open jolespin opened 9 months ago

jolespin commented 9 months ago

Let's say someone ran the following command:

checkm2 predict  --genes --threads 16 --database_path uniref100.KO.1.dmnd --output_directory checkm2_output  --input genomes/ -x faa

In checkm2_output/diamond_output there are several diamond output files. Is there way I can compute the expected number of diamond files based on the number genomes and/or genes in each fasta file?

jolespin commented 9 months ago

I'm seeing this bit here:

    def run(self, protein_files):

        logging.info('Annotating input genomes with DIAMOND using {} threads'.format(self.threads))

        if len(protein_files) <= self.chunksize:
            protein_chunks = self.__concatenate_proteins(protein_files)
            diamond_out = os.path.join(self.diamond_out, "DIAMOND_RESULTS.tsv")
            self.__call_diamond(protein_chunks, diamond_out)

        else:
            #break file list into chunks of size 'chunksize'
            chunk_list = [protein_files[i:i + self.chunksize] for i in range(0, len(protein_files), self.chunksize)]

            for number, chunk in enumerate(chunk_list):
                diamond_out = os.path.join(self.diamond_out, "DIAMOND_RESULTS_{}.tsv".format(number))
                self.__call_diamond(self.__concatenate_proteins(chunk), diamond_out)

        diamond_out_list = [x for x in os.listdir(self.diamond_out) if x.startswith('DIAMOND_RESULTS')]
        if len(diamond_out_list) == 0:
            logging.error("Error: DIAMOND failed to generate output.")
            sys.exit(1)
        else:
            return diamond_out_list

Could be REALLY useful to have something like this:

    from tqdm import tqdm
    def run(self, protein_files):

        logging.info('Annotating input genomes with DIAMOND using {} threads'.format(self.threads))

        if len(protein_files) <= self.chunksize:
            protein_chunks = self.__concatenate_proteins(protein_files)
            diamond_out = os.path.join(self.diamond_out, "DIAMOND_RESULTS.tsv")
            self.__call_diamond(protein_chunks, diamond_out)

        else:
            #break file list into chunks of size 'chunksize'
            chunk_list = [protein_files[i:i + self.chunksize] for i in range(0, len(protein_files), self.chunksize)]

            for number, chunk in tqdm(enumerate(chunk_list), desc="Diamond annotation on chunks", total=len(chunk_list)):
                diamond_out = os.path.join(self.diamond_out, "DIAMOND_RESULTS_{}.tsv".format(number))
                self.__call_diamond(self.__concatenate_proteins(chunk), diamond_out)

        diamond_out_list = [x for x in os.listdir(self.diamond_out) if x.startswith('DIAMOND_RESULTS')]
        if len(diamond_out_list) == 0:
            logging.error("Error: DIAMOND failed to generate output.")
            sys.exit(1)
        else:
            return diamond_out_list