-
Notifications
You must be signed in to change notification settings - Fork 0
/
batch_format_convert.py
67 lines (49 loc) · 1.76 KB
/
batch_format_convert.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
#!/usr/bin/env python
"""
Convert pdbqt files (of ligands or receptors) to pdb files by openbabel.
This script will serach pdbqt files under /path/to/pdbqt, convert and save to the same directory.
requirement:
openbabel 2.4.90 -> obabel
input:
/path/to/pdbqt (path)
usage:
batch_format_convert.py /path/to/pdbqt [-i input_format] [-o output_format]
"""
from glob import glob
import os
import argparse
from functools import partial
# set argparser
parser = argparse.ArgumentParser()
parser.add_argument('pdbqt_path', help='path to pdbqt files')
parser.add_argument('-i', '--input_format', default='pdbqt', help='input format')
parser.add_argument('-o', '--output_format', default='pdb', help='output format')
# arg passed from command line
args = parser.parse_args()
pdbqt_path = args.pdbqt_path
in_format = args.input_format
out_format = args.output_format
def convert(in_file: str, in_format: str = 'pdbqt', out_format: str = 'pdb'):
in_format = in_format.lower()
out_format = out_format.lower()
out_file = in_file.replace(f'.{in_format}', f'.{out_format}')
cmd = f'obabel -i{in_format} {in_file} -o{out_format} -O {out_file}'
print(cmd)
os.system(cmd)
# make partial function
convert_part = partial(convert, in_format=in_format, out_format=out_format)
def main():
if not os.path.exists(pdbqt_path):
raise FileNotFoundError(f'{pdbqt_path} not found')
# use multiprocessing to speed up conversion
import multiprocessing
from multiprocessing import Pool
# check number of cpu
cpu_nr = multiprocessing.cpu_count()
print(f'cpu number: {cpu_nr}')
if cpu_nr >= 8:
cpu_nr = 8
with Pool(cpu_nr) as p:
p.map(convert_part, glob(f'{pdbqt_path}/*.{in_format}'))
if __name__ == '__main__':
main()