Add clickable residue tracks (highlight AAs in 3D)

app.py

@@ -109,7 +109,8 @@ LEGEND_HTML = (
     + _swatch("CDR-H1", CDR_COLORS["H"]["1"]) + _swatch("H2", CDR_COLORS["H"]["2"])
     + _swatch("H3", CDR_COLORS["H"]["3"]) + _swatch("CDR-L1", CDR_COLORS["L"]["1"])
     + _swatch("L2", CDR_COLORS["L"]["2"]) + _swatch("L3", CDR_COLORS["L"]["3"])
-    + _swatch("liability", "magenta") + _swatch("framework", "#cfd8dc")
+    + _swatch("liability", "magenta") + _swatch("your pick", "lime")
+    + _swatch("framework", "#cfd8dc")
     + "</div>"
 )
 
@@ -364,12 +365,30 @@ CHAIN_VIEW = {
     "Light only (VL)": ("L",),
 }
 
+# Per-residue colors for the clickable sequence tracks (CDRs reuse the 3D colors).
+TRACK_COLORS = {
+    "H1": CDR_COLORS["H"]["1"], "H2": CDR_COLORS["H"]["2"], "H3": CDR_COLORS["H"]["3"],
+    "L1": CDR_COLORS["L"]["1"], "L2": CDR_COLORS["L"]["2"], "L3": CDR_COLORS["L"]["3"],
+    "fw": "#eceff1",
+}
+
+
+def build_track(chain_label: str, residues):
+    """Numbered residues -> (HighlightedText tokens, index->imgt map).
+    Each residue is its own token (combine_adjacent=False) so it clicks individually."""
+    tokens, idxmap = [], []
+    for imgt, _ins, aa in residues:
+        cdr = cdr_of(imgt)
+        tokens.append((aa, f"{chain_label}{cdr}" if cdr else "fw"))
+        idxmap.append(imgt)
+    return tokens, idxmap
+
 
 def render_structure(pdb: str, spin: bool = False, highlights=None,
-                     chains=("H", "L")) -> str:
+                     chains=("H", "L"), manual=None) -> str:
     """py3Dmol view -> self-contained HTML in an iframe (renders inside Gradio).
-    Framework grey, CDR loops colored, liability residues as magenta sticks.
-    Only the chains in `chains` are drawn; the rest are hidden."""
+    Framework grey, CDR loops colored, liability residues as magenta sticks,
+    user-clicked residues as lime sticks. Only chains in `chains` are drawn."""
     view = py3Dmol.view(width=760, height=560)
     view.addModel(pdb, "pdb")
     view.setStyle({}, {})  # hide everything; only the selected chains get drawn below
@@ -387,6 +406,13 @@ def render_structure(pdb: str, spin: bool = False, highlights=None,
                     {"chain": chain, "resi": pos},
                     {"stick": {"color": "magenta", "radius": 0.3}},
                 )
+        if manual:  # user-clicked residues, drawn last so they win over liability sticks
+            mpos = manual.get(chain, [])
+            if mpos:
+                view.addStyle(
+                    {"chain": chain, "resi": mpos},
+                    {"stick": {"color": "lime", "radius": 0.3}},
+                )
     view.zoomTo({"chain": list(chains)})  # re-center on the visible chain(s)
     if spin:
         view.spin(True)
@@ -400,9 +426,11 @@ def fold(heavy: str, light: str, chain_choice: str = "Both chains"):
     """Fold VH+VL, scan liabilities, render, and reset the spin toggle."""
     heavy, light = clean(heavy), clean(light)
     reset_btn = gr.update(value="▶ Spin")
+    empty_manual = {"H": [], "L": []}
     if not heavy or not light:
         return ("<p style='padding:1em'>Enter both a heavy and a light chain.</p>",
-                "Need both chains.", "", None, "", {}, False, reset_btn)
+                "Need both chains.", "", None, "", {}, False, reset_btn,
+                [], [], {"H": [], "L": []}, empty_manual)
     try:
         t0 = time.time()
         antibody = PREDICTOR.predict({"H": heavy, "L": light})
@@ -410,7 +438,8 @@ def fold(heavy: str, light: str, chain_choice: str = "Both chains"):
         antibody.save("myabs_fold.pdb")
     except Exception as e:  # OpenMM refinement / numbering can occasionally fail
         return (f"<p style='padding:1em;color:#c00'>Fold failed: {e}</p>",
-                f"Error: {e}", "", None, "", {}, False, reset_btn)
+                f"Error: {e}", "", None, "", {}, False, reset_btn,
+                [], [], {"H": [], "L": []}, empty_manual)
 
     pdb = open("myabs_fold.pdb").read()
     flags, highlights, h3_len, ok = developability(heavy, light, pdb)
@@ -419,7 +448,13 @@ def fold(heavy: str, light: str, chain_choice: str = "Both chains"):
     status = (f"Folded in {dt:.1f} s  ·  VH {len(heavy)} aa / VL {len(light)} aa  ·  "
               f"CDRs highlighted (H: yellow/orange/red, L: cyan/blue).")
     flags_md = format_flags(flags, h3_len, ok)
-    return viewer, status, flags_md, "myabs_fold.pdb", pdb, highlights, False, reset_btn
+
+    # Clickable sequence tracks (index->imgt), and reset manual highlights.
+    h_tokens, h_idx = build_track("H", number_chain(heavy) or [])
+    l_tokens, l_idx = build_track("L", number_chain(light) or [])
+    idxmap = {"H": h_idx, "L": l_idx}
+    return (viewer, status, flags_md, "myabs_fold.pdb", pdb, highlights, False, reset_btn,
+            h_tokens, l_tokens, idxmap, empty_manual)
 
 
 def load_library(name: str):
@@ -435,22 +470,58 @@ def load_library(name: str):
     return entry["H"], entry["L"], prov
 
 
-def toggle_spin(spin: bool, pdb: str, highlights, chain_choice: str = "Both chains"):
+def toggle_spin(spin: bool, pdb: str, highlights, chain_choice="Both chains", manual=None):
     """Flip spin on/off and re-render the stored structure (no re-fold)."""
     spin = not spin
     label = "⏸ Stop" if spin else "▶ Spin"
     if not pdb:
         return spin, gr.update(value=label), gr.update()
     chains = CHAIN_VIEW.get(chain_choice, ("H", "L"))
-    return spin, gr.update(value=label), render_structure(pdb, spin, highlights, chains)
+    return spin, gr.update(value=label), render_structure(pdb, spin, highlights, chains, manual)
 
 
-def select_chains(chain_choice: str, pdb: str, highlights, spin: bool):
+def select_chains(chain_choice: str, pdb: str, highlights, spin: bool, manual=None):
     """Re-render the stored structure showing only the chosen chain(s), no re-fold."""
     if not pdb:
         return gr.update()
     chains = CHAIN_VIEW.get(chain_choice, ("H", "L"))
-    return render_structure(pdb, spin, highlights, chains)
+    return render_structure(pdb, spin, highlights, chains, manual)
+
+
+def _rerender(pdb, highlights, spin, chain_choice, manual):
+    chains = CHAIN_VIEW.get(chain_choice, ("H", "L"))
+    return render_structure(pdb, spin, highlights, chains, manual)
+
+
+def click_residue(chain_label, evt, idxmap, manual, pdb, highlights, spin, chain_choice):
+    """Toggle the clicked residue's highlight and re-render (no re-fold)."""
+    manual = {"H": list(manual.get("H", [])), "L": list(manual.get("L", []))}
+    idx = getattr(evt, "index", None)
+    if isinstance(idx, (list, tuple)):
+        idx = idx[0] if idx else None
+    imap = idxmap.get(chain_label, []) if idxmap else []
+    if not pdb or idx is None or not isinstance(idx, int) or idx < 0 or idx >= len(imap):
+        return manual, gr.update()
+    imgt = imap[idx]
+    cur = manual[chain_label]
+    cur.remove(imgt) if imgt in cur else cur.append(imgt)
+    return manual, _rerender(pdb, highlights, spin, chain_choice, manual)
+
+
+def click_heavy(idxmap, manual, pdb, highlights, spin, chain_choice, evt: gr.SelectData):
+    return click_residue("H", evt, idxmap, manual, pdb, highlights, spin, chain_choice)
+
+
+def click_light(idxmap, manual, pdb, highlights, spin, chain_choice, evt: gr.SelectData):
+    return click_residue("L", evt, idxmap, manual, pdb, highlights, spin, chain_choice)
+
+
+def clear_highlights(pdb, highlights, spin, chain_choice):
+    """Drop all user-clicked residue highlights and re-render."""
+    manual = {"H": [], "L": []}
+    if not pdb:
+        return manual, gr.update()
+    return manual, _rerender(pdb, highlights, spin, chain_choice, manual)
 
 
 with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
