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