from rdkit.Chem import MolFromSmiles, MolToSmiles
from rdkit.Chem import Draw
import image
import copy
import time
import sys
sys.path.insert(0, '../../../')
import molecule_vae
grammar_weights = '../../../pretrained/zinc_vae_str_L56_E100_val.hdf5'
grammar_model = molecule_vae.ZincCharacterModel(grammar_weights)
# We pick the next 50 inputs
next_inputs = sgp.batched_greedy_ei(50, np.min(X_train, 0),
np.max(X_train, 0))
valid_smiles_final = decode_from_latent_space(next_inputs, grammar_model)
from rdkit.Chem import Descriptors
from rdkit.Chem import MolFromSmiles, MolToSmiles
new_features = next_inputs
save_object(valid_smiles_final,
"results/valid_smiles{}.dat".format(iteration))
logP_values = np.loadtxt(
'../../latent_features_and_targets_character/logP_values.txt')
SA_scores = np.loadtxt(
'../../latent_features_and_targets_character/SA_scores.txt')
# 4 September - saving the BNN is optional
# We save the trained BNN
#sys.setrecursionlimit(4000)
# don't have to save the BNN because it won't be necessary to plot the predictions in 56-D space.
# save_object(bb_alpha, "results_QED_solo/bb_alpha{}.dat".format(iteration))
# We pick the next 50 inputs (4 September - variable name bb_alpha_samples not descriptive, it is the size of the batch of collected data points)
bb_alpha_samples = 50
next_inputs = sgp.batched_greedy_ei(bb_alpha, 50, np.min(X_train, 0),
np.max(X_train, 0), bb_alpha_samples)
# We load the decoder to obtain the molecules
preproc = lasp.PreProcessing(dataset='drugs')
enc_dec = lasp.EncoderDecoder()
encoder, decoder = enc_dec.get_functions()
postprocessor = lasp.PostProcessing(enc_dec)
# We collect the molecule statistics
# 4 September - need descriptions of these variables, highly unclear
# 4 September - variables are lists of length = decode_attempts
for ga_iter in ga_pbar:
pred = model.predictor(grid)
ga_pbar.set_description(
'Max pred: {:.4f}, mean pred: {:.4f}, std pred: {:.4f}'
.format(pred.max(), pred.mean(), pred.std()))
grads = torch.autograd.grad([x for x in pred], grid)[0]
grid = grid + ga_lr * grads
pred = pred.detach().cpu().numpy()
selected_idxs = np.argsort(-pred[:, 0])[:batch_size]
next_inputs = grid[selected_idxs]
next_inputs = next_inputs.detach().cpu().numpy()
else:
next_inputs = sgp.batched_greedy_ei(batch_size,
np.min(X_train, 0),
np.max(X_train, 0),
np.mean(X_train, 0),
np.std(X_train, 0),
sample=sample_dist)
valid_arcs_final = util.decode_from_latent_space(
torch.FloatTensor(next_inputs).to(device), model, None, 500, max_n,
False, data_type)
if random_baseline:
if args.sample_dist == 'uniform':
random_inputs = np.random.rand(batch_size, gdim) * (
X_train.max(0) - X_train.min(0)) + X_train.min(0)
elif args.sample_dist == 'normal':
random_inputs = np.random.randn(
batch_size, gdim) * X_train.std(0) + X_train.mean(0)
random_inputs = torch.FloatTensor(random_inputs).to(device)
valid_arcs_random = util.decode_from_latent_space(
pred, uncert = sgp.predict(X_test, 0 * X_test)
error = np.sqrt(np.mean((pred - y_test)**2))
testll = np.mean(sps.norm.logpdf(pred - y_test, scale=np.sqrt(uncert)))
print('Test RMSE: ', error)
print('Test ll: ', testll)
pred, uncert = sgp.predict(X_train, 0 * X_train)
error = np.sqrt(np.mean((pred - y_train)**2))
trainll = np.mean(sps.norm.logpdf(pred - y_train, scale=np.sqrt(uncert)))
print('Train RMSE: ', error)
print('Train ll: ', trainll)
# We pick the next 50 inputs
next_inputs = sgp.batched_greedy_ei(args.bo_batch_size, np.min(X_train, 0),
np.max(X_train, 0))
# We decode the 50 smiles:
# Decode z into smiles
with torch.no_grad():
gen_seq = model.decode(torch.FloatTensor(next_inputs).to(device))
smiles = model.probas_to_smiles(gen_seq)
valid_smiles_final = []
for s in smiles:
s = decoder(s)
m = Chem.MolFromSmiles(s)
if m is None:
valid_smiles_final.append(None)
else:
Chem.Kekulize(m)
s = Chem.MolToSmiles(m, kekuleSmiles=True)
def run_bo_demo():
import sys
sys.path.append('/home/icml18-jtnn')
import pickle
import gzip
from sparse_gp import SparseGP
import scipy.stats as sps
import numpy as np
import os.path
import time
import rdkit
from rdkit.Chem import MolFromSmiles, MolToSmiles
from rdkit.Chem import Descriptors
from rdkit.Chem import PandasTools
import torch
import torch.nn as nn
from jtnn import create_var, JTNNVAE, Vocab
start_time = time.time()
lg = rdkit.RDLogger.logger()
lg.setLevel(rdkit.RDLogger.CRITICAL)
# We define the functions used to load and save objects
def save_object(obj, filename):
result = pickle.dumps(obj)
with gzip.GzipFile(filename, 'wb') as dest:
dest.write(result)
dest.close()
def load_object(filename):
with gzip.GzipFile(filename, 'rb') as source:
result = source.read()
ret = pickle.loads(result)
source.close()
return ret
vocab_path = '../data/vocab.txt'
#save_dir=save_dir
vocab = [x.strip("\r\n ") for x in open(vocab_path)]
vocab = Vocab(vocab)
# print(opts.save_dir)
hidden_size = 450
latent_size = 56
depth = 3
random_seed = 1
model = JTNNVAE(vocab, hidden_size, latent_size, depth)
model.load_state_dict(
torch.load('../molvae/MPNVAE-h450-L56-d3-beta0.005/model.iter-4',
map_location=lambda storage, loc: storage))
#model = model.cuda()
# We load the random seed
np.random.seed(random_seed)
# We load the data (y is minued!)
X = np.loadtxt('latent_features_demo.txt')
y = -np.loadtxt('targets_demo.txt')
y = y.reshape((-1, 1))
n = X.shape[0]
# print(X.shape[1])
permutation = np.random.choice(n, n, replace=False)
# print(n)
X_train = X[permutation, :][0:np.int(np.round(0.8 * n)), :]
X_test = X[permutation, :][np.int(np.round(0.8 * n)):, :]
# print(X_train.shape)
y_train = y[permutation][0:np.int(np.round(0.8 * n))]
y_test = y[permutation][np.int(np.round(0.8 * n)):]
np.random.seed(random_seed)
logP_values = np.loadtxt('logP_values_demo.txt')
SA_scores = np.loadtxt('SA_scores_demo.txt')
cycle_scores = np.loadtxt('cycle_scores_demo.txt')
SA_scores_normalized = (np.array(SA_scores) -
np.mean(SA_scores)) / np.std(SA_scores)
logP_values_normalized = (np.array(logP_values) -
np.mean(logP_values)) / np.std(logP_values)
cycle_scores_normalized = (np.array(cycle_scores) -
np.mean(cycle_scores)) / np.std(cycle_scores)
iteration = 0
while iteration < 1:
# We fit the GP
np.random.seed(iteration * random_seed)
M = 1
sgp = SparseGP(X_train, 0 * X_train, y_train, M)
sgp.train_via_ADAM(X_train,
0 * X_train,
y_train,
X_test,
X_test * 0,
y_test,
minibatch_size=2,
max_iterations=100,
learning_rate=0.001)
pred, uncert = sgp.predict(X_test, 0 * X_test)
error = np.sqrt(np.mean((pred - y_test)**2))
testll = np.mean(sps.norm.logpdf(pred - y_test, scale=np.sqrt(uncert)))
# print 'Test RMSE: ', error
# print 'Test ll: ', testll
pred, uncert = sgp.predict(X_train, 0 * X_train)
error = np.sqrt(np.mean((pred - y_train)**2))
trainll = np.mean(
sps.norm.logpdf(pred - y_train, scale=np.sqrt(uncert)))
# print 'Train RMSE: ', error
# print 'Train ll: ', trainll
# We pick the next 60 inputs
next_inputs = sgp.batched_greedy_ei(60, np.min(X_train, 0),
np.max(X_train, 0))
valid_smiles = []
valid_mols = []
new_features = []
for i in xrange(60):
all_vec = next_inputs[i].reshape((1, -1))
tree_vec, mol_vec = np.hsplit(all_vec, 2)
tree_vec = create_var(torch.from_numpy(tree_vec).float())
mol_vec = create_var(torch.from_numpy(mol_vec).float())
s = model.decode(tree_vec, mol_vec, prob_decode=False)
if s is not None:
valid_smiles.append(s)
# print(MolFromSmiles(s))
valid_mols.append(str(MolFromSmiles(s)))
new_features.append(all_vec)
print len(valid_smiles), "molecules are found"
valid_smiles = valid_smiles[:50]
valid_mols = valid_mols[:50]
new_features = next_inputs[:50]
new_features = np.vstack(new_features)
# save_object(valid_smiles, save_dir + "/valid_smiles{}.dat".format(iteration))
# save_object(valid_mols,save_dir + '/valid_mols{}.png'.format(iteration))
# save_object(mol1
import sascorer
import networkx as nx
from rdkit.Chem import rdmolops
scores = []
for i in range(len(valid_smiles)):
current_log_P_value = Descriptors.MolLogP(
MolFromSmiles(valid_smiles[i]))
current_SA_score = -sascorer.calculateScore(
MolFromSmiles(valid_smiles[i]))
cycle_list = nx.cycle_basis(
nx.Graph(
rdmolops.GetAdjacencyMatrix(MolFromSmiles(
valid_smiles[i]))))
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
current_cycle_score = -cycle_length
current_SA_score_normalized = (
current_SA_score - np.mean(SA_scores)) / np.std(SA_scores)
current_log_P_value_normalized = (
current_log_P_value -
np.mean(logP_values)) / np.std(logP_values)
current_cycle_score_normalized = (
current_cycle_score -
np.mean(cycle_scores)) / np.std(cycle_scores)
score = current_SA_score_normalized + current_log_P_value_normalized + current_cycle_score_normalized
scores.append(-score) #target is always minused
# print valid_smiles
# print scores
# save_object(scores, save_dir + "/scores{}.dat".format(iteration))
if len(new_features) > 0:
X_train = np.concatenate([X_train, new_features], 0)
y_train = np.concatenate([y_train, np.array(scores)[:, None]], 0)
iteration += 1
# print('Seconds taken: %s' %(time.time() -start_time))
all_smiles = []
#all_smiles=valid_smiles+scores+valid_mols
all_smiles.extend(zip(valid_smiles, scores, valid_mols))
all_smiles = [(x, -y, z) for x, y, z in all_smiles]
all_smiles = sorted(all_smiles, key=lambda x: x[1], reverse=True)
# return valid_smiles[0:5],scores[0:5],valid_mols[0:5]
return all_smiles[0:3]
sgp.train_via_ADAM(X_train, 0 * X_train, y_train, X_test, X_test * 0, \
y_test, minibatch_size = 10 * M, max_iterations = cmd_args.num_epochs, learning_rate = args.gp_lr)
# pred, uncert = sgp.predict(X_test, 0 * X_test)
# error = np.sqrt(np.mean((pred - y_test)**2))
# testll = np.mean(sps.norm.logpdf(pred - y_test, scale = np.sqrt(uncert)))
# print 'Test RMSE: ', error
# print 'Test ll: ', testll
# pred, uncert = sgp.predict(X_train, 0 * X_train)
# error = np.sqrt(np.mean((pred - y_train)**2))
# trainll = np.mean(sps.norm.logpdf(pred - y_train, scale = np.sqrt(uncert)))
# print 'Train RMSE: ', error
# print 'Train ll: ', trainll
next_inputs = sgp.batched_greedy_ei(50, np.min(X_train, 0), np.max(X_train, 0))
valid_smiles_final = decode_from_latent_space(next_inputs, model)
save_object(valid_smiles_final, "%s/valid_smiles-seed-%d-iter-%d.dat" % (cmd_args.save_dir, args.seed, iteration))
new_features = next_inputs
scores = []
for i in range(len(valid_smiles_final)):
if valid_smiles_final[ i ] is not None:
current_log_P_value = Descriptors.MolLogP(MolFromSmiles(valid_smiles_final[ i ]))
current_SA_score = -sascorer.calculateScore(MolFromSmiles(valid_smiles_final[ i ]))
cycle_list = nx.cycle_basis(nx.Graph(rdmolops.GetAdjacencyMatrix(MolFromSmiles(valid_smiles_final[ i ]))))
if len(cycle_list) == 0:
cycle_length = 0
else:
cycle_length = max([ len(j) for j in cycle_list ])
if cycle_length <= 6: