Persistent client-side 3Dmol viewer: focus/grey/hide + in-place highlights (orientation preserved)

app.py

@@ -17,10 +17,11 @@ CANDIDATES, not patent-ready antibodies — wet-lab validation still required.
 """
 
 import io
+import itertools
+import json
 import os
 import time
 import gradio as gr
-import py3Dmol
 from ImmuneBuilder import ABodyBuilder2
 
 try:
@@ -358,13 +359,6 @@ def format_flags(flags, h3_len, ok):
 
 
 # -------------------------------------------------------------------- viewer
-# Which chains the viewer draws, for the "Show chains" selector.
-CHAIN_VIEW = {
-    "Both chains": ("H", "L"),
-    "Heavy only (VH)": ("H",),
-    "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"],
@@ -384,77 +378,145 @@ def build_track(chain_label: str, residues):
     return tokens, idxmap
 
 
-def render_structure(pdb: str, spin: bool = False, highlights=None,
-                     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,
-    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
-    for chain in chains:
-        view.addStyle({"chain": chain}, {"cartoon": {"color": "#cfd8dc"}})
-        for cdr, (lo, hi) in CDR_RANGES.items():
-            view.addStyle(
-                {"chain": chain, "resi": list(range(lo, hi + 1))},
-                {"cartoon": {"color": CDR_COLORS[chain][cdr]}},
-            )
-        if highlights:
-            pos = highlights.get(chain, [])
-            if pos:
-                view.addStyle(
-                    {"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)
-    html = view._make_html()
-    esc = html.replace("&", "&").replace('"', """)
-    return f'<iframe style="width:100%;height:580px;border:none;border-radius:8px" srcdoc="{esc}"></iframe>'
+FOCUS_CHOICES = ["Both chains", "Heavy only (VH)", "Light only (VL)"]
+MODE_CHOICES = ["Grey out others", "Hide others"]
+
+_NONCE = itertools.count(1)
+
+
+def build_payload(pdb: str, highlights) -> str:
+    """JSON handed to the client-side viewer: the structure + liability positions.
+    The nonce guarantees the value changes each fold so the .change bridge fires."""
+    highlights = highlights or {}
+    return json.dumps({
+        "pdb": pdb,
+        "highlights": {"H": highlights.get("H", []), "L": highlights.get("L", [])},
+        "n": next(_NONCE),
+    })
+
+
+def toggle_payload(chain_label: str, evt, idxmap) -> str:
+    """JSON telling the client to toggle a lime stick on one clicked residue."""
+    idx = getattr(evt, "index", None)
+    if isinstance(idx, (list, tuple)):
+        idx = idx[0] if idx else None
+    imap = (idxmap or {}).get(chain_label, [])
+    imgt = imap[idx] if isinstance(idx, int) and 0 <= idx < len(imap) else None
+    return json.dumps({"chain": chain_label, "imgt": imgt, "n": next(_NONCE)})
+
+
+# Client-side 3Dmol controller. One persistent viewer; every interaction updates
+# it IN PLACE (no re-render, no zoomTo) so the camera / orientation is preserved.
+# Loaded once via demo.load(js=...); it also injects 3Dmol.js from the CDN.
+_CDR_JS = json.dumps(CDR_COLORS)
+_RANGES_JS = json.dumps({k: list(v) for k, v in CDR_RANGES.items()})
+CONTROLLER_JS = """
+() => {
+  if (window.__myabsReady) return;
+  window.__myabsReady = true;
+  const CDR = __CDR__;
+  const RANGES = __RANGES__;
+  const rlist = (lo,hi) => { let a=[]; for(let i=lo;i<=hi;i++) a.push(i); return a; };
+  const S = window.myabsState = {viewer:null, highlights:{H:[],L:[]}, picks:{H:[],L:[]},
+                                 focus:["H","L"], mode:"grey", spinning:false};
+
+  window.myabsApply = () => {
+    const v = S.viewer; if(!v) return;
+    v.setStyle({}, {});
+    ["H","L"].forEach(ch => {
+      if (S.focus.includes(ch)) {
+        v.addStyle({chain:ch}, {cartoon:{color:"#cfd8dc"}});
+        for (const c in RANGES){ const r=RANGES[c];
+          v.addStyle({chain:ch, resi:rlist(r[0],r[1])}, {cartoon:{color:CDR[ch][c]}}); }
+        const hl=S.highlights[ch]||[]; if(hl.length) v.addStyle({chain:ch, resi:hl},{stick:{color:"magenta",radius:0.3}});
+        const pk=S.picks[ch]||[]; if(pk.length) v.addStyle({chain:ch, resi:pk},{stick:{color:"lime",radius:0.3}});
+      } else if (S.mode === "grey") {
+        v.addStyle({chain:ch}, {cartoon:{color:"#c7ccd1", opacity:0.35}});
+      }
+    });
+    v.render();
+  };
+
+  window.myabsInit = () => {
+    if (S.viewer) return true;
+    const el = document.getElementById("myabs-viewer");
+    if (!el || !window.$3Dmol) return false;
+    S.viewer = $3Dmol.createViewer(el, {backgroundColor:"white"});
+    return true;
+  };
+
+  window.myabsLoad = (payload) => {
+    if (!payload) return;
+    let p; try { p = JSON.parse(payload); } catch(e){ return; }
+    if (!p.pdb) return;
+    if (!window.myabsInit()) { setTimeout(()=>window.myabsLoad(payload), 200); return; }
+    const v = S.viewer;
+    S.highlights = p.highlights || {H:[],L:[]}; S.picks = {H:[],L:[]}; S.spinning = false;
+    v.removeAllModels(); v.addModel(p.pdb, "pdb");
+    window.myabsApply();
+    v.zoomTo(); v.render();   // new molecule: recentering here is expected
+  };
+
+  window.myabsSetFocus = (choice) => {
+    const M = {"Both chains":["H","L"], "Heavy only (VH)":["H"], "Light only (VL)":["L"]};
+    S.focus = M[choice] || ["H","L"]; window.myabsApply();   // no zoomTo -> view kept
+  };
+  window.myabsSetMode = (m) => {
+    S.mode = (m && m.indexOf("Hide")>=0) ? "hide" : "grey"; window.myabsApply();
+  };
+  window.myabsToggleResidue = (payload) => {
+    let p; try { p = JSON.parse(payload); } catch(e){ return; }
+    if (p.imgt == null || !p.chain) return;
+    const arr = S.picks[p.chain] || (S.picks[p.chain]=[]);
+    const i = arr.indexOf(p.imgt); if (i>=0) arr.splice(i,1); else arr.push(p.imgt);
+    window.myabsApply();
+  };
+  window.myabsClear = () => { S.picks = {H:[],L:[]}; window.myabsApply(); };
+  window.myabsToggleSpin = () => {
+    const v = S.viewer; if (!v) return "▶ Spin";
+    S.spinning = !S.spinning; v.spin(S.spinning ? "y" : false);
+    return S.spinning ? "⏸ Stop" : "▶ Spin";
+  };
+
+  if (!window.$3Dmol) {
+    const s = document.createElement("script");
+    s.src = "https://3Dmol.org/build/3Dmol-min.js";
+    s.onload = () => window.myabsInit();
+    document.head.appendChild(s);
+  } else { window.myabsInit(); }
+}
+""".replace("__CDR__", _CDR_JS).replace("__RANGES__", _RANGES_JS)
 
 
 # ---------------------------------------------------------------------- fold
