-
Notifications
You must be signed in to change notification settings - Fork 5
/
sample_multi.py
157 lines (125 loc) · 5.36 KB
/
sample_multi.py
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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import argparse
import os
import torch
from src import utils
from src.model_multi import DDPM
from src.visualizer import save_xyz_file, save_xyz_file_fa
from src.datasets import collate_mr, MultiRDataset_anchor
from tqdm import tqdm
import subprocess
import time
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint', action='store', type=str, required=True)
parser.add_argument('--samples', action='store', type=str, required=True)
parser.add_argument('--data', action='store', type=str, required=False, default=None)
parser.add_argument('--prefix', action='store', type=str, required=True)
parser.add_argument('--n_samples', action='store', type=int, required=True)
parser.add_argument('--n_steps', action='store', type=int, required=False, default=None)
parser.add_argument('--device', action='store', type=str, required=True)
parser.add_argument('--rgroup_size_model', action='store', type=str, required=False, default=None)
args = parser.parse_args()
experiment_name = args.checkpoint.split('/')[-1].replace('.ckpt', '')
output_dir = os.path.join(args.samples, experiment_name)
os.makedirs(output_dir, exist_ok=True)
def check_if_generated(_output_dir, _uuids, n_samples):
generated = True
starting_points = []
for _uuid in _uuids:
uuid_dir = os.path.join(_output_dir, _uuid)
numbers = []
for fname in os.listdir(uuid_dir):
try:
num = int(fname.split('_')[0])
numbers.append(num)
except:
continue
if len(numbers) == 0 or max(numbers) != n_samples - 1:
generated = False
if len(numbers) == 0:
starting_points.append(0)
else:
starting_points.append(max(numbers) - 1)
if len(starting_points) > 0:
starting = min(starting_points)
else:
starting = None
return generated, starting
collate_fn = collate_mr
sample_fn = None
# Loading model form checkpoint (all hparams will be automatically set)
model = DDPM.load_from_checkpoint(args.checkpoint, map_location=args.device)
# Possibility to evaluate on different datasets (e.g., on CASF instead of ZINC)
model.val_data_prefix = args.prefix
# In case <Anonymous> will run my model or vice versa
if args.data is not None:
model.data_path = args.data
# Less sampling steps
if args.n_steps is not None:
model.edm.T = args.n_steps
# Setting up the model
model = model.eval().to(args.device)
model.setup(stage='val')
model.batch_size = 1
# Getting the dataloader
dataloader = model.val_dataloader(collate_fn=collate_fn)
print(f'Dataloader contains {len(dataloader)} batches')
center_of_mass_list = []
time_start = time.time()
for batch_idx, data in enumerate(dataloader):
uuids = []
true_names = []
scaf_names = []
pock_names = []
for uuid in data['uuid']:
uuid = str(uuid)
uuids.append(uuid)
true_names.append(f'{uuid}/true')
scaf_names.append(f'{uuid}/scaf')
pock_names.append(f'{uuid}/pock')
os.makedirs(os.path.join(output_dir, uuid), exist_ok=True)
generated, starting_point = check_if_generated(output_dir, uuids, args.n_samples)
if generated:
print(f'Already generated batch={batch_idx}, max_uuid={max(uuids)}')
continue
if starting_point > 0:
print(f'Generating {args.n_samples - starting_point} for batch={batch_idx}')
h, x, node_mask, scaf_mask = data['one_hot'], data['positions'], data['atom_mask'], data['scaffold_mask']
node_mask = data['atom_mask'] - data['pocket_mask']
scaf_mask = data['scaffold_only_mask']
pock_mask = data['pocket_mask']
save_xyz_file_fa(output_dir, h, x, pock_mask, pock_names)
# Saving ground-truth molecules
save_xyz_file_fa(output_dir, h, x, node_mask, true_names)
# Saving scaffold
save_xyz_file_fa(output_dir, h, x, scaf_mask, scaf_names)
# Sampling and saving generated molecules
for i in tqdm(range(starting_point, args.n_samples), desc=str(batch_idx)):
chain, node_mask, mean = model.sample_chain(data, sample_fn=sample_fn, keep_frames=1)
x = chain[-1][:, :, :model.n_dims]
h = chain[-1][:, :, model.n_dims:]
x += mean
x_rgroup_tmp = x * data['rgroup_mask_batch_new']
x_scaf_ori_tmp = data['positions'] * data['scaffold_mask']
cnt = 0
for k in range(data['batch_new_len_tensor'].shape[0]):
for j in range(data['batch_new_len_tensor'][k]):
x_scaf_ori_tmp[k] += x_rgroup_tmp[cnt]
cnt += 1
h_rgroup_tmp = h * data['rgroup_mask_batch_new']
h_scaf_ori_tmp = data['one_hot'] * data['scaffold_mask']
cnt = 0
for k in range(data['batch_new_len_tensor'].shape[0]):
for j in range(data['batch_new_len_tensor'][k]):
h_scaf_ori_tmp[k] += h_rgroup_tmp[cnt]
cnt += 1
x = x_scaf_ori_tmp
h = h_scaf_ori_tmp
node_mask = data['atom_mask'] - data['pocket_mask']
pred_names = [f'{uuid}/{i}' for uuid in uuids]
save_xyz_file_fa(output_dir, h, x, node_mask, pred_names)
for j in range(len(pred_names)):
out_xyz = f'{output_dir}/{pred_names[j]}_.xyz'
out_sdf = f'{output_dir}/{pred_names[j]}_.sdf'
subprocess.run(f'obabel {out_xyz} -O {out_sdf} 2> /dev/null', shell=True)
time_end = time.time()
print('sample time:', time_end - time_start, 's')