AI4PD/ZymCTRL · Hit handling for bulk generation

Hey there,
thanks for the very easy to use tool. To make the handling of output files as straight forward I wanted to share a quick way of handling the many output files: The script is extracting the lowest perplexity hits for each EC number (in the folder_path) and writes them as lowest_perplexity_{ec_number_batch}.fasta:

import os
from collections import defaultdict
from Bio import SeqIO

# Folder containing .fasta files
folder_path = "Your_path_goes_here"

# Dictionary to store records grouped by EC number batch
batch_records = defaultdict(list)

# Iterate through all .fasta files in the folder
for filename in os.listdir(folder_path):
    if filename.endswith("_0.fasta"):
        file_path = os.path.join(folder_path, filename)
        with open(file_path, "r") as fasta_file:
            for record in SeqIO.parse(fasta_file, "fasta"):
                # Extract perplexity value from the record description
                description_parts = record.description.split()
                perplexity = float(description_parts[-1])

                # Extract EC number batch from the filename
                ec_number_batch = filename.split("_")[0]

                # Store the record and perplexity value as a tuple in the dictionary
                batch_records[ec_number_batch].append((record, perplexity))

# Iterate through the batches and select the record with the lowest perplexity
for ec_number_batch, records in batch_records.items():
    records.sort(key=lambda x: x[1])  # Sort records by perplexity
    lowest_perplexity_record = records[0][0]  # Get the record with the lowest perplexity

    # Write the lowest perpl-exity record to a file for each batch
    output_file_path = f"{folder_path}/lowest_perplexity_{ec_number_batch}.fasta"
    with open(output_file_path, "w") as output_file:
        SeqIO.write(lowest_perplexity_record, output_file, "fasta")