@@ -486,6 +557,15 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
                 "pinch = zoom · two-finger drag = pan.  _Mouse:_ drag = rotate · "
                 "scroll = zoom · right-drag = pan."
             )
+            gr.Markdown(
+                "**Click a residue below to highlight it on the structure (lime stick).** "
+                "Click it again to remove it. CDR residues are pre-colored."
+            )
+            h_track = gr.HighlightedText(label="Heavy chain (VH)", combine_adjacent=False,
+                                         show_legend=False, color_map=TRACK_COLORS)
+            l_track = gr.HighlightedText(label="Light chain (VL)", combine_adjacent=False,
+                                         show_legend=False, color_map=TRACK_COLORS)
+            clear_btn = gr.Button("Clear clicked highlights", size="sm")
 
     gr.Markdown("---")
     flags_md = gr.Markdown()
@@ -493,8 +573,11 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
     pdb_state = gr.State("")       # last folded PDB, for re-render without re-folding
     hi_state = gr.State({})        # liability highlight positions by chain
     spin_state = gr.State(False)   # is the structure currently spinning?
+    idxmap_state = gr.State({"H": [], "L": []})   # track index -> imgt, per chain
+    manual_state = gr.State({"H": [], "L": []})   # user-clicked residues, per chain
 
-    fold_outputs = [viewer, status, flags_md, pdb_file, pdb_state, hi_state, spin_state, spin_btn]
+    fold_outputs = [viewer, status, flags_md, pdb_file, pdb_state, hi_state, spin_state,
+                    spin_btn, h_track, l_track, idxmap_state, manual_state]
 
     go.click(fold, inputs=[h_in, l_in, chain_dd], outputs=fold_outputs)
 
