Files
pymol-open-source/data/startup/apbs_gui/__init__.py
Jarrett Johnson 87d27daafb port over qt6 migrations
Fixes #498
2026-02-24 01:35:00 -05:00

451 lines
14 KiB
Python

'''
PyMOL APBS GUI Plugin
(c) Schrodinger, Inc.
'''
import sys
import os
import importlib
from pymol.Qt import QtCore, QtWidgets
from pymol.Qt.utils import loadUi, AsyncFunc, MainThreadCaller
getOpenFileNames = QtWidgets.QFileDialog.getOpenFileNames
from . import electrostatics
from .qtwidgets import ResizableMessageBox as QMessageBox
class SilentAbort(Exception):
pass
class StdOutCapture:
'''
Redirect stdout and/or stderr to a temporary file until 'release' is called.
'''
def __init__(self, stdout=True, stderr=True):
import tempfile
self.temp = tempfile.TemporaryFile('w+')
self.save_stdout = None
self.save_stderr = None
if stdout:
self.save_stdout = sys.stdout
sys.stdout = self.temp
if stderr:
self.save_stderr = sys.stderr
sys.stderr = self.temp
def release(self):
if self.save_stdout is not None:
sys.stdout = self.save_stdout
if self.save_stderr is not None:
sys.stderr = self.save_stderr
self.temp.seek(0)
content = self.temp.read()
self.temp.close()
return content
def run_impl(form, _self):
'''
Execute the pipeline (prep, apbs, surface vis)
'''
selection = form.input_sele.currentText()
prep_name = ''
map_name = ''
ramp_name = ''
group_name = ['']
def get_name(lineinput):
name = lineinput.text()
if not name:
name = _self.get_unused_name(lineinput.placeholderText())
return name
def do_group(name):
if form.check_no_group.isChecked():
return
if not group_name[0]:
group_name[0] = _self.get_unused_name('run', 1)
_self.group(group_name[0], name)
if form.do_prepare.isChecked():
from . import creating as pc
prep_name = get_name(form.prep_name)
method = form.prep_method.currentText()
warnings = ''
get_res_charge = lambda resn: ((-1) if resn in ('GLU', 'ASP') else
(1) if resn in ('ARG', 'LYS') else 0)
def get_cb_pseudocharge(resn, name):
if name != 'CB':
return 0
return get_res_charge(resn)
try:
state = _self.get_selection_state(selection)
except BaseException as e:
state = -1
if method == 'pdb2pqr':
warnings = pc.pdb2pqr_cli(prep_name, selection,
options=form.pdb2pqr_args.text(),
quiet=0,
state=state,
preserve=form.check_preserve.isChecked(),
fixrna=form.pdb2pqr_fixrna.isChecked(),
_proclist=form._proclist,
exe=form.pdb2pqr_exe.text())
if form.pdb2pqr_ignore_warnings.isChecked():
warnings = ''
elif method == 'protein_assign_charges_and_radii':
_self.create(prep_name, selection)
_self.util.protein_assign_charges_and_radii(prep_name)
elif method.startswith('prepwizard'):
pc.prepwizard(prep_name, selection,
options=form.prepwizard_args.text(),
_proclist=form._proclist,
quiet=0,
preserve=form.check_preserve.isChecked())
_self.alter(prep_name, 'elec_radius = vdw')
elif method.startswith('use formal'):
_self.create(prep_name, selection)
_self.alter(prep_name,
'(elec_radius, partial_charge) = (vdw, formal_charge)')
elif method.startswith('use CB'):
_self.create(prep_name, selection)
_self.alter(
prep_name,
'(elec_radius, partial_charge) = (vdw, getpc(resn, name))',
space={'getpc': get_cb_pseudocharge})
elif method.startswith('use CA'):
_self.create(prep_name, '(%s) & name CA' % (selection))
_self.alter(
prep_name,
'(elec_radius, vdw, partial_charge) = (3.0, 3.0, getpc(resn))',
space={'getpc': get_res_charge})
elif method.startswith('use vdw'):
_self.create(prep_name, selection)
_self.alter(prep_name, 'elec_radius = vdw')
else:
raise ValueError('unknown method: ' + method)
selection = prep_name
form.input_sele.addItem(prep_name)
do_group(prep_name)
if warnings:
@form._callInMainThread
def result():
msgbox = QMessageBox(QMessageBox.Icon.Question, 'Continue?',
method + ' emmitted warnings, do you want to continue?',
QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No , form._dialog)
msgbox.setDetailedText(warnings)
return msgbox.exec()
if result == QMessageBox.StandardButton.No:
raise SilentAbort
if form.do_apbs.isChecked():
map_name = get_name(form.apbs_map)
template = form.apbs_template.toPlainText()
electrostatics.map_new_apbs(
map_name,
selection,
grid=form.apbs_grid.value(),
focus=form.focus_sele.text(),
quiet=0,
preserve=form.check_preserve.isChecked(),
exe=form.apbs_exe.text(),
_template=template,
_proclist=form._proclist,
assign=0)
form.surf_map.addItem(map_name)
do_group(map_name)
if form.do_surface.isChecked():
if not map_name:
map_name = form.surf_map.currentText()
if not map_name:
raise ValueError('missing map')
if not prep_name:
prep_name = _self.get_object_list(selection)[0]
ramp_name = get_name(form.surf_ramp)
v = form.surf_range.value()
_self.ramp_new(ramp_name, map_name, [-v, 0, v])
do_group(ramp_name)
sas = 'Accessible' in form.surf_type.currentText()
_self.set('surface_ramp_above_mode', not sas, prep_name)
_self.set('surface_solvent', sas, prep_name)
_self.set('surface_color', ramp_name, prep_name)
_self.show('surface', selection)
def dialog(_self=None):
if _self is None:
from pymol import cmd as _self
dialog = QtWidgets.QDialog()
uifile = os.path.join(os.path.dirname(__file__), 'apbs.ui')
form = loadUi(uifile, dialog)
form._dialog = dialog
form._proclist = []
def set_apbs_in(contents):
form.apbs_template.setPlainText(contents.strip())
# hide options widgets
form.optarea_prep.setVisible(False)
form.optarea_apbs.setVisible(False)
form.optarea_surf.setVisible(False)
form.optarea_other.setVisible(False)
# pre-fill form with likely data
names = _self.get_object_list()
names += ['(' + n + ')' for n in _self.get_names('public_selections')]
if names:
form.input_sele.clear()
form.input_sele.addItems([('polymer & ' + name)
if _self.count_atoms('polymer & ' + name) > 0
else name for name in names])
form.surf_map.addItems(_self.get_names_of_type('object:map'))
set_apbs_in(electrostatics.template_apbs_in)
# executables
from shutil import which
form.apbs_exe.setText(electrostatics.find_apbs_exe() or 'apbs')
form.pdb2pqr_exe.setText(
which('pdb2pqr') or
which('pdb2pqr30') or
# acellera::htmd-pdb2pqr provides pdb2pqr_cli
which('pdb2pqr_cli') or 'pdb2pqr')
# for async panels
form._callInMainThread = MainThreadCaller()
run_impl_async = AsyncFunc(run_impl)
# "Run" button callback
def run():
form.tabWidget.setEnabled(False)
form.button_ok.clicked.disconnect()
form.button_ok.clicked.connect(abort)
form.button_ok.setText('Abort')
form._capture = StdOutCapture()
# detach from main thread
run_impl_async(form, _self)
# "Run" button "finally" actions (main thread)
@run_impl_async.finished.connect
def run_finally(args):
_, exception = args
form._proclist[:] = []
stdout = form._capture.release()
print(stdout)
form.button_ok.setText('Run')
form.button_ok.clicked.disconnect()
form.button_ok.clicked.connect(run)
form.button_ok.setEnabled(True)
form.tabWidget.setEnabled(True)
if exception is not None:
handle_exception(exception, stdout)
return
quit_msg = "Finished with Success. Close the APBS dialog?"
if QMessageBox.StandardButton.Yes == QMessageBox.question(
form._dialog, 'Finished', quit_msg, QMessageBox.StandardButton.Yes,
QMessageBox.StandardButton.No):
form._dialog.close()
def handle_exception(e, stdout):
if isinstance(e, SilentAbort):
return
msg = str(e) or 'unknown error'
msgbox = QMessageBox(QMessageBox.Icon.Critical, 'Error', msg, QMessageBox.StandardButton.Close, form._dialog)
if stdout.strip():
msgbox.setDetailedText(stdout)
msgbox.exec()
# "Abort" button callback
def abort():
form.button_ok.setEnabled(False)
while form._proclist:
p = form._proclist.pop()
try:
p.terminate()
p.returncode = -15 # SIGTERM
except OSError as e:
print(e)
# selection checker
check_sele_timer = QtCore.QTimer()
check_sele_timer.setSingleShot(True)
# grid auto-value
form.apbs_grid_userchanged = False
form.apbs_grid.setStyleSheet('background: #ff6')
@form.apbs_grid.editingFinished.connect
def _():
form.apbs_grid_userchanged = True
form.apbs_grid.setStyleSheet('')
@check_sele_timer.timeout.connect
def _():
has_props = ['no', 'no']
def callback(partial_charge, elec_radius):
if partial_charge: has_props[0] = 'YES'
if elec_radius > 0: has_props[1] = 'YES'
n = _self.iterate(
form.input_sele.currentText(),
'callback(partial_charge, elec_radius)',
space={'callback': callback})
# grid auto-value (keep map size in the order of 200x200x200)
if n > 1 and not form.apbs_grid_userchanged:
e = _self.get_extent(form.input_sele.currentText())
volume = (e[1][0] - e[0][0]) * (e[1][1] - e[0][1]) * (e[1][2] - e[0][2])
grid = max(0.5, volume ** 0.333 / 200.0)
form.apbs_grid.setValue(grid)
if n < 1:
label = 'Selection is invalid'
color = '#f66'
elif has_props == ['YES', 'YES']:
label = 'No preparation necessary, selection has charges and radii'
form.do_prepare.setChecked(False)
color = '#6f6'
else:
label = 'Selection needs preparation (partial_charge: %s, elec_radius: %s)' % tuple(
has_props)
form.do_prepare.setChecked(True)
color = '#fc6'
form.label_sele_has.setText(label)
form.label_sele_has.setStyleSheet('background: %s; padding: 5' % color)
check_sele_timer.start(0)
@form.apbs_exe_browse.clicked.connect
def _():
fnames = getOpenFileNames(None, filter='apbs (apbs*);;All Files (*)')[0]
if fnames:
form.apbs_exe.setText(fnames[0])
@form.pdb2pqr_exe_browse.clicked.connect
def _():
fnames = getOpenFileNames(None, filter='pdb2pqr (pdb2pqr*);;All Files (*)')[0]
if fnames:
form.pdb2pqr_exe.setText(fnames[0])
# hook up events
form.input_sele.currentIndexChanged.connect(
lambda: check_sele_timer.start(0))
form.input_sele.editTextChanged.connect(
lambda: check_sele_timer.start(1000))
form.button_ok.clicked.connect(run)
# "Register" opens a web browser
@form.button_register.clicked.connect
def _():
import webbrowser
webbrowser.open("http://www.poissonboltzmann.org/")
@form.button_load.clicked.connect
def _():
fnames = getOpenFileNames(
None, filter='APBS Input (*.in);;All Files (*)')[0]
if fnames:
contents = load_apbs_in(form, fnames[0])
set_apbs_in(contents)
@form.button_reset.clicked.connect
def _():
set_apbs_in(electrostatics.template_apbs_in)
form._dialog.show()
form._dialog.resize(500, 600)
def load_apbs_in(form, filename, contents=''):
import shlex
from pymol import cmd, importing
wdir = os.path.dirname(filename)
if not contents:
contents = cmd.file_read(filename)
if not isinstance(contents, str):
contents = contents.decode()
sectionkeys = ('read', 'elec', 'apolar', 'print')
section = ''
insert_write_pot = True
lines = []
for line in contents.splitlines():
a = shlex.split(line)
key = a[0].lower() if a else ''
if not section:
if key in sectionkeys:
section = key
elif key == 'end':
if section == 'elec' and insert_write_pot:
lines.append('write pot dx "{mapfile}"')
section = ''
elif section == 'read':
if len(a) > 2 and key in ('charge', 'kappa', 'mol', 'parm', 'pot'):
filename = os.path.join(wdir, a[2])
if os.path.exists(filename):
format = a[1].lower()
if key == 'mol' and format in ('pqr', 'pdb'):
# load into PyMOL and update selection dropdown
oname = importing.filename_to_objectname(a[2])
oname = cmd.get_unused_name(oname, 0)
cmd.load(filename, oname, format=format)
form.input_sele.addItem(oname)
form.input_sele.setEditText(oname)
# absolute path in input file
a[2] = '"' + filename + '"'
line = ' '.join(a)
else:
QMessageBox.warning(
form._dialog, "Warning",
f'Warning: File "{filename}" does not exist')
elif section == 'elec':
if key == 'write':
if a[1:4] == ['pot', 'dx', "{mapfile}"]:
insert_write_pot = False
lines.append(line)
return '\n'.join(lines)
def __init_plugin__(app=None):
from pymol.plugins import addmenuitemqt as addmenuitem
addmenuitem('APBS Electrostatics', dialog)