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