1
2from __future__ import print_function
3
4from pymol.wizard import Wizard
5from pymol import cmd
6import pymol
7import traceback
8
9sele_prefix = "_pf_s_"
10sele_prefix_len = len(sele_prefix)
11
12dist_prefix = "_pf_d_"
13
14indi_sele = "_indicate_pf"
15
16class Pair_fit(Wizard):
17
18    def __init__(self,_self=cmd):
19
20        Wizard.__init__(self,_self)
21
22        self.memory = 0
23        self.n_pair = 0
24        self.status = 0 # 0 no atoms selections, 1 atom selected
25        self.message = None
26
27        self.selection_mode = cmd.get_setting_int("mouse_selection_mode")
28        cmd.set("mouse_selection_mode",0) # set selection mode to atomic
29        cmd.deselect() # disable the active selection (if any)
30
31    def get_panel(self):
32        return [
33            [ 1, 'Pair Fitting',''],
34            [ 2, 'Fit %d Pairs'%self.n_pair,'cmd.get_wizard().fit()'],
35            [ 2, 'Delete Last Pair','cmd.get_wizard().remove_last()'],
36            [ 2, 'Redraw','cmd.get_wizard().update_dashes()'],
37            [ 2, 'Clear','cmd.get_wizard().clear()'],
38            [ 2, 'Done','cmd.set_wizard()'],
39            ]
40
41    def cleanup(self):
42        self.clear()
43        cmd.set("mouse_selection_mode",self.selection_mode) # restore selection mode
44
45    def clear(self):
46        cmd.delete(sele_prefix+"*")
47        cmd.delete(dist_prefix+"*")
48        cmd.delete(indi_sele)
49        lst = cmd.get_names('selections')
50        self.n_pair = 0
51        self.status = 0
52        self.message = None
53        cmd.unpick()
54        cmd.refresh_wizard()
55
56    def get_prompt(self):
57        self.prompt = None
58        if self.status==0:
59            self.prompt = [ 'Pick the mobile atom...']
60        elif self.status==1:
61            self.prompt = [ 'Pick the target atom...' ]
62        if self.message is not None:
63            self.prompt.append(self.message)
64        return self.prompt
65
66    def set_status(self,status):
67        self.status = status
68        cmd.refresh_wizard()
69
70    def get_sele_list(self,mode='all'):
71        lst = cmd.get_names('selections')
72        lst = [x for x in lst if x[0:sele_prefix_len]==sele_prefix]
73        lst.sort()
74        if mode == 'mobile': # mobile
75            lst=[x for x in lst if x[-1:]=='b']
76        elif mode == 'target': # target
77            lst=[x for x in lst if x[-1:]=='a']
78        return lst
79
80    def fit(self):
81        # build up the pair-wise list of selections
82        cmd.delete(dist_prefix+"*")
83        lst = self.get_sele_list()
84        c = 0
85        args = []
86        while 1:
87            if not len(lst): break
88            a = lst.pop()
89            if not len(lst): break
90            b = lst.pop()
91            args.append(a)
92            args.append(b)
93        # do the fit
94        if len(args):
95            cmd.push_undo(args[0])
96            dist = cmd.pair_fit(*args)
97            self.message = "RMS over %d pairs = %5.3f"%(self.n_pair,dist)
98            cmd.refresh_wizard()
99        self.update_dashes()
100
101    def remove_last(self):
102        # build up the pair-wise list of selections
103        cmd.delete(dist_prefix+"*")
104        lst = self.get_sele_list()
105        if len(lst):
106            cmd.delete(lst.pop())
107            if len(lst):
108                cmd.delete(lst.pop())
109            self.n_pair = self.n_pair - 1
110        self.update_dashes()
111        self.status=0
112        cmd.refresh_wizard()
113
114    def update_dashes(self):
115        cmd.delete(dist_prefix+"*")
116        lst = self.get_sele_list()
117        c = 0
118        while 1:
119            if not len(lst): break
120            a = lst.pop()
121            if not len(lst): break
122            b = lst.pop()
123            name = dist_prefix+str(c)
124            cmd.dist(name,a,b,width=7,length=0.05,gap=0.05)
125            cmd.hide('label',name)
126            cmd.enable(name)
127            c = c + 1
128
129    def check_same_object(self,lst,sele):
130        if not len(lst):
131            return 1
132        else:
133            if cmd.count_atoms("((byobj %s) and %s)"%(lst[0],sele),quiet=1):
134                return 1
135        return 0
136
137    def check_different_object(self,lst,sele):
138        if not len(lst):
139            return 1
140        else:
141            if not cmd.count_atoms("((byobj %s) and %s)"%(lst[0],sele),quiet=1):
142                return 1
143        return 0
144
145    def do_select(self,name): # map selects into picks
146        cmd.unpick()
147        try:
148            cmd.edit(name + " and not " + sele_prefix + "*") # note, using new object name wildcards
149            cmd.delete(name)
150            self.do_pick(0)
151        except pymol.CmdException:
152            traceback.print_exc()
153            pass
154
155    def do_pick(self,bondFlag):
156        if bondFlag:
157            self.message = "Error: please select an atom, not a bond."
158            print(self.message)
159        else:
160            if self.status==0:
161                lst = self.get_sele_list(mode='mobile')
162                if not self.check_same_object(lst,"(pk1)"):
163                    self.message = "Error: must select an atom in the same object as before."
164                    print(self.message)
165                else:
166                    name = sele_prefix + "%02db"%self.n_pair # mobile end in 'b'
167                    cmd.select(name,"(pk1)")
168                    cmd.unpick()
169                    cmd.select(indi_sele,name)
170                    cmd.enable(indi_sele)
171                    self.status = 1
172                    self.message = None
173            elif self.status==1:
174                lst = self.get_sele_list(mode='target')
175                if not self.check_same_object(lst,"(pk1)"):
176                    self.message = "Error: must select an atom in the same object as before."
177                    print(self.message)
178                else:
179                    lst = self.get_sele_list(mode='mobile')
180                    if not self.check_different_object(lst,"(pk1)"):
181                        self.message = "Error: target atom must be in a distinct object."
182                        print(self.message)
183                    else:
184                        name = sele_prefix + "%02da"%self.n_pair # target end in 'a'
185                        cmd.select(name,"(pk1)")
186                        cmd.unpick()
187                        cmd.select(indi_sele,name)
188                        cmd.enable(indi_sele)
189                        self.n_pair = self.n_pair + 1
190                        self.status = 0
191                        self.update_dashes()
192
193        cmd.refresh_wizard()
194