-
Notifications
You must be signed in to change notification settings - Fork 34
Expand file tree
/
Copy pathRetrosynthesis.py
More file actions
121 lines (99 loc) · 5.98 KB
/
Copy pathRetrosynthesis.py
File metadata and controls
121 lines (99 loc) · 5.98 KB
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
113
114
115
116
117
118
119
120
import json
import pandas as pd
import torch
from torch import nn
import sklearn
from rdkit.Chem import PandasTools, AllChem
import dgl
from dgllife.utils import smiles_to_bigraph, WeaveAtomFeaturizer, CanonicalBondFeaturizer
from functools import partial
from scripts.utils import init_featurizer, load_model, collate_molgraphs_test
from scripts.get_edit import combined_edit
from LocalTemplate.template_decoder import *
atom_types = ['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na', 'Ca', 'Fe',
'As', 'Al', 'I', 'B', 'V', 'K', 'Tl', 'Yb', 'Sb', 'Sn', 'Ag', 'Pd', 'Co', 'Se', 'Ti',
'Zn', 'H', 'Li', 'Ge', 'Cu', 'Au', 'Ni', 'Cd', 'In', 'Mn', 'Zr', 'Cr', 'Pt', 'Hg', 'Pb',
'W', 'Ru', 'Nb', 'Re', 'Te', 'Rh', 'Ta', 'Tc', 'Ba', 'Bi', 'Hf', 'Mo', 'U', 'Sm', 'Os', 'Ir',
'Ce', 'Gd', 'Ga', 'Cs']
def predict(model, graph, device):
bg = dgl.batch([graph])
bg.set_n_initializer(dgl.init.zero_initializer)
bg.set_e_initializer(dgl.init.zero_initializer)
bg = bg.to(device)
node_feats = bg.ndata.pop('h').to(device)
edge_feats = bg.edata.pop('e').to(device)
return model(bg, node_feats, edge_feats)
def load_templates(args):
atom_templates = pd.read_csv('%s/atom_templates.csv' % args['data_dir'])
bond_templates = pd.read_csv('%s/bond_templates.csv' % args['data_dir'])
template_infos = pd.read_csv('%s/template_infos.csv' % args['data_dir'])
atom_templates = {atom_templates['Class'][i]: atom_templates['Template'][i] for i in atom_templates.index}
bond_templates = {bond_templates['Class'][i]: bond_templates['Template'][i] for i in bond_templates.index}
template_infos = {template_infos['Template'][i]: {'edit_site': eval(template_infos['edit_site'][i]), 'change_H': eval(template_infos['change_H'][i]), 'change_C': eval(template_infos['change_C'][i]), 'change_S': eval(template_infos['change_S'][i])} for i in template_infos.index}
return atom_templates, bond_templates, template_infos
def remap(mol):
for atom in mol.GetAtoms():
atom.SetAtomMapNum(atom.GetIdx())
return mol
class LocalRetro():
def __init__(self, args):
args['mode'] = 'test'
args = init_featurizer(args)
smiles_to_graph = partial(smiles_to_bigraph, add_self_loop=True)
node_featurizer, edge_featurizer = WeaveAtomFeaturizer(atom_types = atom_types), CanonicalBondFeaturizer(self_loop=True)
self.graph_function = lambda s: smiles_to_graph(s, node_featurizer = node_featurizer, edge_featurizer = edge_featurizer, canonical_atom_order = False)
self.atom_templates, self.bond_templates, self.template_infos = load_templates(args)
self.device = torch.device('cpu')
if torch.cuda.is_available():
self.device = args['device']
self.model = load_model(args)
self.model.eval()
def retrosnythesis(self, smiles, top_k = 10, verbose = False):
graph = self.graph_function(smiles)
with torch.no_grad():
atom_logits, bond_logits, _ = predict(self.model, graph, self.device)
atom_logits = nn.Softmax(dim=1)(atom_logits)
bond_logits = nn.Softmax(dim=1)(bond_logits)
graph = graph.remove_self_loop()
pred_types, pred_sites, pred_scores = combined_edit(graph, atom_logits, bond_logits, top_k)
predictions = [(pred_types[k], pred_sites[k][0], pred_sites[k][1], pred_scores[k]) for k in range(top_k)]
predicted_reactants = [smiles]
predicted_sites = [None]
predicted_templates = [None]
predicted_scores = [None]
for k, prediction in enumerate(predictions):
mol, pred_site, template, template_info, score = read_prediction(smiles, prediction, self.atom_templates, self.bond_templates, self.template_infos, True)
local_template = '>>'.join(['(%s)' % smarts for smarts in template.split('_')[0].split('>>')])
decoded_smiles = decode_localtemplate(mol, pred_site, local_template, template_info)
if verbose:
print ('top %s' % (k+1), pred_site, local_template, score)
try:
decoded_smiles = decode_localtemplate(mol, pred_site, local_template, template_info)
if decoded_smiles == None or decoded_smiles in predicted_reactants:
continue
except Exception as e:
print (e)
continue
predicted_sites.append(pred_site)
predicted_reactants.append(decoded_smiles)
predicted_templates.append(local_template)
predicted_scores.append(score)
results_df = pd.DataFrame({'SMILES': predicted_reactants, 'Predicted site': predicted_sites, 'Local reaction template': predicted_templates, 'Score': predicted_scores})
PandasTools.AddMoleculeColumnToFrame(results_df,'SMILES','Molecule')
remap(results_df['Molecule'][0])
return results_df
# def init_LocalRetro(args):
# args['mode'] = 'test'
# args = init_featurizer(args)
# model = load_model(args)
# atom_templates, bond_templates, template_infos = load_templates(args)
# smiles_to_graph = partial(smiles_to_bigraph, add_self_loop=True)
# atom_types = ['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na', 'Ca', 'Fe',
# 'As', 'Al', 'I', 'B', 'V', 'K', 'Tl', 'Yb', 'Sb', 'Sn', 'Ag', 'Pd', 'Co', 'Se', 'Ti',
# 'Zn', 'H', 'Li', 'Ge', 'Cu', 'Au', 'Ni', 'Cd', 'In', 'Mn', 'Zr', 'Cr', 'Pt', 'Hg', 'Pb',
# 'W', 'Ru', 'Nb', 'Re', 'Te', 'Rh', 'Ta', 'Tc', 'Ba', 'Bi', 'Hf', 'Mo', 'U', 'Sm', 'Os', 'Ir',
# 'Ce', 'Gd', 'Ga', 'Cs']
# node_featurizer = WeaveAtomFeaturizer(atom_types = atom_types)
# edge_featurizer = CanonicalBondFeaturizer(self_loop=True)
# graph_function = lambda s: smiles_to_graph(s, node_featurizer = node_featurizer, edge_featurizer = edge_featurizer, canonical_atom_order = False)
# return model, graph_function, atom_templates, bond_templates, template_infos