-def fold(heavy: str, light: str, chain_choice: str = "Both chains"):
-    """Fold VH+VL, scan liabilities, render, and reset the spin toggle."""
+def fold(heavy: str, light: str):
+    """Fold VH+VL, scan liabilities, and hand the structure to the client viewer."""
     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,
-                [], [], {"H": [], "L": []}, empty_manual)
+        return ("Enter both a heavy and a light chain.", "", None, "",
+                [], [], {"H": [], "L": []}, reset_btn)
     try:
         t0 = time.time()
         antibody = PREDICTOR.predict({"H": heavy, "L": light})
         dt = time.time() - t0
         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,
-                [], [], {"H": [], "L": []}, empty_manual)
+        return (f"Fold failed: {e}", "", None, "",
+                [], [], {"H": [], "L": []}, reset_btn)
 
     pdb = open("myabs_fold.pdb").read()
     flags, highlights, h3_len, ok = developability(heavy, light, pdb)
-    chains = CHAIN_VIEW.get(chain_choice, ("H", "L"))
-    viewer = render_structure(pdb, spin=False, highlights=highlights, chains=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)
 
-    # 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)
+    payload = build_payload(pdb, highlights)  # -> client-side viewer via .change bridge
+    return (status, flags_md, "myabs_fold.pdb", payload,
+            h_tokens, l_tokens, idxmap, reset_btn)
 
 
 def load_library(name: str):
