utils/handle_events.py

import gradio as gr
from utils.handle_files import *
import os

def update_batch_choices(result):
    """
    Update the batch dropdown choices once RFD3 has finished generating structures. results is the output of RFD3.
    Used as event triggered by the completion of RFD3 generation.

    Parameters:
    ----------
    result: list of dicts 
        where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path"
    """
    if result is None:
        return gr.Dropdown(choices = [])
    batches = sorted(list({d["batch"] for d in result}))
    return gr.update(choices=batches, visible=True)

def update_designs(batch, result):
    """
    Update the design dropdown choices once RFD3 has finished generating structures. results is the output of RFD3.
    Used as event triggered by the selection of a batch in the batch dropdown. 
    This logic of trigerring the design dropdown update when a batch is selected comes from the reasoning that different batches may have different numbers of designs.
    There is no reason to think that that will be the case since the number of designs per batch is an input to RFD3, but this might still catch some bugs.

    Parameters:
    ----------
    batch: int
        the batch number selected in the batch dropdown, also takes the batch dropdown object when triggered by an event function.
    result: list of dicts 
        where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path"
    """
    if batch is None:
        return gr.update(choices=[])
    designs = sorted(list({d["design"] for d in result if d["batch"] == batch}))
    return gr.update(choices=designs)

def show_pdb(batch, design, result):
    """
    Show the pdb content of the selected batch and design. 
    Used as event triggered by the show pdb button, when a batch and design are selected in the dropdowns.

    Parameters:
    ----------
    batch: int
        the batch number selected in the batch dropdown, also takes the batch dropdown object when triggered by an event function.
    design: int
        the design number selected in the design dropdown, also takes the design dropdown object when triggered by an event function.
    result: list of dicts
        where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path"
    """
    if batch is None or design is None:
        return gr.update()
    pdb_path= next(d["pdb_path"] for d in result if d["batch"] == int(batch) and d["design"] == int(design))
    with open(pdb_path, 'r') as f:
        pdb_str = f.read()
    print(pdb_path)
    print(pdb_str)
    return gr.update(value=f"Selected Batch: {batch}, Design: {design}, saved at {pdb_path}:\n {pdb_str}", visible=True)

#def validate_config_ready(config_upload):
#    """
#    Once the user presses on the Validate Config button, this function checks whether the config is really ready to go.
#    It is ready if the config file is uploaded.
#
#    Parameters:
#    ----------
#
#    config_upload: gr.File
#        uploaded config file, is None if no file is uploaded
#
#
#    Returns:
#    -------
#    str: message to the user about the status of the config validation
#    str: "upload" if upload, and None if config is not ready for generation.
#    """
#    if config_upload is None:
#        return "Please upload a config file for the generation to be ready.", None
#    return "Config file uploaded!", "upload"
#
#    
#def validate_scaffold_ready_with_file(scaffold_upload):
#    """
#    Once the user presses on the Validate Scaffold button, this function checks whether a file was indeed uploaded. If not, it prompts the user to upload a file or select no scaffold/target for unconditioned generation.
#
#    Parameters:
#    ----------
#    scaffold_upload: gr.File
#        uploaded scaffold file, is None if no file is uploaded
#
#    Returns:
#    -------
#    str: message to the user about the status of the scaffold validation
#    str: status string. "upload" if scaffold file is uploaded and valid, None if not.
#    """
#    if scaffold_upload is not None:
#        return "Scaffold file uploaded and valid!", "upload"
#    else:
#        return "Please upload a scaffold file or press 'No Scaffold/Target' for unconditional generation.", None
#
#
#def give_run_status(num_batches, num_designs_per_batch, config_upload, scaffold_upload):
#    """
#    Once the user presses on the Run Generation button, this function checks whether both config and scaffold are ready, and gives a status message about the generation run that is about to happen. If both config and scaffold are not ready, it prompts the user to ensure they are both validated before running the generation.
#
#    Parameters:
#    ----------
#    num_batches: gr.Number
#        number of batches input by the user manually
#    num_designs_per_batch: gr.Number
#        number of designs per batch input by the user manually
#    config_upload: gr.File
#        uploaded config file, is None if no file is uploaded
#    scaffold_upload: gr.File
#        uploaded scaffold file, is None if no file is uploaded
#
#    Returns:
#    -------
#        gr.update: Gradio update object to update the generation status message in the UI about to be run. If both config and scaffold are not ready, returns an update with a message prompting the user to ensure they are both validated before running the generation.
#
#    
#    """
#    if config_upload is None:        
#        return gr.update(value="Please ensure you have uploaded a configuration file: .yaml or .json")
#    elif scaffold_upload is None:
#        return gr.update(value=f"Running unconditional generation for {num_batches} batches of {num_designs_per_batch}, config uploaded from file {os.path.basename(config_upload)}")
#    else:
#        return gr.update(value=f"Running conditional generation for {num_batches} batches of {num_designs_per_batch}, config uploaded from file {os.path.basename(config_upload)}, scaffold uploaded as {os.path.basename(scaffold_upload)}")