-
Notifications
You must be signed in to change notification settings - Fork 0
/
rawted.py
132 lines (105 loc) · 4.7 KB
/
rawted.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import argparse
import re
from biopython_wrapper import load_pdb_structure, get_chains
from dssr_wrapper import DSSRWrapper, clean_dssr_output
from general_helpers import coordinates_from_pdb_chain
from methods import run_method
from rna_tree import tree_from_version
from zss_wrapper import zss_with_descriptor
def validate_tree_args(trees: list):
tree_regex = re.compile('(v1|v2(\d+,\d+,\d+(,\d+)?)?)')
for tree in trees:
if not tree_regex.match(tree):
print('Input argument tree {} does not match required format. Aborting...'.format(tree))
exit()
def validate_method_args(methods: list):
for method in methods:
if method.upper() not in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']:
print('Input argument method {} does not match any known method. Aborting...'.format(method))
exit()
def validate_args(args):
validate_tree_args(args.trees)
validate_method_args(args.methods)
def run_rawted(args):
trees = args.trees
methods = args.methods
structure_file1 = args.structure1
chain1 = args.chain1
structure_file2 = args.structure2
chain2 = args.chain2
save_folder = args.save_folder
pdb_structure1 = load_pdb_structure(structure_file1)
pdb_structure2 = load_pdb_structure(structure_file2)
dssr1 = DSSRWrapper.load(structure_file1)
dssr2 = DSSRWrapper.load(structure_file2)
clean_dssr_output()
if not chain1:
chain1 = get_chains(pdb_structure2)[0]
else:
chains = get_chains(pdb_structure1)
if chain1 not in chains:
print('Input argument chain1 {} does not match any chains in the input structure ({}). Aborting...'.format(
chain1, chains)
)
exit()
if not chain2:
chain2 = get_chains(pdb_structure2)[0]
else:
chains = get_chains(pdb_structure2)
if chain2 not in chains:
print('Input argument chain2 {} does not match any chains in the input structure ({}). Aborting...'.format(
chain2, chains)
)
exit()
pdb_chain1 = pdb_structure1[0][chain1]
pdb_chain2 = pdb_structure2[0][chain2]
dssr_nts1 = dssr1.nts_for_chain(chain1)
dssr_nts2 = dssr2.nts_for_chain(chain2)
for tree_descriptor in trees:
tree_version = tree_descriptor.split(',')[0]
tree1 = tree_from_version(tree_version).from_nts(dssr_nts1, pdb_chain1)
tree2 = tree_from_version(tree_version).from_nts(dssr_nts2, pdb_chain2)
ted_matrix = zss_with_descriptor(tree1, tree2, tree_descriptor)
for method in methods:
save_args = None
if save_folder:
save_args = {
'structure1': pdb_structure1.copy(),
'structure2': pdb_structure2.copy(),
'chain1': chain1,
'chain2': chain2,
'filename1': '{}/{}_{}_s1.pdb'.format(save_folder, tree_descriptor.replace(',', '-'), method.lower()),
'filename2': '{}/{}_{}_s2.pdb'.format(save_folder, tree_descriptor.replace(',', '-'), method.lower())
}
rmsd, psi, _, _ = run_method(
method,
coordinates_from_pdb_chain(pdb_chain1), coordinates_from_pdb_chain(pdb_chain2),
tree1, tree2,
ted_matrix,
save_args=save_args
)
print('--------------------')
print('Rawted finished.')
print('Structure 1: {}'.format(structure_file1))
print('Structure 2: {}'.format(structure_file2))
print('Tree descriptor: {}'.format(tree_descriptor))
print('Method: {}'.format(method))
print('--------------------')
print('>> Results <<')
print('RMSD: {}'.format(round(rmsd, 4)))
print('PSI: {}'.format(round(psi, 4)))
print('--------------------')
def main():
parser = argparse.ArgumentParser()
parser.add_argument('structure1', type=str, help='Structure 1 to process')
parser.add_argument('chain1', nargs='?', type=str, help='Chain of structure 1 to process')
parser.add_argument('structure2', type=str, help='Structure 2 to process')
parser.add_argument('chain2', nargs='?', type=str, help='Chain of structure 2 to process')
parser.add_argument('--trees', '-t', nargs='*', type=str, default=['v1'], help='Tree arguments')
parser.add_argument('--methods', '-m', nargs='*', type=str, default=['J'], help='RAWTED methods to use')
parser.add_argument('--save_folder', '-s', nargs='?', type=str, help='Save folder')
args = parser.parse_args()
validate_args(args)
run_rawted(args)
if __name__ == '__main__':
main()