Skip to content

Commit 8d93b94

Browse files
kexinhuang12345kexinhuang12345
committed
integrate into DTI with new drug ebcoders
1 parent 5238c74 commit 8d93b94

2 files changed

Lines changed: 81 additions & 40 deletions

File tree

‎DeepPurpose/DTI.py‎

Lines changed: 33 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -210,42 +210,14 @@ def virtual_screening(X_repurpose, target, model, drug_names = None, target_name
210210

211211
return y_pred
212212

213-
## x is a list, len(x)=batch_size, x[i] is tuple, len(x[0])=5
214-
# def mpnn_feature_collate_func(x):
215-
# ## first version
216-
# return [torch.cat([x[j][i] for j in range(len(x))], 0) for i in range(len(x[0]))]
217-
218-
# def mpnn_feature_collate_func(x):
219-
# assert len(x[0]) == 5
220-
# N_atoms_N_bonds = [i[-1] for i in x]
221-
# N_atoms_scope = []
222-
# f_a = torch.cat([x[j][0] for j in range(len(x))], 0)
223-
# f_b = torch.cat([x[j][1] for j in range(len(x))], 0)
224-
# agraph_lst, bgraph_lst = [], []
225-
# Na, Nb = 0, 0
226-
# for j in range(len(x)):
227-
# agraph_lst.append(x[j][2] + Na)
228-
# bgraph_lst.append(x[j][3] + Nb)
229-
# N_atoms_scope.append([Na, x[j][2].shape[0]])
230-
# Na += x[j][2].shape[0]
231-
# Nb += x[j][3].shape[0]
232-
# agraph = torch.cat(agraph_lst, 0)
233-
# bgraph = torch.cat(bgraph_lst, 0)
234-
# return [f_a, f_b, agraph, bgraph, N_atoms_scope]
235-
236-
237-
# def mpnn_collate_func(x):
238-
# #print("len(x) is ", len(x)) ## batch_size
239-
# #print("len(x[0]) is ", len(x[0])) ## 3--- data_process_loader.__getitem__
240-
# mpnn_feature = [i[0] for i in x]
241-
# #print("len(mpnn_feature)", len(mpnn_feature), "len(mpnn_feature[0])", len(mpnn_feature[0]))
242-
# mpnn_feature = mpnn_feature_collate_func(mpnn_feature)
243-
# from torch.utils.data.dataloader import default_collate
244-
# x_remain = [[i[1], i[2]] for i in x]
245-
# x_remain_collated = default_collate(x_remain)
246-
# return [mpnn_feature] + x_remain_collated
247-
# ## used in dataloader
248-
213+
def dgl_collate_func(x):
214+
d, p, y = zip(*x)
215+
print(d)
216+
print(p)
217+
print(y)
218+
import dgl
219+
d = dgl.batch(d)
220+
return d, torch.tensor(p), torch.tensor(y)
249221

250222
class DBTA:
251223
'''
@@ -267,6 +239,23 @@ def __init__(self, **config):
267239
self.model_drug = transformer('drug', **config)
268240
elif drug_encoding == 'MPNN':
269241
self.model_drug = MPNN(config['hidden_dim_drug'], config['mpnn_depth'])
242+
elif drug_encoding == 'DGL_GCN':
243+
self.model_drug = DGL_GCN(in_feats = 74,
244+
hidden_feats = [config['gnn_hid_dim_drug']] * config['gnn_num_layers'],
245+
activation = [config['gnn_activation']] * config['gnn_num_layers'],
246+
predictor_dim = config['hidden_dim_drug'])
247+
elif drug_encoding == 'DGL_NeuralFP':
248+
self.model_drug = DGL_NeuralFP(in_feats = 74,
249+
hidden_feats = [config['gnn_hid_dim_drug']] * config['gnn_num_layers'],
250+
max_degree = config['neuralfp_max_degree'],
251+
activation = [config['gnn_activation']] * config['gnn_num_layers'],
252+
predictor_hidden_size = config['neuralfp_predictor_hid_dim'],
253+
predictor_dim = config['hidden_dim_drug'],
254+
predictor_activation = config['neuralfp_predictor_activation'])
255+
elif drug_encoding == 'DGL_GIN_AttrMasking':
256+
self.model_drug = DGL_GIN_AttrMasking(predictor_dim = config['hidden_dim_drug'])
257+
elif drug_encoding == 'DGL_GIN_ContextPred':
258+
self.model_drug = DGL_GIN_ContextPred(predictor_dim = config['hidden_dim_drug'])
270259
else:
271260
raise AttributeError('Please use one of the available encoding method.')
272261

@@ -308,7 +297,7 @@ def test_(self, data_generator, model, repurposing_mode = False, test = False):
308297
y_label = []
309298
model.eval()
310299
for i, (v_d, v_p, label) in enumerate(data_generator):
311-
if self.drug_encoding == "MPNN" or self.drug_encoding == 'Transformer':
300+
if self.drug_encoding in ["MPNN", 'Transformer', 'DGL_GCN', 'DGL_NeuralFP', 'DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred', 'AttentiveFP']:
312301
v_d = v_d
313302
else:
314303
v_d = v_d.float().to(self.device)
@@ -387,6 +376,8 @@ def train(self, train, val = None, test = None, verbose = True):
387376
'drop_last': False}
388377
if (self.drug_encoding == "MPNN"):
389378
params['collate_fn'] = mpnn_collate_func
379+
elif self.drug_encoding in ['DGL_GCN', 'DGL_NeuralFP', 'DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred', 'AttentiveFP']:
380+
params['collate_fn'] = dgl_collate_func
390381

391382
training_generator = data.DataLoader(data_process_loader(train.index.values, train.Label.values, train, **self.config), **params)
392383
if val is not None:
@@ -402,6 +393,8 @@ def train(self, train, val = None, test = None, verbose = True):
402393

403394
if (self.drug_encoding == "MPNN"):
404395
params_test['collate_fn'] = mpnn_collate_func
396+
elif self.drug_encoding in ['DGL_GCN', 'DGL_NeuralFP', 'DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred', 'AttentiveFP']:
397+
params_test['collate_fn'] = dgl_collate_func
405398
testing_generator = data.DataLoader(data_process_loader(test.index.values, test.Label.values, test, **self.config), **params_test)
406399

407400
# early stopping
@@ -430,7 +423,7 @@ def train(self, train, val = None, test = None, verbose = True):
430423
v_p = v_p
431424
else:
432425
v_p = v_p.float().to(self.device)
433-
if self.drug_encoding == "MPNN" or self.drug_encoding == 'Transformer':
426+
if self.drug_encoding in ["MPNN", 'Transformer', 'DGL_GCN', 'DGL_NeuralFP', 'DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred', 'AttentiveFP']:
434427
v_d = v_d
435428
else:
436429
v_d = v_d.float().to(self.device)
@@ -570,7 +563,8 @@ def predict(self, df_data):
570563

571564
if (self.drug_encoding == "MPNN"):
572565
params['collate_fn'] = mpnn_collate_func
573-
566+
elif self.drug_encoding in ['DGL_GCN', 'DGL_NeuralFP', 'DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred', 'AttentiveFP']:
567+
params['collate_fn'] = dgl_collate_func
574568

575569
generator = data.DataLoader(info, **params)
576570

‎DeepPurpose/utils.py‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -625,6 +625,27 @@ def __init__(self, list_IDs, labels, df, **config):
625625
self.df = df
626626
self.config = config
627627

628+
if self.config['drug_encoding'] in ['DGL_GCN', 'DGL_NeuralFP']:
629+
from dgllife.utils import smiles_to_bigraph, CanonicalAtomFeaturizer, CanonicalBondFeaturizer
630+
self.node_featurizer = CanonicalAtomFeaturizer()
631+
self.edge_featurizer = CanonicalBondFeaturizer(self_loop = True)
632+
from functools import partial
633+
self.fc = partial(smiles_to_bigraph, add_self_loop=True)
634+
635+
elif self.config['drug_encoding'] == 'AttentiveFP':
636+
from dgllife.utils import smiles_to_bigraph, AttentiveFPAtomFeaturizer, AttentiveFPBondFeaturizer
637+
self.node_featurizer = AttentiveFPAtomFeaturizer()
638+
self.edge_featurizer = AttentiveFPBondFeaturizer(self_loop=True)
639+
from functools import partial
640+
self.fc = partial(smiles_to_bigraph, add_self_loop=True)
641+
642+
elif self.config['drug_encoding'] in ['DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred']:
643+
from dgllife.utils import smiles_to_bigraph, PretrainAtomFeaturizer, PretrainBondFeaturizer
644+
self.node_featurizer = PretrainAtomFeaturizer()
645+
self.edge_featurizer = PretrainBondFeaturizer()
646+
from functools import partial
647+
self.fc = partial(smiles_to_bigraph, add_self_loop=True)
648+
628649
def __len__(self):
629650
'Denotes the total number of samples'
630651
return len(self.list_IDs)
@@ -635,6 +656,8 @@ def __getitem__(self, index):
635656
v_d = self.df.iloc[index]['drug_encoding']
636657
if self.config['drug_encoding'] == 'CNN' or self.config['drug_encoding'] == 'CNN_RNN':
637658
v_d = drug_2_embed(v_d)
659+
elif self.config['drug_encoding'] in ['DGL_GCN', 'DGL_NeuralFP', 'DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred', 'AttentiveFP']:
660+
v_d = self.fc(smiles = v_d, node_featurizer = self.node_featurizer, edge_featurizer = self.edge_featurizer)
638661
v_p = self.df.iloc[index]['target_encoding']
639662
if self.config['target_encoding'] == 'CNN' or self.config['target_encoding'] == 'CNN_RNN':
640663
v_p = protein_2_embed(v_p)
@@ -650,7 +673,27 @@ def __init__(self, list_IDs, labels, df, **config):
650673
self.list_IDs = list_IDs
651674
self.df = df
652675
self.config = config
653-
print(df.columns.values)
676+
677+
if self.config['drug_encoding'] in ['DGL_GCN', 'DGL_NeuralFP']:
678+
from dgllife.utils import smiles_to_bigraph, CanonicalAtomFeaturizer, CanonicalBondFeaturizer
679+
self.node_featurizer = CanonicalAtomFeaturizer()
680+
self.edge_featurizer = CanonicalBondFeaturizer(self_loop = True)
681+
from functools import partial
682+
self.fc = partial(smiles_to_bigraph, add_self_loop=True)
683+
684+
elif self.config['drug_encoding'] == 'AttentiveFP':
685+
from dgllife.utils import smiles_to_bigraph, AttentiveFPAtomFeaturizer, AttentiveFPBondFeaturizer
686+
self.node_featurizer = AttentiveFPAtomFeaturizer()
687+
self.edge_featurizer = AttentiveFPBondFeaturizer(self_loop=True)
688+
from functools import partial
689+
self.fc = partial(smiles_to_bigraph, add_self_loop=True)
690+
691+
elif self.config['drug_encoding'] in ['DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred']:
692+
from dgllife.utils import smiles_to_bigraph, PretrainAtomFeaturizer, PretrainBondFeaturizer
693+
self.node_featurizer = PretrainAtomFeaturizer()
694+
self.edge_featurizer = PretrainBondFeaturizer()
695+
from functools import partial
696+
self.fc = partial(smiles_to_bigraph, add_self_loop=True)
654697
def __len__(self):
655698
'Denotes the total number of samples'
656699
return len(self.list_IDs)
@@ -661,9 +704,13 @@ def __getitem__(self, index):
661704
v_d = self.df.iloc[index]['drug_encoding_1']
662705
if self.config['drug_encoding'] == 'CNN' or self.config['drug_encoding'] == 'CNN_RNN':
663706
v_d = drug_2_embed(v_d)
707+
elif self.config['drug_encoding'] in ['DGL_GCN', 'DGL_NeuralFP', 'DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred', 'AttentiveFP']:
708+
v_d = self.fc(smiles = v_d, node_featurizer = self.node_featurizer, edge_featurizer = self.edge_featurizer)
664709
v_p = self.df.iloc[index]['drug_encoding_2']
665710
if self.config['drug_encoding'] == 'CNN' or self.config['drug_encoding'] == 'CNN_RNN':
666711
v_p = drug_2_embed(v_p)
712+
elif self.config['drug_encoding'] in ['DGL_GCN', 'DGL_NeuralFP', 'DGL_GIN_AttrMasking', 'DGL_GIN_ContextPred', 'AttentiveFP']:
713+
v_p = self.fc(smiles = v_p, node_featurizer = self.node_featurizer, edge_featurizer = self.edge_featurizer)
667714
y = self.labels[index]
668715
return v_d, v_p, y
669716

0 commit comments

Comments
 (0)