Update app.py (#2)

- Update app.py (54c4c689941bd2e25d2b4d73d34a5344cf9c83ff)


Co-authored-by: Simon Duerr <simonduerr@users.noreply.huggingface.co>

app.py

@@ -50,8 +50,36 @@ def molecule(input_pdb):
                 let element = $("#container");
                 let config = { backgroundColor: "white" };
                 let viewer = $3Dmol.createViewer(element, config);
+                let colorAlpha = function (atom) {
+                    if (atom.b < 50) {
+                        return "OrangeRed";
+                    } else if (atom.b < 70) {
+                        return "Gold";
+                    } else if (atom.b < 90) {
+                        return "MediumTurquoise";
+                    } else {
+                        return "Blue";
+                    }
+                };
+                
                 viewer.addModel(pdb, "pdb");
-                viewer.getModel(0).setStyle({}, { cartoon: { colorscheme: {"prop":"resi","gradient":"roygb", "min": 26, "max": 250} } });
+                // set plddt coloring
+                viewer.getModel(0).setStyle({cartoon: { colorfunc: colorAlpha }});
+                // display pLDDT tooltips when hovering over atoms
+                viewer.getModel(0).setHoverable({}, true,
+                        function (atom, viewer, event, container) {
+                            if (!atom.label) {
+                                atom.label = viewer.addLabel(atom.resn + atom.resi + " pLDDT=" + atom.b, { position: atom, backgroundColor: "mintcream", fontColor: "black" });
+                            }
+                        },
+                        function (atom, viewer) {
+                            if (atom.label) {
+                                viewer.removeLabel(atom.label);
+                                delete atom.label;
+                            }
+                        }
+                    );
+                }
                 viewer.zoomTo();
                 viewer.render();
                 viewer.zoom(1.2, 2000);