Add fold confidence: color-by-predicted-error mode + readout (completes CHECK-VIABLE part 1)

app.py

@@ -429,16 +429,31 @@ CONTROLLER_JS = """
   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};
+                                 focus:["H","L"], mode:"grey", spinning:false, colorMode:"cdr"};
+
+  // Per-residue predicted error (PDB B-factor, Å) -> color. 0 = confident (blue),
+  // >= CONF_MAX = uncertain (red), through yellow. Matches AlphaFold-style intuition.
+  const CONF_MAX = 1.5;
+  const confColor = (atom) => {
+    const t = Math.max(0, Math.min(1, (atom.b || 0) / CONF_MAX));
+    let r,g,b;
+    if (t < 0.5){ const u=t/0.5; r=Math.round(43+(240-43)*u); g=Math.round(131+(200-131)*u); b=Math.round(186+(50-186)*u); }
+    else { const u=(t-0.5)/0.5; r=Math.round(240+(215-240)*u); g=Math.round(200+(48-200)*u); b=Math.round(50+(39-50)*u); }
+    return "rgb("+r+","+g+","+b+")";
+  };
 
   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]}}); }
+        if (S.colorMode === "confidence") {
+          v.addStyle({chain:ch}, {cartoon:{colorfunc: confColor}});   // color by predicted error
+        } else {
+          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") {
@@ -475,6 +490,9 @@ CONTROLLER_JS = """
   window.myabsSetMode = (m) => {
     S.mode = (m && m.indexOf("Hide")>=0) ? "hide" : "grey"; window.myabsApply();
   };