@@ -504,16 +587,25 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
     )
 
     spin_btn.click(
-        toggle_spin, inputs=[spin_state, pdb_state, hi_state, chain_dd],
+        toggle_spin, inputs=[spin_state, pdb_state, hi_state, chain_dd, manual_state],
         outputs=[spin_state, spin_btn, viewer],
     )
 
     # Switch which chain(s) are shown, re-rendering the stored fold (no re-fold).
     chain_dd.change(
-        select_chains, inputs=[chain_dd, pdb_state, hi_state, spin_state],
+        select_chains, inputs=[chain_dd, pdb_state, hi_state, spin_state, manual_state],
         outputs=viewer,
     )
 
+    # Click a residue in a sequence track -> toggle its lime highlight on the structure.
+    click_inputs = [idxmap_state, manual_state, pdb_state, hi_state, spin_state, chain_dd]
+    h_track.select(click_heavy, inputs=click_inputs, outputs=[manual_state, viewer])
+    l_track.select(click_light, inputs=click_inputs, outputs=[manual_state, viewer])
+    clear_btn.click(
+        clear_highlights, inputs=[pdb_state, hi_state, spin_state, chain_dd],
+        outputs=[manual_state, viewer],
+    )
+
 if __name__ == "__main__":
     # Local dev defaults to 127.0.0.1:7900. On an HF Docker Space the Dockerfile
     # sets GRADIO_SERVER_NAME=0.0.0.0 and GRADIO_SERVER_PORT=7860 (app_port).