rename results fields, modularize event handling

app.py

@@ -11,7 +11,7 @@ from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
 from utils import download_weights
 from utils.pipelines import test_rfd3_from_notebook, unconditional_generation
 #from gradio_molecule3d import Molecule3D
-
+from utils.handle_events import *
 
 download_weights()
 
@@ -56,7 +56,7 @@ with gr.Blocks(title="RFD3 Test") as demo:
     gen_btn = gr.Button("Run Unconditional Generation")
 
 
-    # New visualize section
+    # Section to inspect PDB of generated structures
     with gr.Row():
         batch_dropdown = gr.Dropdown(
             choices=[], 
@@ -68,46 +68,21 @@ with gr.Blocks(title="RFD3 Test") as demo:
             label="Select Design",
             visible=True
         )
-        visualize_btn = gr.Button("Visualize", visible=True)
+        show_pdb_btn = gr.Button("Show PDB content", visible=True)
         
     display_state = gr.Textbox(label="Selected Batch and Design", visible=True)
     display_state.value = "Please Select a Batch and Design number to show sequence"
-    #viewer = Molecule3D(visible=True)
 
-    def update_batch_choices(result):
-        if result is None:
-            return gr.Dropdown(choices = [])
-        batches = sorted(list({d["batch"] for d in result}))
-        return gr.update(choices=batches, visible=True)
         
     gen_btn.click(unconditional_generation, inputs=[num_batches, num_designs_per_batch, length], outputs=[gen_directory, gen_results]).then(
         update_batch_choices,
         inputs=gen_results,
         outputs=batch_dropdown)
 
-
-
-    def update_designs(batch, result):
-        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)
-    
-
-    
     batch_dropdown.change(update_designs, inputs=[batch_dropdown, gen_results], outputs=[design_dropdown])
     design_dropdown.change()
-
-    def visualize_selection(batch, design, result):
-        if batch is None or design is None:
-            return gr.update()
-        pdb_path= next(d["pdb"] 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()
-        return gr.update(value=f"Selected Batch: {batch}, Design: {design}, saved at {pdb_str}:\n {pdb_str}", visible=True)
-
     
-    visualize_btn.click(visualize_selection, inputs=[batch_dropdown, design_dropdown, gen_results], outputs=display_state)
+    show_pdb_btn.click(show_pdb, inputs=[batch_dropdown, design_dropdown, gen_results], outputs=display_state)
 
 
     #def load_viewer(batch, design, result):

utils/handle_events.py

@@ -0,0 +1,56 @@
+import gradio as gr
+
+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"] 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()
+    return gr.update(value=f"Selected Batch: {batch}, Design: {design}, saved at {pdb_str}:\n {pdb_str}", visible=True)

utils/pipelines.py

@@ -43,7 +43,6 @@ def test_rfd3_from_notebook():
         return f"Error: {str(e)}"
     
 
-
 # Initialize engine and run generation
 @spaces.GPU(duration=300) 
 def unconditional_generation(num_batches, num_designs_per_batch, length):
@@ -72,7 +71,7 @@ def unconditional_generation(num_batches, num_designs_per_batch, length):
         for batch in range(num_batches):
             for design in range(num_designs_per_batch):
                 file_name = os.path.join(directory, f"_{batch}_model_{design}.cif.gz")
-                results.append({"batch": batch, "design": design, "file": file_name, "pdb": mcif_gz_to_pdb_string_gemmi(file_name)})
+                results.append({"batch": batch, "design": design, "cif_path": file_name, "pdb_path": mcif_gz_to_pdb(file_name)})
             
         print(results)
         return directory, results
@@ -89,17 +88,18 @@ def collect_outputs(gen_directory, num_batches, num_designs_per_batch):
         return f"Error: {str(e)}"
     
 
-def mcif_gz_to_pdb_string_gemmi(file_path: str) -> str:
+def mcif_gz_to_pdb(file_path: str) -> str:
     """
-    Converts a .mcif.gz file to a PDB-formatted string.
-    
-    Args:
-        file_path (str): Path to the .mcif.gz file.
+    Converts a .mcif.gz file to pdb and saves it to the same directory. Returns the path to the pdb file.
     
-    Returns:
-        str: PDB content as string.
+    Parameters:
+    ----------
+    file_path: str,
+        Path to the .mcif.gz file.
     
-    Requires: pip install gemmi
+    Returns
+    -------
+        str: path to the generated pdb file.
     """
     st = gemmi.read_structure(file_path)
     st.setup_entities()  # Recommended for consistent entity handling [web:18]