Spaces:
Runtime error
Runtime error
File size: 3,091 Bytes
753e275 605bb4b 753e275 |
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 |
import os
import shutil
import argparse
from diffab.tools.dock.hdock import HDockAntibody
from diffab.tools.runner.design_for_pdb import args_factory, design_for_pdb
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--antigen', type=str, required=True)
parser.add_argument('--antibody', type=str, default='./data/examples/3QHF_Fv.pdb')
parser.add_argument('--heavy', type=str, default='H', help='Chain id of the heavy chain.')
parser.add_argument('--light', type=str, default='L', help='Chain id of the light chain.')
parser.add_argument('--hdock_bin', type=str, default='./bin/hdock')
parser.add_argument('--createpl_bin', type=str, default='./bin/createpl')
parser.add_argument('-n', '--num_docks', type=int, default=10)
parser.add_argument('-c', '--config', type=str, default='./configs/test/codesign_single.yml')
parser.add_argument('-o', '--out_root', type=str, default='./results')
parser.add_argument('-t', '--tag', type=str, default='')
parser.add_argument('-s', '--seed', type=int, default=None)
parser.add_argument('-d', '--device', type=str, default='cuda')
parser.add_argument('-b', '--batch_size', type=int, default=16)
args = parser.parse_args()
print('[INFO] Start docking...')
hdock_missing = []
if not os.path.exists(args.hdock_bin):
hdock_missing.append(args.hdock_bin)
if not os.path.exists(args.createpl_bin):
hdock_missing.append(args.createpl_bin)
if len(hdock_missing) > 0:
print("[WARNING] The following HDOCK applications are missing:")
for f in hdock_missing:
print(f" > {f}")
print("Please download HDOCK from http://huanglab.phys.hust.edu.cn/software/hdocklite/ "
"and put `hdock` and `createpl` to the above path.")
exit()
antigen_name = os.path.basename(os.path.splitext(args.antigen)[0])
docked_pdb_dir = os.path.join(os.path.splitext(args.antigen)[0] + '_dock')
os.makedirs(docked_pdb_dir, exist_ok=True)
docked_pdb_paths = []
for fname in os.listdir(docked_pdb_dir):
if fname.endswith('.pdb'):
docked_pdb_paths.append(os.path.join(docked_pdb_dir, fname))
if len(docked_pdb_paths) < args.num_docks:
with HDockAntibody() as dock_session:
dock_session.set_antigen(args.antigen)
dock_session.set_antibody(args.antibody)
docked_tmp_paths = dock_session.dock()
for i, tmp_path in enumerate(docked_tmp_paths[:args.num_docks]):
dest_path = os.path.join(docked_pdb_dir, f"{antigen_name}_Ab_{i:04d}.pdb")
shutil.copyfile(tmp_path, dest_path)
print(f'[INFO] Copy {tmp_path} -> {dest_path}')
docked_pdb_paths.append(dest_path)
for pdb_path in docked_pdb_paths:
current_args = vars(args)
current_args['tag'] += antigen_name
design_args = args_factory(
pdb_path = pdb_path,
**current_args,
)
design_for_pdb(design_args)
if __name__ == '__main__':
main()
|