1 2from __future__ import print_function 3 4from pymol.wizard import Wizard 5from pymol import cmd 6import pymol 7 8default_mode = 'labchg' 9 10class Charge(Wizard): 11 12 def __init__(self,_self=cmd): 13 Wizard.__init__(self,_self) 14 15 self.modes = [ 16 'labchg', 17 'addchg', 18 'cpychg', 19 'zrochg', 20 'mzochg', 21 'cbachg', 22 'movchg', 23 'rbachg', 24 'sumchg', 25 ] 26 27 self.mode = default_mode 28 self.status = 0 29 30 self.mode_name = { 31 'labchg':'Show', 32 'cpychg':'Copy', 33 'addchg':'Add', 34 'zrochg':'Zero', 35 'mzochg':'Move & Zero (atom)', 36 'cbachg':'Move & Zero (resi)', 37 'movchg':'Move & Remove (atom)', 38 'rbachg':'Move & Remove (resi)', 39 'sumchg':'Get Total Charge' 40 } 41 42 # initialize mode menu 43 44 smm = [] 45 smm.append([ 2, 'Atom Charge Mode', '' ]) 46 for a in self.modes: 47 smm.append([ 1, self.mode_name[a], 'cmd.get_wizard().set_mode("'+a+'")']) 48 49 self.menu['mode']=smm 50 51 self.memory = 0 52 53 _self.edit_mode() 54 55 56 def get_panel(self): 57 return [ 58 [ 1, 'Charge Wizard',''], 59 [ 3, self.mode_name[self.mode],'mode'], 60 [ 2, 'Clear','cmd.get_wizard().clear()'], 61 [ 2, 'Done','cmd.set_wizard()'], 62 ] 63 64 def cleanup(self): 65 global default_mode 66 default_mode = self.mode 67 self.clear() 68 69 def clear(self): 70 self.set_status(0) 71 if 'wcharge' in self.cmd.get_names('selections'): 72 if self.mode!='sumchg': 73 self.cmd.edit("wcharge") 74 self.cmd.label("pkmol",'') # fastest clear command 75 else: 76 self.cmd.label("wcharge",'') # fastest clear command 77 self.cmd.delete("wcharge") 78 self.cmd.unpick() 79 self.cmd.unpick() 80 self.cmd.refresh_wizard() 81 82 def get_prompt(self): 83 self.prompt = None 84 if self.mode == 'cpychg': 85 if self.status==0: 86 self.prompt = [ 'Pick source atom...' ] 87 elif self.status==1: 88 self.prompt = [ 'Pick destination atom on which to assign charge %6.4f'%self.partial_charge ] 89 if self.mode == 'addchg': 90 if self.status==0: 91 self.prompt = [ 'Pick source atom...' ] 92 elif self.status==1: 93 self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ] 94 if self.mode == 'mzochg': 95 if self.status==0: 96 self.prompt = [ 'Pick source atom to copy and zero...' ] 97 elif self.status==1: 98 self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ] 99 if self.mode == 'movchg': 100 if self.status==0: 101 self.prompt = [ 'Pick source atom to copy and destroy...' ] 102 elif self.status==1: 103 self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ] 104 if self.mode == 'sumchg': 105 if self.status==0: 106 self.prompt = [ 'Pick an atom on the chain...' ] 107 if self.status==1: 108 self.prompt = [ 'Total charge on the chain is %6.4f'%self.partial_charge, 109 'Pick an atom on the chain...' ] 110 if self.mode == 'cbachg': 111 if self.status==0: 112 self.prompt = [ 'Pick source residue to copy and zero...' ] 113 elif self.status==1: 114 self.prompt = [ 'Pick destination residue on which to add charges.'] 115 if self.mode == 'rbachg': 116 if self.status==0: 117 self.prompt = [ 'Pick source residue to copy and remove...' ] 118 elif self.status==1: 119 self.prompt = [ 'Pick destination residue on which to add charges.'] 120 if self.mode == 'zrochg': 121 self.prompt = [ 'Pick atom on which to zero charge...' ] 122 123 124 if self.mode == 'labchg': 125 self.prompt = [ 'Pick atom on which to show charge...' ] 126 127 if "wcharge" in self.cmd.get_names('selections'): 128 pymol.stored.charge = 0 129 if self.cmd.iterate("(byres wcharge)", 130 "stored.charge = stored.charge + partial_charge"): 131 self.prompt.insert(0,"Total charge on the residue is %6.4f"%pymol.stored.charge) 132 133 return self.prompt 134 135 def set_mode(self,mode): 136 if mode in self.modes: 137 self.mode = mode 138 self.status = 0 139 self.cmd.refresh_wizard() 140 141 def set_status(self,status): 142 self.status = status 143 self.cmd.refresh_wizard() 144 145 def do_pick(self,bondFlag): 146 if bondFlag: 147 print(" Error: please select a single atom") 148 149 if self.mode == 'cpychg': 150 # picking up 151 if self.status==0: 152 if self.cmd.iterate("(pk1)","stored.charge = partial_charge"): 153 self.partial_charge = pymol.stored.charge 154 self.status = 1 155 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 156 self.cmd.select("wcharge","(pk1)") 157 self.cmd.unpick() 158 self.cmd.enable("wcharge") 159 160 # dropping off 161 elif self.status==1: 162 pymol.stored.charge=self.partial_charge 163 if self.cmd.alter("(pk1)","partial_charge = stored.charge"): 164 self.status = 0 165 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 166 self.cmd.select("wcharge","(pk1)") 167 self.cmd.unpick() 168 self.cmd.enable("wcharge") 169 170 if self.mode == 'cbachg' or self.mode == 'rbachg': 171 # picking up 172 if self.status==0: 173 pymol.stored.chg_dict = {} 174 if self.cmd.iterate("(byres pk1)","stored.chg_dict[name] = partial_charge"): 175 self.charge_dict = pymol.stored.chg_dict 176 self.status = 1 177 self.cmd.label("(byres pk1)","'%6.4f'%partial_charge") 178 self.cmd.select("wcharge","(byres pk1)") 179 self.cmd.unpick() 180 self.cmd.enable("wcharge") 181 182 # dropping off 183 elif self.status==1: 184 pymol.stored.valid_atoms = [] 185 if self.cmd.iterate("(byres pk1)","stored.valid_atoms.append(name)"): 186 kees = list(self.charge_dict.keys()) 187 valid_dict = {} 188 for a in pymol.stored.valid_atoms: 189 if a in kees: 190 valid_dict[a] = 1 191 pymol.stored.chg_dict = self.charge_dict 192 # copy/add charges 193 for a in valid_dict.keys(): 194 self.cmd.alter("((byres pk1) and name %s)"%a, 195 "partial_charge = partial_charge + stored.chg_dict[name]") 196 if self.mode == 'rbachg': 197 self.cmd.remove("((wcharge) and name %s)"%a) 198 else: 199 self.cmd.alter("((wcharge) and name %s)"%a, 200 "partial_charge = 0") 201 self.status = 0 202 # update labels 203 self.cmd.label("(wcharge or (byres pk1))","'%6.4f'%partial_charge") 204 # show which atoms had charges moved 205 self.cmd.select("wcharge","(none)") 206 for a in valid_dict.keys(): 207 self.cmd.select("wcharge","(wcharge or ((byres pk1) and name %s))"%a) 208 self.cmd.unpick() 209 self.cmd.enable("wcharge") 210 211 if self.mode == 'addchg': 212 # picking up 213 if self.status==0: 214 if self.cmd.iterate("(pk1)","stored.charge = partial_charge"): 215 self.partial_charge = pymol.stored.charge 216 self.status = 1 217 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 218 self.cmd.select("wcharge","(pk1)") 219 self.cmd.unpick() 220 self.cmd.enable("wcharge") 221 222 # dropping off 223 elif self.status==1: 224 pymol.stored.charge=self.partial_charge 225 if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"): 226 self.status = 0 227 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 228 self.cmd.select("wcharge","(pk1)") 229 self.cmd.unpick() 230 self.cmd.enable("wcharge") 231 232 if self.mode == 'mzochg': 233 # picking up 234 if self.status==0: 235 if self.cmd.iterate("(pk1)","stored.charge = partial_charge"): 236 self.partial_charge = pymol.stored.charge 237 self.status = 1 238 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 239 self.cmd.select("wcharge","(pk1)") 240 self.cmd.unpick() 241 self.cmd.enable("wcharge") 242 243 # dropping off 244 elif self.status==1: 245 pymol.stored.charge=self.partial_charge 246 if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"): 247 self.status = 0 248 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 249 self.cmd.alter("(wcharge)","partial_charge=0") 250 self.cmd.label("(wcharge)","'%6.4f'%partial_charge") 251 self.cmd.select("wcharge","(pk1)") 252 self.cmd.unpick() 253 self.cmd.enable("wcharge") 254 255 if self.mode == 'movchg': 256 # picking up 257 if self.status==0: 258 if self.cmd.iterate("(pk1)","stored.charge = partial_charge"): 259 self.partial_charge = pymol.stored.charge 260 self.status = 1 261 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 262 self.cmd.select("wcharge","(pk1)") 263 self.cmd.unpick() 264 self.cmd.enable("wcharge") 265 266 # dropping off 267 elif self.status==1: 268 pymol.stored.charge=self.partial_charge 269 if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"): 270 self.status = 0 271 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 272 self.cmd.remove("wcharge") 273 self.cmd.select("wcharge","(pk1)") 274 self.cmd.unpick() 275 self.cmd.enable("wcharge") 276 277 if self.mode == 'zrochg': 278 if self.cmd.alter("(pk1)","partial_charge = 0.0"): 279 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 280 self.cmd.select("wcharge","(pk1)") 281 self.cmd.unpick() 282 self.cmd.enable("wcharge") 283 284 if self.mode == 'labchg': 285 self.cmd.label("(pk1)","'%6.4f'%partial_charge") 286 self.cmd.select("wcharge","(pk1)") 287 self.cmd.unpick() 288 self.cmd.enable("wcharge") 289 290 if self.mode == 'sumchg': 291 pymol.stored.charge = 0.0 292 if self.cmd.iterate("(pkmol)","stored.charge = stored.charge + partial_charge"): 293 self.partial_charge = pymol.stored.charge 294 self.status = 1 295 self.cmd.select("wcharge","(pkmol)") 296 self.cmd.unpick() 297 self.cmd.enable("wcharge") 298 299 self.cmd.refresh_wizard() 300