mirror of
https://github.com/Saoge123/PocketFlow.git
synced 2026-06-04 12:44:22 +08:00
158 lines
7.6 KiB
Python
158 lines
7.6 KiB
Python
import argparse
|
|
import time
|
|
from pocket_flow import PocketFlow, Generate
|
|
from pocket_flow.utils import *
|
|
from pocket_flow.utils import mask_node, Protein, ComplexData, ComplexData
|
|
import collections
|
|
|
|
|
|
def parameter():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('-pkt', '--pocket', type=str, default='None', help='the pdb file of pocket in receptor')
|
|
parser.add_argument('--ckpt', type=str, default='./ckpt/ZINC-pretrained-255000.pt', help='the path of saved model')
|
|
parser.add_argument('-n', '--num_gen', type=int, default=100, help='the number of generateive molecule')
|
|
parser.add_argument('--name', type=str, default='receptor', help='receptor name')
|
|
parser.add_argument('-d', '--device', type=str, default='cuda:0', help='cuda:x or cpu')
|
|
parser.add_argument('-at', '--atom_temperature', type=float, default=1.0, help='temperature for atom sampling')
|
|
parser.add_argument('-bt', '--bond_temperature', type=float, default=1.0, help='temperature for bond sampling')
|
|
parser.add_argument('--max_atom_num', type=int, default=40, help='the max atom number for generation')
|
|
parser.add_argument('-ft', '--focus_threshold', type=float, default=0.5, help='the threshold of probility for focus atom')
|
|
parser.add_argument('-cm', '--choose_max', type=bool, default=True, help='whether choose the atom that has the highest prob as focus atom')
|
|
parser.add_argument('--min_dist_inter_mol', type=float, default=3.0, help='inter-molecular dist cutoff between protein and ligand.')
|
|
parser.add_argument('--bond_length_range', type=str, default=(1.0,2.0), help='the range of bond length for mol generation.')
|
|
parser.add_argument('-mdb', '--max_double_in_6ring', type=int, default=0, help='')
|
|
parser.add_argument('--with_print', type=bool, default=True, help='whether print SMILES in generative process')
|
|
parser.add_argument('--root_path', type=str, default='gen_results', help='the root path for saving results')
|
|
parser.add_argument('--readme', '-rm', type=str, default='None', help='description of this genrative task')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
if __name__ == '__main__':
|
|
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
|
|
|
|
args = parameter()
|
|
'''args.pocket = 'test_samples/test_pocket10/1bvr_C_rec_pocket10-surf.pdb'
|
|
args.ckpt = 'PubChem-pretrained-315000.pt'
|
|
args.choose_max = 1
|
|
args.root_path = 'gen_results/'
|
|
args.name = '1bvr'
|
|
args.num_gen = 10'''
|
|
'''
|
|
#args.name = 'HAT1'
|
|
args.root_path = 'gen_results_new'
|
|
args.device = 'cuda:0'
|
|
args.focus_threshold = 0.5'''
|
|
if args.name == 'receptor':
|
|
args.name = args.pocket.split('/')[-1].split('-')[0]
|
|
## Load Target
|
|
assert args.pocket != 'None', 'Please specify pocket !'
|
|
assert args.ckpt != 'None', 'Please specify model !'
|
|
pdb_file = args.pocket
|
|
|
|
|
|
pro_dict = Protein(pdb_file).get_atom_dict(removeHs=True, get_surf=True)
|
|
lig_dict = Ligand.empty_dict()
|
|
data = ComplexData.from_protein_ligand_dicts(
|
|
protein_dict=torchify_dict(pro_dict),
|
|
ligand_dict=torchify_dict(lig_dict),
|
|
)
|
|
|
|
## init transform
|
|
protein_featurizer = FeaturizeProteinAtom()
|
|
ligand_featurizer = FeaturizeLigandAtom(atomic_numbers=[6,7,8,9,15,16,17,35,53])
|
|
focal_masker = FocalMaker(r=6.0, num_work=16, atomic_numbers=[6,7,8,9,15,16,17,35,53])
|
|
atom_composer = AtomComposer(knn=16, num_workers=16, for_gen=True, use_protein_bond=True)
|
|
|
|
## transform data
|
|
data = RefineData()(data)
|
|
data = LigandCountNeighbors()(data)
|
|
data = protein_featurizer(data)
|
|
data = ligand_featurizer(data)
|
|
node4mask = torch.arange(data.ligand_pos.size(0))
|
|
data = mask_node(data, torch.empty([0], dtype=torch.long), node4mask, num_atom_type=9, y_pos_std=0.)
|
|
#data = focal_masker.run(data)
|
|
data = atom_composer.run(data)
|
|
|
|
## Load model
|
|
print('Loading model ...')
|
|
'''if args.device != 'cpu':
|
|
device = 'cuda:{}'.format(args.device)
|
|
else:
|
|
device = args.device'''
|
|
device = args.device
|
|
ckpt = torch.load(args.ckpt, map_location=device)
|
|
|
|
'''d=collections.OrderedDict()
|
|
for i in ckpt['model']:
|
|
if 'pos_predictor.gvp' in i:
|
|
new_i = i.replace('gvp','mlp')
|
|
d[new_i] = ckpt['model'][i]
|
|
else:
|
|
d[i] = ckpt['model'][i]
|
|
ckpt['model'] = d'''
|
|
|
|
config = ckpt['config']
|
|
model = PocketFlow(config).to(device)
|
|
model.load_state_dict(ckpt['model'])
|
|
'''torch.save({
|
|
'config': ckpt['config'],
|
|
'model': model.state_dict(),
|
|
'optimizer': ckpt['optimizer'],
|
|
'scheduler': ckpt['scheduler'],
|
|
'iteration': ckpt['iteration'],
|
|
}, './PubChem-pretrained-315000.pt')'''
|
|
print('Generating molecules ...')
|
|
temperature = [args.atom_temperature, args.bond_temperature]
|
|
# print(args.bond_length_range, type(args.bond_length_range))
|
|
if isinstance(args.bond_length_range, str):
|
|
args.bond_length_range = eval(args.bond_length_range)
|
|
generate = Generate(model, atom_composer.run, temperature=temperature, atom_type_map=[6,7,8,9,15,16,17,35,53],
|
|
num_bond_type=4, max_atom_num=args.max_atom_num, focus_threshold=args.focus_threshold,
|
|
max_double_in_6ring=args.max_double_in_6ring, min_dist_inter_mol=args.min_dist_inter_mol,
|
|
bond_length_range=args.bond_length_range, choose_max=args.choose_max, device=device)
|
|
start = time.time()
|
|
generate.generate(data, num_gen=args.num_gen, rec_name=args.name, with_print=args.with_print,
|
|
root_path=args.root_path)
|
|
os.system('cp {} {}'.format(args.ckpt, generate.out_dir))
|
|
#print('cp {} {}'.format(args.ckpt, generate.out_dir))
|
|
#with open(generate.out_dir+'/readme.txt', 'w') as fw:
|
|
# fw.write('Model copied from {}\n'.format(args.ckpt))
|
|
#if str(args.readme) != 'None':
|
|
gen_config = '\n'.join(['{}: {}'.format(k,v) for k,v in args.__dict__.items()])
|
|
with open(generate.out_dir + '/readme.txt', 'w') as fw:
|
|
fw.write(gen_config)
|
|
end = time.time()
|
|
print('Time: {}'.format(timewait(end-start)))
|
|
'''
|
|
'python main_generate.py -pkt {} --ckpt {} -n {} --name {} \
|
|
-root_path {} -d {} -at {} -bt {} --max_atom_num {} -ft {} -cm {} --with_print {}'
|
|
|
|
python main_generate.py -pkt pockets/METTL16/METTL16-pocket10-norm.pdb --ckpt ../PocketFlow/saved_model/finetuning_new-215000.pt -n 100 --name METTL16 -root_path gen_results -d 0 -at 1 -bt 1 --max_atom_num 35 -ft 0.5 -cm True --with_print True
|
|
'''
|
|
|
|
'''def make_cmd(para_dict):
|
|
cmd_temp = 'python main_generate.py -pkt {} --ckpt {} -n {} --name {} --root_path {} -d {} -at {} -bt {} --max_atom_num {} -ft {} -cm {} --with_print {}'
|
|
cmd = cmd_temp.format(
|
|
para_dict.pkt, para_dict.ckpt, para_dict.num, para_dict.name, para_dict.root_path,
|
|
para_dict.device, para_dict.atom_temp, para_dict.bond_temp, para_dict.max_atom_num,
|
|
para_dict.focus_threshold, para_dict.choose_max, para_dict.with_print
|
|
)
|
|
return cmd
|
|
|
|
para_dict = Dict({})
|
|
para_dict.pkt = 'pockets/METTL16/METTL16-pocket10-norm.pdb'
|
|
para_dict.ckpt = '../PocketFlow/saved_model/finetuning_new-215000.pt'
|
|
para_dict.num = 10000
|
|
para_dict.name = 'METTL16'
|
|
para_dict.root_path = 'gen_results'
|
|
para_dict.device = 'cuda:0'
|
|
para_dict.atom_temp = 1.0
|
|
para_dict.bond_temp = 1.0
|
|
para_dict.max_atom_num = 35
|
|
para_dict.focus_threshold = 0.5
|
|
para_dict.choose_max = True
|
|
para_dict.with_print = False
|
|
cmd_1 = '\n'.join([make_cmd(para_dict) for _ in range(5)])
|
|
open('gbp_pocket_flow_with_edge/gen_{}.sh'.format(para_dict.name),'w').write(cmd_1)'''
|