/
align_export.py
61 lines (50 loc) · 1.96 KB
/
align_export.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
import ast
from pymol import cmd, selector
from argparse import ArgumentParser
def modelChainTuple(string):
try:
model,chain=map(str, string.split(','))
return model,chain
except:
print("Type error when reading model-chain tuples")
raise TypeError
parser = ArgumentParser(description='Specify models.chains to align or extend pymol api',)
parser.add_argument('--model-chain-tuples', nargs='+', type=modelChainTuple)
parser.add_argument('--pmlxtnd')
args =parser.parse_args()
print(args)
"""
-fetch the structures(arg)
-parseout the chains of interest(arg)
-align in the same coordinate frame
-save as a pdb
"""
def chain_align_save(*args, **kwargs):
args = [ast.literal_eval(kvpair) for kvpair in args]
for pair in args:
cmd.fetch(str.lower(pair[0]))
create_subchain_object(pair[0], pair[1])
cmd.delete(str.lower(pair[0]))
for mobile in [ '{}.{}'.format(model, chain) for (model, chain) in args ][1:]:
cmd.super(mobile, "{}.{}".format(args[0][0], args[0][1]),reset=1,transform=1,quiet=0)
cmd.reset()
cmd.save('alignment.cif')
def create_subchain_object(pdbid, subchain):
tempname = 'selection_{}.{}'.format(pdbid, subchain)
cmd.select(tempname ,'m. {} and c. {}'.format( pdbid, subchain))
cmd.create('{}.{}'.format(pdbid,subchain), tempname)
cmd.delete(tempname)
if args.pmlxtnd:
cmd.extend('create_subchain_object', create_subchain_object)
cmd.extend('chain_align_save', chain_align_save);
print("Added two commands to pymol.cmd")
else:
mcpairs=args.model_chain_tuples
for pair in mcpairs:
cmd.fetch(str.lower(pair[0]))
create_subchain_object(pair[0], pair[1])
cmd.delete(str.lower(pair[0]))
for mobile in [ '{}.{}'.format(model, chain) for (model, chain) in mcpairs ][1:]:
cmd.super(mobile, "{}.{}".format(mcpairs[0][0], mcpairs[0][1]),reset=1,transform=1,quiet=0)
cmd.reset()
cmd.save('alignment.cif')