metal3d / app.py
simonduerr's picture
Update app.py
c4025fd
import gradio as gr
import urllib
import re
import sys
import warnings
import torch
import torch.nn as nn
import ipywidgets as widgets
from ipywidgets import interact, fixed
from utils.helpers import *
from utils.voxelization import processStructures
from utils.model import Model
import numpy as np
import os
import moleculekit
print(moleculekit.__version__)
def update(inp, file, mode):
try:
pdb_file = file.name
except:
print("using pdbfile")
try:
pdb_file = inp
if (
re.match(
"[OPQ][0-9][A-Z0-9]{3}[0-9]|[A-NR-Z][0-9]([A-Z][A-Z0-9]{2}[0-9]){1,2}",
pdb_file,
).group()
== pdb_file
):
urllib.request.urlretrieve(
f"https://alphafold.ebi.ac.uk/files/AF-{pdb_file}-F1-model_v2.pdb",
f"files/{pdb_file}.pdb",
)
except AttributeError:
if len(inp) == 4:
pdb_file = inp
urllib.request.urlretrieve(
f"http://files.rcsb.org/download/{pdb_file.lower()}.pdb1",
f"files/{pdb_file}.pdb",
)
else:
return "pdb code must be 4 letters or Uniprot code does not match", ""
if mode == "All residues":
ids = get_all_protein_resids(
f"files/{pdb_file}.pdb",
)
else:
ids = get_all_metalbinding_resids(f"files/{pdb_file}.pdb")
voxels, prot_centers, prot_N, prots = processStructures(pdb_file, ids)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
voxels.to(device)
print(voxels.shape)
model = Model()
model.to(device)
model.load_state_dict(torch.load("weights/metal_0.5A_v3_d0.2_16Abox.pth"))
model.eval()
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
output = model(voxels)
print(output.shape)
prot_v = np.vstack(prot_centers)
output_v = output.flatten().cpu().detach().numpy()
bb = get_bb(prot_v)
gridres = 0.5
grid, box_N = create_grid_fromBB(bb, voxelSize=gridres)
probability_values = get_probability_mean(grid, prot_v, output_v)
print(probability_values.shape)
write_cubefile(
bb,
probability_values,
box_N,
outname=f"output/metal_{pdb_file}.cube",
gridres=gridres,
)
message = find_unique_sites(
probability_values,
grid,
writeprobes=True,
probefile=f"output/probes_{pdb_file}.pdb",
threshold=7,
p=0.15,
)
return message, molecule(
f"files/{pdb_file}.pdb",
f"output/probes_{pdb_file}.pdb",
f"output/metal_{pdb_file}.cube",
)
def test():
x = """<!DOCTYPE html>
<html>
<head>
<meta http-equiv="content-type" content="text/html; charset=UTF-8" />
</head>
<body>
<script src="https://3Dmol.org/build/3Dmol-min.js" async></script> <div style="height: 400px; width: 400px; position: relative;" class="viewer_3Dmoljs" data-pdb="2POR" data-backgroundcolor="0xffffff" data-style="stick" ></div>
</body></html>"""
return f"""<iframe style="width: 100%; height: 480px" name="result" allow="midi; geolocation; microphone; camera;
display-capture; encrypted-media;" sandbox="allow-modals allow-forms
allow-scripts allow-same-origin allow-popups
allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
def read_mol(molpath):
with open(molpath, "r") as fp:
lines = fp.readlines()
mol = ""
for l in lines:
mol += l
return mol
def molecule(pdb, probes, cube):
mol = read_mol(pdb)
probes = read_mol(probes)
cubefile = read_mol(cube)
x = (
"""<!DOCTYPE html>
<html>
<head>
<meta http-equiv="content-type" content="text/html; charset=UTF-8" />
<style>
body{
font-family:sans-serif
}
.mol-container {
width: 100%;
height: 400px;
position: relative;
}
.slider{
width:80%;
margin:0 auto
}
.slidercontainer{
display:flex;
}
.slidercontainer > * + * {
margin-left: 0.5rem;
}
#isovalue{
text-align:right}
</style>
<script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/rangeslider.js/2.3.3/rangeslider.min.js" integrity="sha512-BUlWdwDeJo24GIubM+z40xcj/pjw7RuULBkxOTc+0L9BaGwZPwiwtbiSVzv31qR7TWx7bs6OPTE5IyfLOorboQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
</head>
<body>
<div class="slidercontainer">
<span>Isovalue </span>
<span id="isovalue">0.5</span>
<input class="slider" type="range" id="rangeslider" min="0" max="1" step="0.025" value=0.5>
</div>
<div id="container" class="mol-container"></div>
<script>
let viewer = null;
let voldata = null;
$(document).ready(function () {
let element = $("#container");
let config = { backgroundColor: "white" };
viewer = $3Dmol.createViewer( element, config );
viewer.ui.initiateUI();
let data = `"""
+ mol
+ """`
viewer.addModel( data, "pdb" );
let cubefile = `"""
+ cubefile
+ """`
voldata = new $3Dmol.VolumeData(cubefile, "cube");
viewer.addIsosurface(voldata, { isoval: 0.7 , color: "blue", alpha: 0.85, smoothness: 1 });
viewer.getModel(0).setStyle({}, {cartoon: {color: "grayCarbon"}});
let probes =`"""
+ probes
+ """`
viewer.addModel(probes, "pdb");
viewer.getModel(1).setStyle({ "resn": "ZN" }, { "sphere": { }});
viewer.getModel(1).setHoverable({}, true,
function (atom, viewer, event, container) {
if (!atom.label) {
atom.label = viewer.addLabel("ZN p=" + atom.pdbline.substring(55, 60), { 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(0.8, 2000);
});
</script>
<script>
$("#rangeslider").rangeslider().on("change", function (el) {
isoval = parseFloat(el.target.value);
$("#isovalue").text(el.target.value)
viewer.addIsosurface(voldata, { isoval: parseFloat(el.target.value), color: "blue", alpha: 0.85, smoothness: 1 });
viewer.render();
});
</script>
</body></html>"""
)
return f"""<iframe style="width: 100%; height: 480px" name="result" allow="midi; geolocation; microphone; camera;
display-capture; encrypted-media;" sandbox="allow-modals allow-forms
allow-scripts allow-same-origin allow-popups
allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
metal3d = gr.Blocks()
with metal3d:
gr.Markdown("# Metal3D")
gr.Markdown(
"""
Details about implementation and code available here:
>Duerr, Levy and Roethlisberger, Predicting zinc ion location using deep learning, BioRxiv, 2022 "
"""
)
with gr.Group():
inp = gr.Textbox(
placeholder="PDB Code or Uniprot identifier", label="Input molecule"
)
gr.Markdown("or upload a file")
file = gr.File(file_count="single", type="file")
mode = gr.Radio(
["All metalbinding residues (ASP, CYS, GLU, HIS)", "All residues"],
label="Residues to use for prediction",
)
btn = gr.Button("Run")
gr.Markdown("# Output")
out = gr.Textbox(label="status")
mol = gr.HTML()
btn.click(fn=update, inputs=[inp, file, mode], outputs=[out, mol])
metal3d.launch()