Add plugin config, small UI improvements (#14)

* Add plugin config, small UI improvements

* Fix workflow diff

* Reformat

* Add link to job to pools table, visualization tweaks

* Return dict in refolding input

* Update ovo/app/components/input_components.py

Co-authored-by: Copilot <[email protected]>

* Update ovo/app/pages/designs/explorer.py

Co-authored-by: Copilot <[email protected]>

* Update ovo/core/utils/formatting.py

Co-authored-by: Copilot <[email protected]>

* Fix param metadata loop

---------

Co-authored-by: Copilot <[email protected]>

Author: Prihoda, David

ovo/app/components/input_components.py

@@ -6,20 +6,20 @@
 from ovo.app.components.molstar_custom_component import molstar_custom_component, StructureVisualization
 from ovo.app.components.molstar_custom_component.dataclasses import ContigSegment
 from ovo.core.database import WorkflowTypes, Workflow
-from ovo.core.utils.residue_selection import parse_selections
+from ovo.core.utils.residue_selection import parse_selections, from_segments_to_hotspots
 from ovo.core.utils.formatting import safe_filename
 from ovo.core.utils.pdb import get_pdb, filter_pdb_str
 
 from ovo import storage
 from ovo.core.utils.pdb import add_glycan_to_pdb
 
 
-def initialize_workflow(page_key: str, workflow_name: str, include_subclasses=False) -> Workflow:
+def initialize_workflow(page_key: str, workflow_name: str, include_subclasses: bool = False) -> Workflow:
     """Initialize a workflow in the session state based on dropdown with available workflow variants (subclasses).
 
     :param page_key: Key identifying the workflow page, used as key in session_state.workflows
     :param workflow_name: Base name of the workflow (e.g. "rfdiffusion-end-to-end")
-
+    :param include_subclasses: Whether to let user choose from subclasses of the base workflow as variants
     """
     variant_names = WorkflowTypes.get_subclass_names(workflow_name) if include_subclasses else [workflow_name]
     if len(variant_names) > 1:
@@ -47,8 +47,14 @@ def initialize_workflow(page_key: str, workflow_name: str, include_subclasses=Fa
 
 def pdb_input_component(old_pdb_code: str | None) -> tuple[str, bytes] | None:
     st.write("Enter PDB code, UniProt ID or upload your own structure file")
-    with st.columns(2)[0]:
-        new_pdb_code = st.text_input("PDB code or UniProt ID", value=old_pdb_code, key="input_pdb_code")
+    with st.container(horizontal=True, vertical_alignment="bottom"):
+        new_pdb_code = st.text_input(
+            "Download by ID",
+            placeholder="Input PDB code or UniProt ID...",
+            value=old_pdb_code,
+            key="input_pdb_code",
+            width=270,
+        )
         st.button("Confirm")
 
     uploaded_file = st.file_uploader("...or upload a PDB file", type=["pdb", "pdb1"], key="input_pdb_file")
@@ -72,9 +78,21 @@ def sequence_selection_fragment(
     input_name: str,
     color="chain-id",
     representation_type="cartoon+ball-and-stick",
+    write_segments: bool = True,
     fixed_segments: list[ContigSegment] | None = None,
     **selection_kwargs,
 ):
+    """Fragment to display input structure and allow residue selection via molstar component.
+
+    :param page_key: Key identifying the workflow page, typically __file__
+    :param input_name: Name of the input structure, used for download filename
+    :param color: Initial color scheme for molstar component
+    :param representation_type: Initial representation type for molstar component
+    :param write_segments: Whether to write selected segments (True) or individual residues (False)
+    :param fixed_segments: If provided, only these segments will be shown in the structure (others hidden)
+    :param selection_kwargs: Additional kwargs passed to workflow.get/set_selected_segments methods
+    """
+
     workflow = st.session_state.workflows[page_key]
 
     if page_key not in st.session_state.selection_history:
@@ -85,25 +103,11 @@ def sequence_selection_fragment(
 
     selection_history = st.session_state.selection_history[page_key][input_name]
 
-    st.write(
-        "Click residues in sequence or structure to add them to the selection. Shift+Click to select a residue range."
-    )
-
-    force_reload = False
-    if (
-        st.button(":material/undo: Undo selection", disabled=len(selection_history) <= 1)
-        and len(selection_history) >= 2
-    ):
-        # Go one selection backwards
-        selection_history.pop(-1)
-        previous_value = selection_history.pop(-1)
-        st.write(f"Setting selection back to: {', '.join(previous_value)}")
-        workflow.set_selected_segments(previous_value, **selection_kwargs)
-        force_reload = True
-
     pdb_input_string = storage.read_file_str(workflow.get_input_pdb_path())
 
-    if fixed_segments and not st.checkbox("Show full input structure"):
+    if fixed_segments and not st.toggle(
+        "Show full input structure", help="By default, only the fixed input segments are shown."
+    ):
         pdb_input_string = filter_pdb_str(
             pdb_input_string,
             segments=[seg.value for seg in fixed_segments],
@@ -115,8 +119,9 @@ def sequence_selection_fragment(
         if st.toggle(
             "Show glycosylation sites",
             key="show_glycosylation_sites_input",
-            help="The positions of glycans are strictly based on the 5eli NAG template; "
-            "nonexistent links might be displayed.",
+            help="Visualize glycosylation sites on the structure based on glycosylation motifs in the sequence (NXS, NXT, X = any amino acid except P). "
+            "The positions of glycans are approximate, based on the 5eli NAG template, "
+            "nonexistent bonds might be displayed.",
         ):
             pdb_input_string, glycosylated_residues = add_glycan_to_pdb(pdb_input_string)
             if glycosylated_residues:
@@ -153,6 +158,23 @@ def sequence_selection_fragment(
             label_visibility="collapsed",
         )
 
+    with st.container(horizontal=True, vertical_alignment="center"):
+        force_reload = False
+        if (
+            st.button(":material/undo: Undo selection", disabled=len(selection_history) <= 1)
+            and len(selection_history) >= 2
+        ):
+            # Go one selection backwards
+            selection_history.pop(-1)
+            previous_value = selection_history.pop(-1)
+            workflow.set_selected_segments(previous_value, **selection_kwargs)
+            force_reload = True
+
+        selection_container = st.empty()
+        selection_container.write(
+            "Click residues in sequence or structure to add them to the selection. Shift+Click to select a residue range."
+        )
+
     selected_str = molstar_custom_component(
         structures=[
             StructureVisualization(
@@ -186,4 +208,9 @@ def sequence_selection_fragment(
     if not force_reload and selection and (not selection_history or selection != selection_history[-1]):
         selection_history.append(selection)
 
-    st.write(f"Selected segments: {', '.join(selection)}")
+    if selection:
+        if write_segments:
+            selection_container.write(f"Selected segments: {'/'.join(selection)}")
+        else:
+            residues = from_segments_to_hotspots(selection)
+            selection_container.write(f"Selected residues: {residues}")

ovo/app/components/molstar_custom_component/contigs_parser.py

@@ -172,28 +172,29 @@ def parse_contigs_trb(self, contigs_trb: str | dict = "") -> list[ContigSegment]
 
         contigs = "/0 ".join(newContigs)
 
-        reference_tuples = []
         if "complex_con_ref_pdb_idx" in trb_file:
-            for tuple in trb_file["complex_con_ref_pdb_idx"]:
-                newList = [tuple[0], tuple[1], False]
-                reference_tuples.append(newList)
+            ref_idx = trb_file["complex_con_ref_pdb_idx"]
         else:
             assert "con_ref_pdb_idx" in trb_file, "Could not find con_ref_pdb_idx in the .trb file"
-            for tuple in trb_file["con_ref_pdb_idx"]:
-                newList = [tuple[0], tuple[1], False]
-                reference_tuples.append(newList)
+            ref_idx = trb_file["con_ref_pdb_idx"]
 
-        # Do the same for the mapped tuples.
-        mapped_tuples = []
         if "complex_con_hal_pdb_idx" in trb_file:
-            for tuple in trb_file["complex_con_hal_pdb_idx"]:
-                newList = [tuple[0], tuple[1], False]
-                mapped_tuples.append(newList)
+            hal_idx = trb_file["complex_con_hal_pdb_idx"]
         else:
-            assert "con_hal_pdb_idx" in trb_file, "Could not find con_ref_pdb_idx in the .trb file"
-            for tuple in trb_file["con_hal_pdb_idx"]:
-                newList = [tuple[0], tuple[1], False]
-                mapped_tuples.append(newList)
+            assert "con_hal_pdb_idx" in trb_file, "Could not find con_hal_pdb_idx in the .trb file"
+            hal_idx = trb_file["con_hal_pdb_idx"]
+
+        return self.parse_contigs_ref(contigs, ref_idx=ref_idx, hal_idx=hal_idx)
+
+    def parse_contigs_ref(self, contigs: str, ref_idx: list[tuple], hal_idx: list[tuple]):
+        """Parse the contigs and return segments with output numbering."""
+        reference_tuples = []
+        for tuple in ref_idx:
+            reference_tuples.append([tuple[0], tuple[1], False])
+
+        mapped_tuples = []
+        for tuple in hal_idx:
+            mapped_tuples.append([tuple[0], tuple[1], False])
 
         return self._parse_contigs_with_tuples(
             contigs=contigs, mapped_tuples=mapped_tuples, reference_tuples=reference_tuples, include_generated=True

ovo/app/components/navigation.py

@@ -226,7 +226,7 @@ def project_round_selector(
     return (selected_round_ids, None) if allow_design_input else selected_round_ids
 
 
-def pool_selector_table(pools_table: pd.DataFrame) -> list[str]:
+def pool_selector_table(pools_table: pd.DataFrame, project_id: str) -> list[str]:
     key = "_".join(pools_table["ID"])
 
     # Prepare table for data editor
@@ -248,6 +248,10 @@ def pool_selector_table(pools_table: pd.DataFrame) -> list[str]:
         default_selection = st.query_params.get("pool_ids").split(",") if st.query_params.get("pool_ids") else []
         st.session_state[default_selection_key] = default_selection
     pools_table.insert(0, "Selected", [pool_id in default_selection for pool_id in pools_table["ID"]])
+    pools_table["Job ID"] = [
+        f"./jobs?pool_ids={pool_id}&project_id={project_id}" if not pd.isna(job_id) else None
+        for pool_id, job_id in pools_table.set_index("ID")["Job ID"].items()
+    ]
 
     edited = st.data_editor(
         pools_table,
@@ -257,6 +261,7 @@ def pool_selector_table(pools_table: pd.DataFrame) -> list[str]:
         disabled=data_columns,
         column_config={
             "Selected": "",
+            "Job ID": st.column_config.LinkColumn(display_text=r"Job ↑", disabled=True),
         },
     )
 

ovo/app/components/submission_components.py

@@ -25,8 +25,6 @@ def pool_submission_inputs(page_key: str):
     rounds_by_id = get_or_create_project_rounds(project_id=st.session_state.project.id)
     round_ids = list(rounds_by_id.keys())
 
-    left, right, _ = st.columns([1, 1, 2], vertical_alignment="bottom")
-
     if st.session_state.pool_inputs.get(page_key):
         round_id, pool_name, pool_description = st.session_state.pool_inputs[page_key]
         if round_id is None:
@@ -39,6 +37,7 @@ def pool_submission_inputs(page_key: str):
         round_id = st.session_state["new_round_id"]
         del st.session_state["new_round_id"]
 
+    left, right, _ = st.columns([1, 1, 2], vertical_alignment="bottom")
     with left:
         round_id = st.selectbox(
             "Project Round",

ovo/app/components/workflow_visualization_components.py

@@ -4,6 +4,7 @@
 
 from ovo import db, storage
 from ovo.app.components.custom_elements import wrapped_columns
+from ovo.app.components.molstar_custom_component.dataclasses import ContigSegment
 from ovo.core.database import (
     Design,
     descriptors,
@@ -203,6 +204,8 @@ def rfdiffusion_scaffold_design_visualization(design_id: str | None):
         elif prediction_descriptor.b_factor_value == "fractional_plddt":
             pdb_str = pdb_to_mmcif(prediction_pdb, "-", True, fractional_plddt=True)
             description += " (colored by pLDDT confidence)"
+        else:
+            pdb_str = prediction_pdb
 
         molstar_custom_component(
             structures=[
@@ -718,6 +721,27 @@ def visualize_rfdiffusion_design_sequence(design_id: str):
         if workflow.rfdiffusion_params.inpaint_seq
         else []
     )
+    visualize_scaffold_alignment(
+        input_seq_by_resno=input_seq_by_resno,
+        designed_sequences=designed_sequences,
+        parsed_segments=parsed_segments,
+        inpainted_positions=inpainted_positions,
+    )
+
+
+def visualize_scaffold_alignment(
+    input_seq_by_resno: dict[str, dict[str, str]],
+    designed_sequences: dict[str, str],
+    parsed_segments: list[ContigSegment],
+    inpainted_positions: list[str],
+):
+    """Visualize alignment between input and designed sequences based on contig segments.
+
+    :param input_seq_by_resno: Mapping from chain ID -> residue number (as string) -> amino acid.
+    :param designed_sequences: Mapping from chain ID -> designed sequence.
+    :param parsed_segments: List of ContigSegment objects representing the contig segments.
+    :param inpainted_positions: List of positions (e.g., 'A12') that were inpainted.
+    """
     html = [
         f'<div style="font-family: monospace; font-size: 14px; margin: 10px 2px; display: inline-block;">'
         f"&nbsp;&nbsp;&nbsp;&nbsp;Region<br />"

ovo/app/pages/bindcraft/bindcraft_binder_design.py

@@ -160,7 +160,7 @@ def hotspots_step():
         st.error("Please provide an input structure in the input structure step.")
         return
 
-    sequence_selection_fragment(__file__, workflow.input_name, color="hydrophobicity")
+    sequence_selection_fragment(__file__, workflow.input_name, color="hydrophobicity", write_segments=False)
 
 
 @st.fragment()

ovo/app/pages/designs/designs.py

@@ -45,7 +45,7 @@
             st.write("No pools created yet in this round")
             st.stop()
 
-    selected_pool_ids = pool_selector_table(pools_table)
+    selected_pool_ids = pool_selector_table(pools_table, st.session_state.project.id)
 
 views = {
     "🔵 Explorer": "ovo.app.pages.designs.explorer:explorer_fragment",

ovo/app/pages/designs/explorer.py

@@ -9,7 +9,7 @@
 )
 from ovo.app.components.navigation import design_navigation_selector
 from ovo.app.utils.cached_db import get_cached_design, get_cached_pool, get_cached_design_job
-from ovo.core.database.models import Design, Pool, DesignWorkflow, WorkflowTypes
+from ovo.core.database.models import Design, Pool, DesignWorkflow, WorkflowTypes, UnknownWorkflow
 from ovo.app.utils.cached_db import get_cached_pools, get_cached_design_jobs
 
 
@@ -70,6 +70,12 @@ def design_visualization_fragment(selected_design_ids: list[str]):
     design_job = get_cached_design_job(pool.design_job_id) if pool.design_job_id else None
     WorkflowType = type(design_job.workflow) if design_job and design_job.workflow else DesignWorkflow
 
+    if design_job and design_job.workflow.is_instance(UnknownWorkflow):
+        st.warning(
+            f"Falling back to basic structure and sequence visualization, "
+            f"failed to load workflow information: {design_job.workflow.error}"
+        )
+
     WorkflowType.visualize_single_design_structures(design_id)
 
     st.markdown("### Sequence")

ovo/app/pages/rfdiffusion/binder_design.py

@@ -235,7 +235,7 @@ def hotspots_step():
         st.error("Please provide an input structure in the input structure step.")
         return
 
-    sequence_selection_fragment(__file__, workflow.input_name, color="hydrophobicity")
+    sequence_selection_fragment(__file__, workflow.input_name, color="hydrophobicity", write_segments=False)
 
 
 @st.fragment()

ovo/app/pages/rfdiffusion/binder_diversification.py

@@ -339,7 +339,9 @@ def selection_step():
         )
         return
 
-    sequence_selection_fragment(__file__, workflow.input_name, contig_index=0, partial_diffusion=True)
+    sequence_selection_fragment(
+        __file__, workflow.input_name, write_segments=False, contig_index=0, partial_diffusion=True
+    )
 
 
 @st.fragment()