@@ -470,58 +532,13 @@ def load_library(name: str):
     return entry["H"], entry["L"], prov
 
 
-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, manual)
+def click_heavy(idxmap, evt: gr.SelectData):
+    """Map a heavy-chain track click to its IMGT position for the client to toggle."""
+    return toggle_payload("H", evt, idxmap)
 
 
-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, 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)
+def click_light(idxmap, evt: gr.SelectData):
+    return toggle_payload("L", evt, idxmap)
 
 
 with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
@@ -546,11 +563,17 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
             status = gr.Markdown()
             pdb_file = gr.File(label="Download structure (.pdb)")
         with gr.Column(scale=3):
-            chain_dd = gr.Radio(
-                choices=list(CHAIN_VIEW.keys()), value="Both chains",
-                label="Show chains", scale=1,
+            with gr.Row():
+                focus_dd = gr.Radio(choices=FOCUS_CHOICES, value="Both chains",
+                                    label="Focus chain")
+                mode_dd = gr.Radio(choices=MODE_CHOICES, value="Grey out others",
+                                   label="The non-focused chain is…")
+            # Persistent 3Dmol viewer container. Updated in place by the client-side
+            # controller (CONTROLLER_JS) so interactions never reset the camera.
+            viewer = gr.HTML(
+                '<div id="myabs-viewer" style="width:100%;height:580px;position:relative;'
+                'border:1px solid #eee;border-radius:8px"></div>'
             )
-            viewer = gr.HTML()
             gr.HTML(LEGEND_HTML)
             gr.Markdown(
                 "**Rotate it yourself** — _Touchscreen:_ one-finger drag = rotate · "
@@ -570,41 +593,35 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
     gr.Markdown("---")
     flags_md = gr.Markdown()
 
-    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
+    idxmap_state = gr.State({"H": [], "L": []})  # track index -> imgt, per chain
+    # Hidden bridges: server writes JSON here, a .change(js=...) hands it to the viewer.
+    load_box = gr.Textbox(visible=False)     # fold -> myabsLoad
+    toggle_box = gr.Textbox(visible=False)   # residue click -> myabsToggleResidue
 
-    fold_outputs = [viewer, status, flags_md, pdb_file, pdb_state, hi_state, spin_state,
-                    spin_btn, h_track, l_track, idxmap_state, manual_state]
+    fold_outputs = [status, flags_md, pdb_file, load_box, h_track, l_track,
+                    idxmap_state, spin_btn]
 
-    go.click(fold, inputs=[h_in, l_in, chain_dd], outputs=fold_outputs)
+    go.click(fold, inputs=[h_in, l_in], outputs=fold_outputs)
 
     # Pick a therapeutic -> load its sequences + provenance -> auto-fold.
     lib_dd.change(load_library, inputs=lib_dd, outputs=[h_in, l_in, provenance]).then(
-        fold, inputs=[h_in, l_in, chain_dd], outputs=fold_outputs
+        fold, inputs=[h_in, l_in], outputs=fold_outputs
     )
 
-    spin_btn.click(
-        toggle_spin, inputs=[spin_state, pdb_state, hi_state, chain_dd, manual_state],
-        outputs=[spin_state, spin_btn, viewer],
-    )
+    # --- client-side viewer controls (no server round-trip, camera preserved) ---
+    load_box.change(None, inputs=[load_box], js="(p) => window.myabsLoad(p)")
+    toggle_box.change(None, inputs=[toggle_box], js="(p) => window.myabsToggleResidue(p)")
+    focus_dd.change(None, inputs=[focus_dd], js="(c) => window.myabsSetFocus(c)")
+    mode_dd.change(None, inputs=[mode_dd], js="(m) => window.myabsSetMode(m)")
+    spin_btn.click(None, outputs=[spin_btn], js="() => window.myabsToggleSpin()")
+    clear_btn.click(None, js="() => window.myabsClear()")
 
-    # 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, manual_state],
-        outputs=viewer,
-    )
+    # Residue click: server maps track index -> IMGT, client applies the lime stick.
+    h_track.select(click_heavy, inputs=[idxmap_state], outputs=[toggle_box])
+    l_track.select(click_light, inputs=[idxmap_state], outputs=[toggle_box])
 
-    # 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],
-    )
+    # Load 3Dmol.js + define the controller once, on page load.
+    demo.load(None, js=CONTROLLER_JS)
 
 if __name__ == "__main__":
     # Local dev defaults to 127.0.0.1:7900. On an HF Docker Space the Dockerfile