Spaces:
Sleeping
Sleeping
File size: 4,739 Bytes
a3ea5d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import torch
import streamlit as st
import sys, os
import rdkit
import rdkit.Chem as Chem
from rdkit.Chem.Draw import MolToImage
from rdkit.Chem import Descriptors
import sascorer
import networkx as nx
os.environ['KMP_DUPLICATE_LIB_OK']='True'
sys.path.append('%s/fast_jtnn/' % os.path.dirname(os.path.realpath(__file__)))
from mol_tree import Vocab, MolTree
from jtprop_vae import JTPropVAE
from molbloom import buy
lg = rdkit.RDLogger.logger()
lg.setLevel(rdkit.RDLogger.CRITICAL)
st.header('Junction Tree Variational Autoencoder for Molecular Graph Generation (JTVAE)')
st.subheader('Wengong Jin, Regina Barzilay, Tommi Jaakkola')
descrip = '''
We seek to automate the design of molecules based on specific chemical properties. In computational terms, this task involves continuous embedding and generation of molecular graphs. Our primary contribution is the direct realization of molecular graphs, a task previously approached by generating linear SMILES strings instead of graphs. Our junction tree variational autoencoder generates molecular graphs in two phases, by first generating a tree-structured scaffold over chemical substructures, and then combining them into a molecule with a graph message passing network. This approach allows us to incrementally expand molecules while maintaining chemical validity at every step. We evaluate our model on multiple tasks ranging from molecular generation to optimization. Across these tasks, our model outperforms previous state-of-the-art baselines by a significant margin.
[https://arxiv.org/abs/1802.04364](https://arxiv.org/abs/1802.04364)'''
with st.expander('About'):
st.markdown(descrip)
st.text_input('Enter a SMILES string:','CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',key='smiles')
def penalized_logp_standard(mol):
logP_mean = 2.4399606244103639873799239
logP_std = 0.9293197802518905481505840
SA_mean = -2.4485512208785431553792478
SA_std = 0.4603110476923852334429910
cycle_mean = -0.0307270378623088931402396
cycle_std = 0.2163675785228087178335699
log_p = Descriptors.MolLogP(mol)
SA = -sascorer.calculateScore(mol)
# cycle score
cycle_list = nx.cycle_basis(nx.Graph(Chem.rdmolops.GetAdjacencyMatrix(mol)))
if len(cycle_list) == 0:
cycle_length = 0
else:
cycle_length = max([len(j) for j in cycle_list])
if cycle_length <= 6:
cycle_length = 0
else:
cycle_length = cycle_length - 6
cycle_score = -cycle_length
# print(logP_mean)
standardized_log_p = (log_p - logP_mean) / logP_std
standardized_SA = (SA - SA_mean) / SA_std
standardized_cycle = (cycle_score - cycle_mean) / cycle_std
return standardized_log_p + standardized_SA + standardized_cycle
mol = Chem.MolFromSmiles(st.session_state.smiles)
if mol is None:
st.write('SMILES is invalid. Please enter a valid SMILES.')
else:
st.write('Molecule:')
st.image(MolToImage(mol,size=(300,300)))
score = penalized_logp_standard(mol)
st.write('Penalized logP score: %.5f' % (score))
if mol is not None:
st.slider('Choose learning rate: ',0.0,10.0,0.4,key='lr')
st.slider('Choose similarity cutoff: ',0.0,3.0,0.4,key='sim_cutoff')
st.slider('Choose number of iterations: ',1,100,80,key='n_iter')
vocab = [x.strip("\r\n ") for x in open('./vocab.txt')]
vocab = Vocab(vocab)
if st.button('Optimize'):
st.write('Testing')
model = JTPropVAE(vocab, 450, 56, 20, 3)
model.load_state_dict(torch.load('./model.iter-685000',map_location=torch.device('cpu')))
new_smiles,sim = model.optimize(st.session_state.smiles, sim_cutoff=st.session_state.sim_cutoff, lr=st.session_state.lr, num_iter=st.session_state.n_iter)
del model
if new_smiles is None:
st.write('Cannot optimize.')
else:
st.write('New SMILES:')
st.code(new_smiles)
new_mol = Chem.MolFromSmiles(new_smiles)
if new_mol is None:
st.write('New SMILES is invalid.')
else:
st.write('New SMILES molecule:')
st.image(MolToImage(new_mol,size=(300,300)))
new_score = penalized_logp_standard(new_mol)
st.write('New penalized logP score: %.5f' % (new_score))
st.write('Caching ZINC20 if necessary...')
if buy(new_smiles, catalog='zinc20',canonicalize=True):
st.write('This molecule exists.')
st.caption('Checked by molbloom.')
else:
st.write('THIS MOLECULE DOES NOT EXIST!')
st.caption('Checked by molbloom.')
|