Spaces:
Running
Running
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, custom_resids, clustering_threshold): | |
try: | |
filepath = 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", | |
) | |
filepath = 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", | |
) | |
filepath = f"files/{pdb_file}.pdb" | |
else: | |
return "pdb code must be 4 letters or Uniprot code does not match", "" | |
identifier = os.path.basename(filepath) | |
if mode == "All residues": | |
print('using all residues') | |
ids = get_all_protein_resids(filepath) | |
elif len(custom_resids)!=0: | |
print('using listed residues', custom_resids) | |
ids=get_all_resids_from_list(filepath,custom_resids.replace(","," ")) | |
else: | |
print('using metalbinding') | |
ids = get_all_metalbinding_resids(filepath) | |
print(filepath) | |
print(ids) | |
try: | |
voxels, prot_centers, prot_N, prots = processStructures(filepath, ids) | |
except Exception as e: | |
return "Error", f"""<div class="text-center mt-4"> Something went wrong with the voxelization, reset custom residues and other input fiels and check error message <br> <br> <code>{e}</code></div>""" | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
voxels.to(device) | |
model = Model() | |
model.to(device) | |
model.load_state_dict(torch.load("weights/metal_0.5A_v3_d0.2_16Abox.pth", map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"))) | |
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_{identifier}.cube", | |
gridres=gridres, | |
) | |
message = find_unique_sites( | |
probability_values, | |
grid, | |
writeprobes=True, | |
probefile=f"output/probes_{identifier}.pdb", | |
threshold=7, | |
p=clustering_threshold, | |
) | |
return message, molecule( | |
filepath, | |
f"output/probes_{identifier}.pdb", | |
f"output/metal_{identifier}.cube", | |
) | |
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" /> | |
<link rel="stylesheet" href="https://unpkg.com/flowbite@1.4.5/dist/flowbite.min.css" /> | |
<style> | |
body{ | |
font-family:sans-serif | |
} | |
.mol-container { | |
width: 100%; | |
height: 600px; | |
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 my-8"> | |
<span>Isovalue </span> | |
<span id="isovalue">0.5</span> | |
<input class="slider text-blue-400" type="range" id="rangeslider" min="0" max="1" step="0.05" value=0.5> | |
</div> | |
<div id="container" class="mol-container"></div> | |
<div class="flex items-center justify-center my-4"> | |
<div class="px-4"> | |
<label for="sidechain" class="relative inline-flex items-center mb-4 cursor-pointer "> | |
<input id="sidechain"type="checkbox" class="sr-only peer"> | |
<div class="w-11 h-6 bg-gray-200 rounded-full peer peer-focus:ring-4 peer-focus:ring-blue-300 dark:peer-focus:ring-blue-800 dark:bg-gray-700 peer-checked:after:translate-x-full peer-checked:after:border-white after:absolute after:top-0.5 after:left-[2px] after:bg-white after:border-gray-300 after:border after:rounded-full after:h-5 after:w-5 after:transition-all dark:border-gray-600 peer-checked:bg-blue-600"></div> | |
<span class="ml-3 text-sm font-medium text-gray-900 dark:text-gray-300">Show side chains</span> | |
</label> | |
</div> | |
<div class="px-4"> | |
<label for="pdbmetal" class="relative inline-flex items-center mb-4 cursor-pointer "> | |
<input id="pdbmetal" type="checkbox" class="sr-only peer" checked> | |
<div class="w-11 h-6 bg-gray-200 rounded-full peer peer-focus:ring-4 peer-focus:ring-blue-300 dark:peer-focus:ring-blue-800 dark:bg-gray-700 peer-checked:after:translate-x-full peer-checked:after:border-white after:absolute after:top-0.5 after:left-[2px] after:bg-white after:border-gray-300 after:border after:rounded-full after:h-5 after:w-5 after:transition-all dark:border-gray-600 peer-checked:bg-blue-600"></div> | |
<span class="ml-3 text-sm font-medium text-gray-900 dark:text-gray-300">Show PDB metals</span> | |
</label> | |
</div> | |
<div class="px-4"> | |
<label for="probes" class="relative inline-flex items-center mb-4 cursor-pointer "> | |
<input id="probes" type="checkbox" class="sr-only peer" checked> | |
<div class="w-11 h-6 bg-gray-200 rounded-full peer peer-focus:ring-4 peer-focus:ring-blue-300 dark:peer-focus:ring-blue-800 dark:bg-gray-700 peer-checked:after:translate-x-full peer-checked:after:border-white after:absolute after:top-0.5 after:left-[2px] after:bg-white after:border-gray-300 after:border after:rounded-full after:h-5 after:w-5 after:transition-all dark:border-gray-600 peer-checked:bg-blue-600"></div> | |
<span class="ml-3 text-sm font-medium text-gray-900 dark:text-gray-300">Show Probes</span> | |
</label> | |
</div> | |
</div> | |
<div class="flex items-center justify-center my-4"> | |
<button type="button" class="text-gray-900 bg-white hover:bg-gray-100 border border-gray-200 focus:ring-4 focus:outline-none focus:ring-gray-100 font-medium rounded-lg text-sm px-5 py-2.5 text-center inline-flex items-center dark:focus:ring-gray-600 dark:bg-gray-800 dark:border-gray-700 dark:text-white dark:hover:bg-gray-700 mr-2 mb-2" id="download"> | |
<svg class="w-6 h-6 mr-2 -ml-1" fill="none" stroke="currentColor" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"></path></svg> | |
Download predictions | |
</button> | |
</div> | |
<script> | |
let viewer = null; | |
let voldata = null; | |
let shape = null; | |
let sidechain = null; | |
let metal = 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"); | |
shape = viewer.addIsosurface(voldata, { isoval: 0.5 , color: "blue", alpha: 0.85, smoothness: 1 }); | |
viewer.getModel(0).setStyle({}, {cartoon: {}}); | |
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); | |
$("#sidechain").change(function () { | |
if (this.checked) { | |
BB = ["C", "O", "N"] | |
viewer.getModel(0).setStyle( {"and": [{resn: ["GLY", "PRO"], invert: true},{atom: BB, invert: true},]},{stick: {hidden:false, colorscheme: "WhiteCarbon", radius: 0.3}, cartoon: {}}); | |
viewer.render() | |
$("#pdbmetal").prop( "checked", false ); | |
} else { | |
BB = ["C", "O", "N"] | |
viewer.getModel(0).setStyle({"and": [{resn: ["GLY", "PRO"], invert: true},{atom: BB, invert: true},]},{stick: {colorscheme: "WhiteCarbon",hidden:true, radius: 0.3}, cartoon: {}}); | |
viewer.render() | |
$("#pdbmetal").prop( "checked", false ); | |
} | |
}); | |
$("#pdbmetal").change(function () { | |
if (this.checked) { | |
viewer.getModel(0).setStyle({ "resn": ["ZN","MG","NA","FE", "NI","MN","CA", "CU", "CU1"] }, { "sphere": {hidden:false}}); | |
viewer.render() | |
} else { | |
viewer.getModel(0).setStyle({ "resn": ["ZN","MG","NA","FE","NI", "MN","CA", "CU", "CU1"] }, { "sphere": {hidden:true}}); | |
viewer.render() | |
} | |
}); | |
$("#probes").change(function () { | |
if (this.checked) { | |
viewer.getModel(1).setStyle({ "resn": "ZN" }, { "sphere": { }}); | |
viewer.addStyle() | |
viewer.render() | |
} else { | |
viewer.getModel(1).setStyle({}); | |
viewer.render() | |
} | |
}); | |
$("#download").click(function () { | |
download("protein.pdb", data); | |
download("metaldensity.cube", cubefile); | |
download("probes.pdb", probes); | |
}) | |
}); | |
function download(filename, text) { | |
var element = document.createElement("a"); | |
element.setAttribute("href", "data:text/plain;charset=utf-8," + encodeURIComponent(text)); | |
element.setAttribute("download", filename); | |
element.style.display = "none"; | |
document.body.appendChild(element); | |
element.click(); | |
document.body.removeChild(element); | |
} | |
</script> | |
<script> | |
$("#rangeslider").rangeslider().on("change", function (el) { | |
isoval = parseFloat(el.target.value); | |
$("#isovalue").text(el.target.value) | |
console.log("Change isosurface to "+el.target.value) | |
viewer.removeShape(shape) | |
shape=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: 1000px" 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 set_examples(example): | |
n,code, resids = example | |
return [ | |
n, | |
code, | |
resids | |
] | |
metal3d = gr.Blocks() | |
with metal3d: | |
gr.Markdown("# Metal3D") | |
with gr.Tabs(): | |
with gr.TabItem("Input"): | |
inp = gr.Textbox( placeholder="PDB Code or Uniprot identifier or upload file below", label="Input molecule" | |
) | |
file = gr.File(file_count="single", type="file") | |
with gr.TabItem("Settings"): | |
with gr.Row(): | |
mode = gr.Radio( | |
["All metalbinding residues (ASP, CYS, GLU, HIS)", "All residues"], | |
label="Residues to use for prediction", | |
) | |
custom_resids = gr.Textbox(placeholder="Comma separated list of residues", label="Custom residues") | |
with gr.Row(): | |
clustering_threshold = gr.Slider(minimum=0.15,maximum=1, value=0.15,step=0.05, label="Clustering threshold") | |
distance_cutoff = gr.Slider(minimum=1,maximum=10, value=7,step=1, label="Clustering distance cutoff") | |
btn = gr.Button("Run") | |
n = gr.Textbox(label="Label",visible=False) | |
examples = gr.Dataset(components=[n,inp, custom_resids], | |
samples=[["HCA2", "2CBA", ""], | |
["Nickel in GB1 dimer", "6F5N", ""], | |
["Zebrafish palmitoyltransferase ZDHHC15B PDB", "6BMS", ""], | |
["Human palmitoyltransferase ZDHHC23 AlphaFold", "Q8IYP9", "280,273,263,260,274,277,274,287"]], | |
) | |
examples.click(fn=set_examples, inputs=examples, outputs=examples.components) | |
#gr.Markdown( | |
# """ <small>Inference using CPU-only, can be quite slow for more than 20 residues. Use Colab notebook for GPU acceleration</small> | |
#""" | |
#) | |
gr.Markdown("# Output") | |
out = gr.Textbox(label="status") | |
mol = gr.HTML() | |
btn.click(fn=update, inputs=[inp, file, mode, custom_resids, clustering_threshold], outputs=[out, mol]) | |
metal3d.launch(share=True) | |