|
import gradio as gr |
|
import os |
|
|
|
import copy |
|
import os |
|
import torch |
|
|
|
import time |
|
from argparse import ArgumentParser, Namespace, FileType |
|
from rdkit.Chem import RemoveHs |
|
from functools import partial |
|
import numpy as np |
|
import pandas as pd |
|
from rdkit import RDLogger |
|
from rdkit.Chem import MolFromSmiles, AddHs |
|
from torch_geometric.loader import DataLoader |
|
import yaml |
|
import sys |
|
import csv |
|
|
|
csv.field_size_limit(sys.maxsize) |
|
|
|
print(torch.__version__) |
|
os.makedirs("data/esm2_output", exist_ok=True) |
|
os.makedirs("results", exist_ok=True) |
|
from datasets.process_mols import ( |
|
read_molecule, |
|
generate_conformer, |
|
write_mol_with_coords, |
|
) |
|
from datasets.pdbbind import PDBBind |
|
from utils.diffusion_utils import t_to_sigma as t_to_sigma_compl, get_t_schedule |
|
from utils.sampling import randomize_position, sampling |
|
from utils.utils import get_model |
|
from utils.visualise import PDBFile |
|
from tqdm import tqdm |
|
from datasets.esm_embedding_preparation import esm_embedding_prep |
|
import subprocess |
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
with open(f"workdir/paper_score_model/model_parameters.yml") as f: |
|
score_model_args = Namespace(**yaml.full_load(f)) |
|
|
|
with open(f"workdir/paper_confidence_model/model_parameters.yml") as f: |
|
confidence_args = Namespace(**yaml.full_load(f)) |
|
|
|
import shutil |
|
|
|
t_to_sigma = partial(t_to_sigma_compl, args=score_model_args) |
|
|
|
model = get_model(score_model_args, device, t_to_sigma=t_to_sigma, no_parallel=True) |
|
state_dict = torch.load( |
|
f"workdir/paper_score_model/best_ema_inference_epoch_model.pt", |
|
map_location=torch.device("cpu"), |
|
) |
|
model.load_state_dict(state_dict, strict=True) |
|
model = model.to(device) |
|
model.eval() |
|
|
|
confidence_model = get_model( |
|
confidence_args, |
|
device, |
|
t_to_sigma=t_to_sigma, |
|
no_parallel=True, |
|
confidence_mode=True, |
|
) |
|
state_dict = torch.load( |
|
f"workdir/paper_confidence_model/best_model_epoch75.pt", |
|
map_location=torch.device("cpu"), |
|
) |
|
confidence_model.load_state_dict(state_dict, strict=True) |
|
confidence_model = confidence_model.to(device) |
|
confidence_model.eval() |
|
|
|
|
|
def get_pdb(pdb_code="", filepath=""): |
|
try: |
|
return filepath.name |
|
except AttributeError as e: |
|
if pdb_code is None or pdb_code == "": |
|
return None |
|
else: |
|
os.system(f"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb") |
|
return f"{pdb_code}.pdb" |
|
|
|
|
|
def get_ligand(smiles="", filepath=""): |
|
if smiles is None or smiles == "": |
|
try: |
|
return filepath.name |
|
except AttributeError as e: |
|
return None |
|
else: |
|
return smiles |
|
|
|
|
|
def read_mol(molpath): |
|
with open(molpath, "r") as fp: |
|
lines = fp.readlines() |
|
mol = "" |
|
for l in lines: |
|
mol += l |
|
return mol |
|
|
|
|
|
def molecule(input_pdb, ligand_pdb, original_ligand): |
|
|
|
structure = read_mol(input_pdb) |
|
mol = read_mol(ligand_pdb) |
|
|
|
try: |
|
ligand = read_mol(original_ligand.name) |
|
_, ext = os.path.splitext(original_ligand.name) |
|
lig_str_1 = """let original_ligand = `""" + ligand + """`""" |
|
lig_str_2 = f""" |
|
viewer.addModel( original_ligand, "{ext[1:]}" ); |
|
viewer.getModel(2).setStyle({{stick:{{colorscheme:"greenCarbon"}}}});""" |
|
except AttributeError as e: |
|
ligand = None |
|
lig_str_1 = "" |
|
lig_str_2 = "" |
|
|
|
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: 600px; |
|
height: 600px; |
|
position: relative; |
|
mx-auto:0 |
|
} |
|
.mol-container select{ |
|
background-image:None; |
|
} |
|
.green{ |
|
width:20px; |
|
height:20px; |
|
background-color:#33ff45; |
|
display:inline-block; |
|
} |
|
.magenta{ |
|
width:20px; |
|
height:20px; |
|
background-color:magenta; |
|
display:inline-block; |
|
} |
|
</style> |
|
<script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script> |
|
</head> |
|
<body> |
|
<button id="startanimation">Replay diffusion process</button> |
|
<button id="togglesurface">Toggle surface representation</button> |
|
<div> |
|
<span class="green"></span> Uploaded ligand position |
|
<span class="magenta"></span> Predicted ligand position |
|
</div> |
|
<div id="container" class="mol-container"></div> |
|
|
|
<script> |
|
let ligand = `""" |
|
+ mol |
|
+ """` |
|
let structure = `""" |
|
+ structure |
|
+ """` |
|
""" |
|
+ lig_str_1 |
|
+ """ |
|
|
|
let viewer = null; |
|
let surface = false; |
|
let surf = null; |
|
$(document).ready(function () { |
|
let element = $("#container"); |
|
let config = { backgroundColor: "white" }; |
|
viewer = $3Dmol.createViewer(element, config); |
|
viewer.addModel( structure, "pdb" ); |
|
viewer.setStyle({}, {cartoon: {color: "gray"}}); |
|
viewer.zoomTo(); |
|
viewer.zoom(0.7); |
|
viewer.addModelsAsFrames(ligand, "pdb"); |
|
viewer.animate({loop: "forward",reps: 1}); |
|
|
|
viewer.getModel(1).setStyle({stick:{colorscheme:"magentaCarbon"}}); |
|
""" |
|
+ lig_str_2 |
|
+ """ |
|
viewer.render(); |
|
|
|
}) |
|
|
|
$("#startanimation").click(function() { |
|
viewer.animate({loop: "forward",reps: 1}); |
|
}); |
|
$("#togglesurface").click(function() { |
|
if (surface != true) { |
|
surf = viewer.addSurface($3Dmol.SurfaceType.VDW, { "opacity": 0.9, "color": "white" }, { model: 0 }); |
|
surface = true; |
|
} else { |
|
viewer.removeAllSurfaces() |
|
surface = false; |
|
} |
|
}); |
|
</script> |
|
</body></html>""" |
|
) |
|
|
|
return f"""<iframe style="width: 100%; height: 700px" 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>""" |
|
|
|
|
|
import sys |
|
|
|
|
|
def esm(protein_path, out_file): |
|
print("running esm") |
|
esm_embedding_prep(out_file, protein_path) |
|
|
|
os.environ["HOME"] = "esm/model_weights" |
|
subprocess.call( |
|
f"python esm/scripts/extract.py esm2_t33_650M_UR50D {out_file} data/esm2_output --repr_layers 33 --include per_tok", |
|
shell=True, |
|
env=os.environ, |
|
) |
|
|
|
|
|
def update(inp, file, ligand_inp, ligand_file, n_it): |
|
pdb_path = get_pdb(inp, file) |
|
ligand_path = get_ligand(ligand_inp, ligand_file) |
|
|
|
esm( |
|
pdb_path, |
|
f"data/{os.path.basename(pdb_path)}_prepared_for_esm.fasta", |
|
) |
|
tr_schedule = get_t_schedule(inference_steps=n_it) |
|
rot_schedule = tr_schedule |
|
tor_schedule = tr_schedule |
|
print("common t schedule", tr_schedule) |
|
( |
|
failures, |
|
skipped, |
|
confidences_list, |
|
names_list, |
|
run_times, |
|
min_self_distances_list, |
|
) = ( |
|
0, |
|
0, |
|
[], |
|
[], |
|
[], |
|
[], |
|
) |
|
N = 10 |
|
protein_path_list = [pdb_path] |
|
ligand_descriptions = [ligand_path] |
|
no_random = False |
|
ode = False |
|
no_final_step_noise = False |
|
out_dir = "results/" |
|
test_dataset = PDBBind( |
|
transform=None, |
|
root="", |
|
protein_path_list=protein_path_list, |
|
ligand_descriptions=ligand_descriptions, |
|
receptor_radius=score_model_args.receptor_radius, |
|
cache_path="data/cache", |
|
remove_hs=score_model_args.remove_hs, |
|
max_lig_size=None, |
|
c_alpha_max_neighbors=score_model_args.c_alpha_max_neighbors, |
|
matching=False, |
|
keep_original=False, |
|
popsize=score_model_args.matching_popsize, |
|
maxiter=score_model_args.matching_maxiter, |
|
all_atoms=score_model_args.all_atoms, |
|
atom_radius=score_model_args.atom_radius, |
|
atom_max_neighbors=score_model_args.atom_max_neighbors, |
|
esm_embeddings_path="data/esm2_output", |
|
require_ligand=True, |
|
num_workers=1, |
|
keep_local_structures=False, |
|
) |
|
test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False) |
|
confidence_test_dataset = PDBBind( |
|
transform=None, |
|
root="", |
|
protein_path_list=protein_path_list, |
|
ligand_descriptions=ligand_descriptions, |
|
receptor_radius=confidence_args.receptor_radius, |
|
cache_path="data/cache", |
|
remove_hs=confidence_args.remove_hs, |
|
max_lig_size=None, |
|
c_alpha_max_neighbors=confidence_args.c_alpha_max_neighbors, |
|
matching=False, |
|
keep_original=False, |
|
popsize=confidence_args.matching_popsize, |
|
maxiter=confidence_args.matching_maxiter, |
|
all_atoms=confidence_args.all_atoms, |
|
atom_radius=confidence_args.atom_radius, |
|
atom_max_neighbors=confidence_args.atom_max_neighbors, |
|
esm_embeddings_path="data/esm2_output", |
|
require_ligand=True, |
|
num_workers=1, |
|
) |
|
confidence_complex_dict = {d.name: d for d in confidence_test_dataset} |
|
for idx, orig_complex_graph in tqdm(enumerate(test_loader)): |
|
if ( |
|
confidence_model is not None |
|
and not ( |
|
confidence_args.use_original_model_cache |
|
or confidence_args.transfer_weights |
|
) |
|
and orig_complex_graph.name[0] not in confidence_complex_dict.keys() |
|
): |
|
skipped += 1 |
|
print( |
|
f"HAPPENING | The confidence dataset did not contain {orig_complex_graph.name[0]}. We are skipping this complex." |
|
) |
|
continue |
|
try: |
|
data_list = [copy.deepcopy(orig_complex_graph) for _ in range(N)] |
|
randomize_position( |
|
data_list, |
|
score_model_args.no_torsion, |
|
no_random, |
|
score_model_args.tr_sigma_max, |
|
) |
|
pdb = None |
|
lig = orig_complex_graph.mol[0] |
|
visualization_list = [] |
|
for graph in data_list: |
|
pdb = PDBFile(lig) |
|
pdb.add(lig, 0, 0) |
|
pdb.add( |
|
( |
|
orig_complex_graph["ligand"].pos |
|
+ orig_complex_graph.original_center |
|
) |
|
.detach() |
|
.cpu(), |
|
1, |
|
0, |
|
) |
|
pdb.add( |
|
(graph["ligand"].pos + graph.original_center).detach().cpu(), |
|
part=1, |
|
order=1, |
|
) |
|
visualization_list.append(pdb) |
|
|
|
start_time = time.time() |
|
if confidence_model is not None and not ( |
|
confidence_args.use_original_model_cache |
|
or confidence_args.transfer_weights |
|
): |
|
confidence_data_list = [ |
|
copy.deepcopy(confidence_complex_dict[orig_complex_graph.name[0]]) |
|
for _ in range(N) |
|
] |
|
else: |
|
confidence_data_list = None |
|
|
|
data_list, confidence = sampling( |
|
data_list=data_list, |
|
model=model, |
|
inference_steps=n_it, |
|
tr_schedule=tr_schedule, |
|
rot_schedule=rot_schedule, |
|
tor_schedule=tor_schedule, |
|
device=device, |
|
t_to_sigma=t_to_sigma, |
|
model_args=score_model_args, |
|
no_random=no_random, |
|
ode=ode, |
|
visualization_list=visualization_list, |
|
confidence_model=confidence_model, |
|
confidence_data_list=confidence_data_list, |
|
confidence_model_args=confidence_args, |
|
batch_size=1, |
|
no_final_step_noise=no_final_step_noise, |
|
) |
|
ligand_pos = np.asarray( |
|
[ |
|
complex_graph["ligand"].pos.cpu().numpy() |
|
+ orig_complex_graph.original_center.cpu().numpy() |
|
for complex_graph in data_list |
|
] |
|
) |
|
run_times.append(time.time() - start_time) |
|
|
|
if confidence is not None and isinstance( |
|
confidence_args.rmsd_classification_cutoff, list |
|
): |
|
confidence = confidence[:, 0] |
|
if confidence is not None: |
|
confidence = confidence.cpu().numpy() |
|
re_order = np.argsort(confidence)[::-1] |
|
confidence = confidence[re_order] |
|
confidences_list.append(confidence) |
|
ligand_pos = ligand_pos[re_order] |
|
write_dir = ( |
|
f'{out_dir}/index{idx}_{data_list[0]["name"][0].replace("/","-")}' |
|
) |
|
os.makedirs(write_dir, exist_ok=True) |
|
confidences = [] |
|
for rank, pos in enumerate(ligand_pos): |
|
mol_pred = copy.deepcopy(lig) |
|
if score_model_args.remove_hs: |
|
mol_pred = RemoveHs(mol_pred) |
|
if rank == 0: |
|
write_mol_with_coords( |
|
mol_pred, pos, os.path.join(write_dir, f"rank{rank+1}.sdf") |
|
) |
|
confidences.append(confidence[rank]) |
|
write_mol_with_coords( |
|
mol_pred, |
|
pos, |
|
os.path.join( |
|
write_dir, f"rank{rank+1}_confidence{confidence[rank]:.2f}.sdf" |
|
), |
|
) |
|
self_distances = np.linalg.norm( |
|
ligand_pos[:, :, None, :] - ligand_pos[:, None, :, :], axis=-1 |
|
) |
|
self_distances = np.where( |
|
np.eye(self_distances.shape[2]), np.inf, self_distances |
|
) |
|
min_self_distances_list.append(np.min(self_distances, axis=(1, 2))) |
|
|
|
filenames = [] |
|
if confidence is not None: |
|
for rank, batch_idx in enumerate(re_order): |
|
visualization_list[batch_idx].write( |
|
os.path.join(write_dir, f"rank{rank+1}_reverseprocess.pdb") |
|
) |
|
filenames.append( |
|
os.path.join(write_dir, f"rank{rank+1}_reverseprocess.pdb") |
|
) |
|
else: |
|
for rank, batch_idx in enumerate(ligand_pos): |
|
visualization_list[batch_idx].write( |
|
os.path.join(write_dir, f"rank{rank+1}_reverseprocess.pdb") |
|
) |
|
filenames.append( |
|
os.path.join(write_dir, f"rank{rank+1}_reverseprocess.pdb") |
|
) |
|
names_list.append(orig_complex_graph.name[0]) |
|
except Exception as e: |
|
print("Failed on", orig_complex_graph["name"], e) |
|
failures += 1 |
|
return None |
|
|
|
zippath = shutil.make_archive( |
|
os.path.join("results", os.path.basename(pdb_path)), "zip", write_dir |
|
) |
|
print("Zipped outputs to", zippath) |
|
labels = [ |
|
f"rank {i+1}, confidence {confidences[i]:.2f}" for i in range(len(filenames)) |
|
] |
|
|
|
torch.cuda.empty_cache() |
|
return ( |
|
molecule(pdb_path, filenames[0], ligand_file), |
|
gr.Dropdown.update(choices=labels, value=labels[0]), |
|
filenames, |
|
pdb_path, |
|
zippath, |
|
) |
|
|
|
|
|
def updateView(out, filenames, pdb, ligand_file): |
|
print("updating view") |
|
i = out |
|
print(i) |
|
i = int(i.split(",")[0].replace("rank", "")) - 1 |
|
return molecule(pdb, filenames[i], ligand_file) |
|
|
|
|
|
demo = gr.Blocks() |
|
|
|
with demo: |
|
gr.Markdown("# DiffDock") |
|
gr.Markdown( |
|
">**DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking**, Corso, Gabriele and Stärk, Hannes and Jing, Bowen and Barzilay, Regina and Jaakkola, Tommi, arXiv:2210.01776 [GitHub](https://github.com/gcorso/diffdock)" |
|
) |
|
gr.Markdown("") |
|
with gr.Box(): |
|
with gr.Row(): |
|
with gr.Column(): |
|
gr.Markdown("## Protein") |
|
inp = gr.Textbox( |
|
placeholder="PDB Code or upload file below", label="Input structure" |
|
) |
|
file = gr.File(file_count="single", label="Input PDB") |
|
with gr.Column(): |
|
gr.Markdown("## Ligand") |
|
ligand_inp = gr.Textbox( |
|
placeholder="Provide SMILES input or upload mol2/sdf file below", |
|
label="SMILES string", |
|
) |
|
ligand_file = gr.File(file_count="single", label="Input Ligand") |
|
n_it = gr.Slider( |
|
minimum=10, maximum=40, label="Number of inference steps", step=1 |
|
) |
|
|
|
btn = gr.Button("Run predictions") |
|
|
|
gr.Markdown("## Output") |
|
pdb = gr.Variable() |
|
filenames = gr.Variable() |
|
out = gr.Dropdown(interactive=True, label="Ranked samples") |
|
mol = gr.HTML() |
|
output_file = gr.File(file_count="single", label="Output files") |
|
gr.Examples( |
|
[ |
|
[ |
|
"6w70", |
|
"examples/6w70.pdb", |
|
"COc1ccc(cc1)n2c3c(c(n2)C(=O)N)CCN(C3=O)c4ccc(cc4)N5CCCCC5=O", |
|
"examples/6w70_ligand.sdf", |
|
10, |
|
], |
|
[ |
|
"6moa", |
|
"examples/6moa_protein_processed.pdb", |
|
"", |
|
"examples/6moa_ligand.sdf", |
|
10, |
|
], |
|
[ |
|
"", |
|
"examples/6o5u_protein_processed.pdb", |
|
"", |
|
"examples/6o5u_ligand.sdf", |
|
10, |
|
], |
|
[ |
|
"", |
|
"examples/6o5u_protein_processed.pdb", |
|
"[NH3+]C[C@H]1O[C@H](O[C@@H]2[C@@H]([NH3+])C[C@H]([C@@H]([C@H]2O)O[C@H]2O[C@H](CO)[C@H]([C@@H]([C@H]2O)[NH3+])O)[NH3+])[C@@H]([C@H]([C@@H]1O)O)O", |
|
"examples/6o5u_ligand.sdf", |
|
10, |
|
], |
|
[ |
|
"", |
|
"examples/6o5u_protein_processed.pdb", |
|
"", |
|
"examples/6o5u_ligand.sdf", |
|
10, |
|
], |
|
[ |
|
"", |
|
"examples/6ahs_protein_processed.pdb", |
|
"", |
|
"examples/6ahs_ligand.sdf", |
|
10, |
|
], |
|
], |
|
[inp, file, ligand_inp, ligand_file, n_it], |
|
[mol, out, filenames, pdb, output_file], |
|
|
|
|
|
) |
|
btn.click( |
|
fn=update, |
|
inputs=[inp, file, ligand_inp, ligand_file, n_it], |
|
outputs=[mol, out, filenames, pdb, output_file], |
|
) |
|
out.change(fn=updateView, inputs=[out, filenames, pdb, ligand_file], outputs=mol) |
|
demo.launch() |
|
|
