297 lines
12 KiB
Python
297 lines
12 KiB
Python
from pymol.wizard import Wizard
|
|
from pymol import cmd
|
|
import pymol
|
|
|
|
default_mode = 'labchg'
|
|
|
|
class Charge(Wizard):
|
|
|
|
def __init__(self,_self=cmd):
|
|
Wizard.__init__(self,_self)
|
|
|
|
self.modes = [
|
|
'labchg',
|
|
'addchg',
|
|
'cpychg',
|
|
'zrochg',
|
|
'mzochg',
|
|
'cbachg',
|
|
'movchg',
|
|
'rbachg',
|
|
'sumchg',
|
|
]
|
|
|
|
self.mode = default_mode
|
|
self.status = 0
|
|
|
|
self.mode_name = {
|
|
'labchg':'Show',
|
|
'cpychg':'Copy',
|
|
'addchg':'Add',
|
|
'zrochg':'Zero',
|
|
'mzochg':'Move & Zero (atom)',
|
|
'cbachg':'Move & Zero (resi)',
|
|
'movchg':'Move & Remove (atom)',
|
|
'rbachg':'Move & Remove (resi)',
|
|
'sumchg':'Get Total Charge'
|
|
}
|
|
|
|
# initialize mode menu
|
|
|
|
smm = []
|
|
smm.append([ 2, 'Atom Charge Mode', '' ])
|
|
for a in self.modes:
|
|
smm.append([ 1, self.mode_name[a], 'cmd.get_wizard().set_mode("'+a+'")'])
|
|
|
|
self.menu['mode']=smm
|
|
|
|
self.memory = 0
|
|
|
|
_self.edit_mode()
|
|
|
|
|
|
def get_panel(self):
|
|
return [
|
|
[ 1, 'Charge Wizard',''],
|
|
[ 3, self.mode_name[self.mode],'mode'],
|
|
[ 2, 'Clear','cmd.get_wizard().clear()'],
|
|
[ 2, 'Done','cmd.set_wizard()'],
|
|
]
|
|
|
|
def cleanup(self):
|
|
global default_mode
|
|
default_mode = self.mode
|
|
self.clear()
|
|
|
|
def clear(self):
|
|
self.set_status(0)
|
|
if 'wcharge' in self.cmd.get_names('selections'):
|
|
if self.mode!='sumchg':
|
|
self.cmd.edit("wcharge")
|
|
self.cmd.label("pkmol",'') # fastest clear command
|
|
else:
|
|
self.cmd.label("wcharge",'') # fastest clear command
|
|
self.cmd.delete("wcharge")
|
|
self.cmd.unpick()
|
|
self.cmd.unpick()
|
|
self.cmd.refresh_wizard()
|
|
|
|
def get_prompt(self):
|
|
self.prompt = None
|
|
if self.mode == 'cpychg':
|
|
if self.status==0:
|
|
self.prompt = [ 'Pick source atom...' ]
|
|
elif self.status==1:
|
|
self.prompt = [ 'Pick destination atom on which to assign charge %6.4f'%self.partial_charge ]
|
|
if self.mode == 'addchg':
|
|
if self.status==0:
|
|
self.prompt = [ 'Pick source atom...' ]
|
|
elif self.status==1:
|
|
self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ]
|
|
if self.mode == 'mzochg':
|
|
if self.status==0:
|
|
self.prompt = [ 'Pick source atom to copy and zero...' ]
|
|
elif self.status==1:
|
|
self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ]
|
|
if self.mode == 'movchg':
|
|
if self.status==0:
|
|
self.prompt = [ 'Pick source atom to copy and destroy...' ]
|
|
elif self.status==1:
|
|
self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ]
|
|
if self.mode == 'sumchg':
|
|
if self.status==0:
|
|
self.prompt = [ 'Pick an atom on the chain...' ]
|
|
if self.status==1:
|
|
self.prompt = [ 'Total charge on the chain is %6.4f'%self.partial_charge,
|
|
'Pick an atom on the chain...' ]
|
|
if self.mode == 'cbachg':
|
|
if self.status==0:
|
|
self.prompt = [ 'Pick source residue to copy and zero...' ]
|
|
elif self.status==1:
|
|
self.prompt = [ 'Pick destination residue on which to add charges.']
|
|
if self.mode == 'rbachg':
|
|
if self.status==0:
|
|
self.prompt = [ 'Pick source residue to copy and remove...' ]
|
|
elif self.status==1:
|
|
self.prompt = [ 'Pick destination residue on which to add charges.']
|
|
if self.mode == 'zrochg':
|
|
self.prompt = [ 'Pick atom on which to zero charge...' ]
|
|
|
|
|
|
if self.mode == 'labchg':
|
|
self.prompt = [ 'Pick atom on which to show charge...' ]
|
|
|
|
if "wcharge" in self.cmd.get_names('selections'):
|
|
pymol.stored.charge = 0
|
|
if self.cmd.iterate("(byres wcharge)",
|
|
"stored.charge = stored.charge + partial_charge"):
|
|
self.prompt.insert(0,"Total charge on the residue is %6.4f"%pymol.stored.charge)
|
|
|
|
return self.prompt
|
|
|
|
def set_mode(self,mode):
|
|
if mode in self.modes:
|
|
self.mode = mode
|
|
self.status = 0
|
|
self.cmd.refresh_wizard()
|
|
|
|
def set_status(self,status):
|
|
self.status = status
|
|
self.cmd.refresh_wizard()
|
|
|
|
def do_pick(self,bondFlag):
|
|
if bondFlag:
|
|
print(" Error: please select a single atom")
|
|
|
|
if self.mode == 'cpychg':
|
|
# picking up
|
|
if self.status==0:
|
|
if self.cmd.iterate("(pk1)","stored.charge = partial_charge"):
|
|
self.partial_charge = pymol.stored.charge
|
|
self.status = 1
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
# dropping off
|
|
elif self.status==1:
|
|
pymol.stored.charge=self.partial_charge
|
|
if self.cmd.alter("(pk1)","partial_charge = stored.charge"):
|
|
self.status = 0
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
if self.mode == 'cbachg' or self.mode == 'rbachg':
|
|
# picking up
|
|
if self.status==0:
|
|
pymol.stored.chg_dict = {}
|
|
if self.cmd.iterate("(byres pk1)","stored.chg_dict[name] = partial_charge"):
|
|
self.charge_dict = pymol.stored.chg_dict
|
|
self.status = 1
|
|
self.cmd.label("(byres pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(byres pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
# dropping off
|
|
elif self.status==1:
|
|
pymol.stored.valid_atoms = []
|
|
if self.cmd.iterate("(byres pk1)","stored.valid_atoms.append(name)"):
|
|
kees = list(self.charge_dict.keys())
|
|
valid_dict = {}
|
|
for a in pymol.stored.valid_atoms:
|
|
if a in kees:
|
|
valid_dict[a] = 1
|
|
pymol.stored.chg_dict = self.charge_dict
|
|
# copy/add charges
|
|
for a in valid_dict.keys():
|
|
self.cmd.alter("((byres pk1) and name %s)"%a,
|
|
"partial_charge = partial_charge + stored.chg_dict[name]")
|
|
if self.mode == 'rbachg':
|
|
self.cmd.remove("((wcharge) and name %s)"%a)
|
|
else:
|
|
self.cmd.alter("((wcharge) and name %s)"%a,
|
|
"partial_charge = 0")
|
|
self.status = 0
|
|
# update labels
|
|
self.cmd.label("(wcharge or (byres pk1))","'%6.4f'%partial_charge")
|
|
# show which atoms had charges moved
|
|
self.cmd.select("wcharge","(none)")
|
|
for a in valid_dict.keys():
|
|
self.cmd.select("wcharge","(wcharge or ((byres pk1) and name %s))"%a)
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
if self.mode == 'addchg':
|
|
# picking up
|
|
if self.status==0:
|
|
if self.cmd.iterate("(pk1)","stored.charge = partial_charge"):
|
|
self.partial_charge = pymol.stored.charge
|
|
self.status = 1
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
# dropping off
|
|
elif self.status==1:
|
|
pymol.stored.charge=self.partial_charge
|
|
if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"):
|
|
self.status = 0
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
if self.mode == 'mzochg':
|
|
# picking up
|
|
if self.status==0:
|
|
if self.cmd.iterate("(pk1)","stored.charge = partial_charge"):
|
|
self.partial_charge = pymol.stored.charge
|
|
self.status = 1
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
# dropping off
|
|
elif self.status==1:
|
|
pymol.stored.charge=self.partial_charge
|
|
if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"):
|
|
self.status = 0
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.alter("(wcharge)","partial_charge=0")
|
|
self.cmd.label("(wcharge)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
if self.mode == 'movchg':
|
|
# picking up
|
|
if self.status==0:
|
|
if self.cmd.iterate("(pk1)","stored.charge = partial_charge"):
|
|
self.partial_charge = pymol.stored.charge
|
|
self.status = 1
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
# dropping off
|
|
elif self.status==1:
|
|
pymol.stored.charge=self.partial_charge
|
|
if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"):
|
|
self.status = 0
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.remove("wcharge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
if self.mode == 'zrochg':
|
|
if self.cmd.alter("(pk1)","partial_charge = 0.0"):
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
if self.mode == 'labchg':
|
|
self.cmd.label("(pk1)","'%6.4f'%partial_charge")
|
|
self.cmd.select("wcharge","(pk1)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
if self.mode == 'sumchg':
|
|
pymol.stored.charge = 0.0
|
|
if self.cmd.iterate("(pkmol)","stored.charge = stored.charge + partial_charge"):
|
|
self.partial_charge = pymol.stored.charge
|
|
self.status = 1
|
|
self.cmd.select("wcharge","(pkmol)")
|
|
self.cmd.unpick()
|
|
self.cmd.enable("wcharge")
|
|
|
|
self.cmd.refresh_wizard()
|