leadopt.py 1.49 KB
Newer Older
Harrison Green's avatar
init  
Harrison Green committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
'''
fragment prediction CLI tool
'''

import argparse
import os

import torch
import torch.nn as nn
import torch.nn.functional as F

import h5py

from leadopt.models.voxel import VoxelFingerprintNet2
from leadopt.infer import infer_all
hgarrereyn's avatar
hgarrereyn committed
16
from leadopt.pretrained import MODELS
Harrison Green's avatar
init  
Harrison Green committed
17
18
19
20
21
22
23


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('-r', '--receptor', required=True, help='Receptor file (.pdb)')
    parser.add_argument('-l', '--ligand', required=True, help='Ligand file (.sdf)')
jdurrant's avatar
jdurrant committed
24

Harrison Green's avatar
init  
Harrison Green committed
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
    parser.add_argument('-n', '--num_samples', type=int, default=16, help='Number of random rotation samples to use per prediction (default: 16)')
    parser.add_argument('-k', '--num_suggestions', type=int, default=25, help='Number of suggestions per fragment')

    parser.add_argument('-m', '--model', default=[k for k in MODELS][0], choices=[k for k in MODELS])
    parser.add_argument('-mp', '--model_path', required=True)
    parser.add_argument('-d', '--data_path', required=True)

    args = parser.parse_args()

    # load model
    m = MODELS[args.model].load(args.model_path)
    fingerprints, smiles = MODELS[args.model].get_fingerprints(args.data_path)

    # run infer step
    res = infer_all(
jdurrant's avatar
jdurrant committed
40
41
42
43
44
        model=m,
        fingerprints=fingerprints,
        smiles=smiles,
        rec_path=args.receptor,
        lig_path=args.ligand,
Harrison Green's avatar
init  
Harrison Green committed
45
46
47
48
49
50
51
52
53
54
55
        num_samples=args.num_samples,
        k=args.num_suggestions
    )

    print(res)


if __name__=='__main__':
    main()


jdurrant's avatar
jdurrant committed
56
# leadopt.py -r my_receptor.pdb -l my_ligand.sdf