+  window.myabsSetColorMode = (m) => {
+    S.colorMode = (m && m.indexOf("onfidence")>=0) ? "confidence" : "cdr"; window.myabsApply();
+  };
   window.myabsToggleResidue = (payload) => {
     let p; try { p = JSON.parse(payload); } catch(e){ return; }
     if (p.imgt == null || !p.chain) return;
@@ -504,6 +522,48 @@ CONTROLLER_JS = """
 
 
 # ---------------------------------------------------------------------- fold
+def fold_confidence(pdb: str) -> str:
+    """Readout of ABodyBuilder2's per-residue predicted error (PDB B-factor = RMS
+    spread across the 4-model ensemble, Å; lower = models agree = more confident)."""
+    if not HAVE_SASA or not pdb:
+        return ""
+    try:
+        model = PDBParser(QUIET=True).get_structure("ab", io.StringIO(pdb))[0]
+    except Exception:
+        return ""
+    per, h3 = [], []
+    for chain in model:
+        for res in chain:
+            bs = [a.get_bfactor() for a in res]
+            if not bs:
+                continue
+            b = sum(bs) / len(bs)
+            per.append(b)
+            if chain.id == "H" and 105 <= res.id[1] <= 117:
+                h3.append(b)
+    if not per:
+        return ""
+
+    def band(x):
+        return "high" if x < 0.5 else ("moderate" if x < 1.5 else "low")
+
+    mean = sum(per) / len(per)
+    out = [f"**Fold confidence:** mean predicted error **{mean:.2f} Å** "
+           f"({band(mean)} confidence)."]
+    if h3:
+        h3m = sum(h3) / len(h3)
+        note = ""
+        if h3m >= 1.5:
+            note = " — the least certain region, so treat its liability flags with extra caution"
+        elif h3m >= 0.8:
+            note = " — moderate certainty"
+        out.append(f"CDR-H3 **{h3m:.2f} Å**{note}.")
+    out.append("_Predicted error = spread across the 4-model ensemble (lower = models agree). "
+               "Set the viewer's \"Color by\" to **Fold confidence** to see it on the structure "
+               "(blue = confident → red = uncertain)._")
+    return " ".join(out)
+
+
 MAX_CHAIN_LEN = 250  # antibody variable domains are ~110-130 aa; a generous DoS cap
 
 
@@ -511,8 +571,8 @@ 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")
-    # trailing 7 outputs (everything after `status`) for the early-return paths
-    tail = ("", None, "", [], [], {"H": [], "L": []}, reset_btn)
+    # trailing 8 outputs (everything after `status`) for the early-return paths
+    tail = ("", "", None, "", [], [], {"H": [], "L": []}, reset_btn)
     if not heavy or not light:
         return ("Enter both a heavy and a light chain.",) + tail
     if len(heavy) > MAX_CHAIN_LEN or len(light) > MAX_CHAIN_LEN:
@@ -534,13 +594,14 @@ def fold(heavy: str, light: str):
     flags, highlights, h3_len, ok = developability(heavy, light, pdb)
     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).")
+    conf_md = fold_confidence(pdb)
     flags_md = format_flags(flags, h3_len, ok)
 
     h_tokens, h_idx = build_track("H", number_chain(heavy) or [], highlights.get("H", []))
     l_tokens, l_idx = build_track("L", number_chain(light) or [], highlights.get("L", []))
     idxmap = {"H": h_idx, "L": l_idx}
     payload = build_payload(pdb, highlights)  # -> client-side viewer via .change bridge
-    return (status, flags_md, out_path, payload,
+    return (status, conf_md, flags_md, out_path, payload,
             h_tokens, l_tokens, idxmap, reset_btn)
 
 
@@ -586,6 +647,7 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
                 go = gr.Button("Fold", variant="primary")
                 spin_btn = gr.Button("▶ Spin")
             status = gr.Markdown()
+            conf_md = gr.Markdown()
             pdb_file = gr.File(label="Download structure (.pdb)")
         with gr.Column(scale=3):
             with gr.Row():
@@ -593,6 +655,8 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
                                     label="Focus chain")
                 mode_dd = gr.Radio(choices=MODE_CHOICES, value="Grey out others",
                                    label="The non-focused chain is…")
+                color_dd = gr.Radio(choices=["CDR regions", "Fold confidence"],
+                                    value="CDR regions", label="Color by")
             # Persistent 3Dmol viewer container. Updated in place by the client-side
             # controller (CONTROLLER_JS) so interactions never reset the camera.
             viewer = gr.HTML(
@@ -600,6 +664,14 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
                 'border:1px solid #eee;border-radius:8px"></div>'
             )
             gr.HTML(LEGEND_HTML)
+            gr.HTML(
+                "<div style='font-size:13px;margin:2px 0'><b>Confidence scale</b> "
+                "(when Color by → Fold confidence): "
+                "<span style='color:rgb(43,131,186)'>■&nbsp;confident</span> → "
+                "<span style='color:rgb(240,200,50)'>■</span> → "
+                "<span style='color:rgb(215,48,39)'>■&nbsp;uncertain</span> "
+                "&nbsp;(predicted error, Å)</div>"
+            )
             gr.Markdown(
                 "**Rotate it yourself** — _Touchscreen:_ one-finger drag = rotate · "
                 "pinch = zoom · two-finger drag = pan.  _Mouse:_ drag = rotate · "
@@ -624,7 +696,7 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
     load_box = gr.Textbox(visible=False)     # fold -> myabsLoad
     toggle_box = gr.Textbox(visible=False)   # residue click -> myabsToggleResidue
 
-    fold_outputs = [status, flags_md, pdb_file, load_box, h_track, l_track,
+    fold_outputs = [status, conf_md, flags_md, pdb_file, load_box, h_track, l_track,
                     idxmap_state, spin_btn]
 
     go.click(fold, inputs=[h_in, l_in], outputs=fold_outputs)
@@ -639,6 +711,7 @@ with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
     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)")
+    color_dd.change(None, inputs=[color_dd], js="(m) => window.myabsSetColorMode(m)")
     spin_btn.click(None, outputs=[spin_btn], js="() => window.myabsToggleSpin()")
     clear_btn.click(None, js="() => window.myabsClear()")