import schrodinger
from schrodinger import project
from schrodinger.application.jaguar import input as jaginput
from schrodinger.Qt import QtCore
from schrodinger.Qt.QtCore import Qt
from schrodinger.ui.qt import pop_up_widgets
from schrodinger.ui.qt.standard.colors import LightModeColors
from ... import input_tab_widgets_pka
from ... import utils as gui_utils
from . import base_widgets
from . import charge_selector
from .base_widgets import ATOMS_ROLE
from .base_widgets import SORT_ROLE
maestro = schrodinger.get_maestro()
CHARGE_WEIGHTS_ROLE = Qt.UserRole + 100
MARKER_SETTINGS = {"color": "green", "alt_color": "yellow"}
[docs]class ChargeConstraintsColumns(object):
HEADERS = ["Atom", "ID", "Entry Title", "Charge"]
NUM_COLS = len(HEADERS)
ATOMS, ID, TITLE, CHARGE = list(range(NUM_COLS))
[docs]class ChargeConstraintRow(base_widgets.SubTabRow):
"""
Data about a charge constraint setting
:cvar DEFAULT_WEIGHT: default atom weight in charge constraint.
:vartype DEFAULT_WEIGHT: float
"""
DEFAULT_WEIGHT = 1.0
[docs] def __init__(self, entry_id=None, title=None, charge=None, weights=None):
super(ChargeConstraintRow, self).__init__(entry_id, title)
self.atom_names = None
self.atom_nums = None
self.name_to_num = None
self.charge = charge
self.weights = weights
if self.weights:
self.updateAtomNames(list(weights))
[docs] def copy(self):
"""
Create a new row object that is a copy of this row
:rtype: `ChargeConstraintRow`
:return: A row item that is a copy of this row.
"""
eid = self.entry_id
title = self.title
charge = self.charge
weights = self.weights.copy()
return ChargeConstraintRow(entry_id=eid,
title=title,
charge=charge,
weights=weights)
[docs] def updateAtomNames(self, names):
"""
Change the atom names for this constraint
:param names: The list of atom names
:type names: list
"""
self.atom_names = names
self._updateChargeWeights(names)
self._updateAtomNumbers(names)
def _updateChargeWeights(self, names):
"""
Updates the atom weights dictionary when the atom names change.
:param names: The list of atom names
:type names: list
"""
if names is None:
self.weights = None
elif self.weights is not None:
self.weights = {k: v for k, v in self.weights.items() if k in names}
for atom in names:
if atom not in self.weights:
self.weights[atom] = self.DEFAULT_WEIGHT
else:
self.weights = {a: self.DEFAULT_WEIGHT for a in names}
def _updateAtomNumbers(self, names):
"""
Updates the atom index dictionary when the atom names change.
:param names: The list of atom names
:type names: list
"""
if names:
proj = maestro.project_table_get()
struc = proj[self.entry_id].getStructure()
jaginput.apply_jaguar_atom_naming(struc)
self.name_to_num = {
atom.name: atom.index
for atom in struc.atom
if atom.name in names
}
self.atom_nums = list(self.name_to_num.values())
else:
self.name_to_num = None
self.atom_nums = None
[docs] def weightsByNum(self):
"""
Get a dictionary of {atom index: atom weight}. (As opposed to
`weights`, which is {atom name: atom weight})
:return: Dictionary of {atom index: atom weight}
:rtype: dict
"""
return {
self.name_to_num[name]: self.weights[name]
for name in self.atom_names
}
[docs]class ChargeConstraintsTableView(base_widgets.SubTabTableView):
"""
The view for the charge constraints.
"""
COLUMN = ChargeConstraintsColumns
MARKER_SETTINGS = MARKER_SETTINGS
addJaguarMarker = QtCore.pyqtSignal(list, dict)
removeJaguarMarker = QtCore.pyqtSignal(list)
[docs] def __init__(self, parent=None):
super(ChargeConstraintsTableView, self).__init__(parent)
self._setDelegates(parent)
def _setDelegates(self, parent):
"""
Set all delegates for the charge constraints view.
:param parent: Qt parent widget
:type parent: `PyQt5.QtWidgets.QWidget`
"""
self._atom_delegate = AtomSelectionDelegate(parent)
self._atom_delegate.removeJaguarMarker.connect(
self.removeJaguarMarker.emit)
self._atom_delegate.addJaguarMarker.connect(self._emitAddMarker)
self.setItemDelegateForColumn(self.COLUMN.ATOMS, self._atom_delegate,
True)
self._charge_delegate = ChargeSelectorDelegate(parent)
self.setItemDelegateForColumn(self.COLUMN.CHARGE, self._charge_delegate)
def _setHighlightingForAtoms(self, atoms, highlight):
for cur_atom in atoms:
self.setMarkerHighlighting.emit([cur_atom], highlight)
[docs] def dataChanged(self, topleft, bottomright, role):
"""
If the atoms data in a selected cell changes, make sure that the newly
created markers are highlighted
"""
if not topleft.column() <= self.COLUMN.ATOMS <= bottomright.column():
return
changed_rows = list(range(topleft.row(), bottomright.row() + 1))
selected_rows = {index.row() for index in self.selectedIndexes()}
if selected_rows.intersection(changed_rows):
selection = self.selectionModel().selection()
self._emitForSelection(selection, True)
def _emitAddMarker(self, atoms, index):
settings = self.MARKER_SETTINGS
settings["highlight"] = index in self.selectedIndexes()
self.addJaguarMarker.emit(atoms, settings)
[docs]class ChargeConstraintsModel(base_widgets.SubTabModel):
COLUMN = ChargeConstraintsColumns
UNEDITABLE = (COLUMN.ID, COLUMN.TITLE)
ROW_CLASS = ChargeConstraintRow
MARKER_SETTINGS = MARKER_SETTINGS
def _otherData(self, col, proj_row, role):
# See parent class for documentation
if role == CHARGE_WEIGHTS_ROLE:
return proj_row.weights
def _displayAndSortData(self, col, charge_row, role):
# See parent class for documentation
if col == self.COLUMN.CHARGE:
if role == Qt.DisplayRole:
if charge_row.charge is not None:
# Use Python's stringifying rather than Qt's so that we
# always get a decimal point and at least one digit after it
# even for whole numbers
return str(charge_row.charge)
elif role == SORT_ROLE:
return charge_row.charge
else:
return super(ChargeConstraintsModel,
self)._displayAndSortData(col, charge_row, role)
def _toolTipData(self, col, row_data):
# See parent class for documentation
if col == self.COLUMN.CHARGE:
if row_data.weights:
# weights can be empty if the user requests a tooltip before
# finishing the initial atom selection
return self._weightsToolTip(row_data.weights)
def _weightsToolTip(self, weights):
"""
Return tool tip text for the Charge column that displays atom weights
:param weights: A dictionary of {atom name: atom weight}
:type weights: dict
:return: The tool tip text
:rtype: str
"""
pad_right = "style='padding-right: 10px;'"
msg = ("<table><tr><th align=left %s>Atom</th>"
"<th align=left>Weight</th></tr>" % pad_right)
sorted_weights = sorted(
weights.items(), key=lambda x: gui_utils.atom_name_sort_key(x[0]))
for atom, weight in sorted_weights:
msg += ("<tr><td %s>%s</td><td>%0.2f</td></tr>" %
(pad_right, atom, weight))
msg += "</table>"
return msg
def _setData(self, col, row_data, value, role, row_num):
# See parent class and Qt documentation
if role == Qt.EditRole:
if col == self.COLUMN.ATOMS:
self._setAtomsData(row_data, value, row_num)
elif col == self.COLUMN.ID:
row_data.entry_id = value
elif col == self.COLUMN.TITLE:
row_data.title = value
elif col == self.COLUMN.CHARGE:
row_data.charge = value
elif role == CHARGE_WEIGHTS_ROLE:
row_data.weights = value
return True
def _setAtomsData(self, row_data, value, row_num):
"""
Set data for the Atoms column
:param row_data: The ROW_CLASS instance to modify
:type row_data: ROW_CLASS
:param value: The value to set
:param value: object
:param row_num: The row number
:type row_num: int
"""
self.removeJaguarMarkerForRow(row_data)
row_data.updateAtomNames(value)
self.addJaguarMarkerForRow(row_data)
if row_data.charge is None:
charge = (atom.formal_charge for atom in row_data.getAtoms())
row_data.charge = float(sum(charge))
charge_index = self.index(row_num, self.COLUMN.CHARGE)
self.dataChanged.emit(charge_index, charge_index)
[docs] def addJaguarMarkerForRow(self, row):
# See parent class and Qt documentation
atoms = row.getAtoms()
for cur_atom in atoms:
self.addJaguarMarker.emit([cur_atom], self.MARKER_SETTINGS)
[docs] def removeJaguarMarkerForRow(self, row):
# See parent class and Qt documentation
atoms = row.getAtoms()
for cur_atom in atoms:
self.removeJaguarMarker.emit([cur_atom])
[docs] def copyRow(self, row):
"""
Copy row into a new row in this model
:type row: `ChargeConstraintRow`
:param row: The row item containing the data to copy
"""
eid = row.entry_id
title = row.title
charge = row.charge
weights = row.weights.copy()
self.addRow(entry_id=eid, title=title, charge=charge, weights=weights)
[docs]class AtomSelectionDelegate(input_tab_widgets_pka.AtomSelectionDelegate):
"""
A delegate for selecting atoms for the charge constraint.
:note: We currently don't allow the user to type in atom names. If the user
doesn't click on an atom in the workspace, then we have no way to determine
which entry the constraint refers to. Invalid atoms names (i.e. those not
in the form <element><index>) would also cause issues for the weights pop-
up. Additionally, we'd have to add atom name validation to the model in
order to color cells with errors and to give a tool tip explaining the
problem, as is done for pKa atoms. We'd also have to implement validate()
for the sub-tab. It would be possible to solve these issues, but for now
we simply require that the user select atoms from the workspace.
:ivar addJaguarMarker: A signal emitted when a workspace marker should be added.
Emitted with:
- The list of atoms to add the marker for (list)
- The index that the atom is being marked for. (Used to determine whether
the marker should be highlighted or not.) (`QtCore.QModelIndex`)
:vartype addJaguarMarker: `PyQt5.QtCore.pyqtSignal`
:ivar removeJaguarMarker: A signal emitted when a workspace marker should be
removed. Emitted with:
- The list of atoms to remove the marker for (list)
:vartype removeJaguarMarker: `PyQt5.QtCore.pyqtSignal`
"""
MAESTRO_STATUS_MESSAGE = ("Pick an atom to be included in the charge "
"constraint")
TOOL_TIP_INSTRUCTIONS = ("Click an atoms in the workspace to\n"
"set them as a charge constraint atom.\n"
"Click again to remove.")
TOOL_TIP_WRONG_EID = ("The atom you selected is not\n"
"part of this structure.")
# Don't display "Double-click to edit...", since we'll be in the editor
# whenever the cell is empty
DEFAULT_DATA = {}
addJaguarMarker = QtCore.pyqtSignal(list, QtCore.QModelIndex)
removeJaguarMarker = QtCore.pyqtSignal(list)
[docs] def __init__(self, parent):
super(AtomSelectionDelegate, self).__init__(parent)
self.closeEditor.connect(self.clearRowIfEmpty)
def _ensureEntryIncluded(self, index):
# See parent class for documentation
eid, eid_index = self._getEid(index)
if not eid:
# The entry ID won't be set for the row if the user hasn't selected
# an initial atom yet
return
proj = maestro.project_table_get()
proj_row = proj[eid]
if proj_row.in_workspace == project.NOT_IN_WORKSPACE:
proj_row.includeOnly()
def _setEditorAtom(self, ws_atom_num, editor, index):
"""
Respond to the user clicking an atom in the workspace. If the atom does
not belong to this structure, display a warning. Otherwise, enter the
Jaguar-ified atom name into the line edit and display a marker on the
atom. If atom is already included in the list of charge constraint
atoms we remove it from the list.
:param ws_atom_num: The atom number (relative to the workspace
structure) that was selected.
:type ws_atom_num: int
:param editor: The line edit to enter the atom name into
:type editor: `PyQt5.QtWidgets.QLineEdit`
:param index: The index of the cell being edited
:type index: `PyQt5.QtCore.QModelIndex`
"""
editor_eid, eid_index = self._getEid(index)
atom_name, atom_num, eid, title = gui_utils.get_atom_info(ws_atom_num)
model = index.model()
# If this is the first atom picked set entry id for the row to be
# equal to the atom's entry id.
if editor_eid is None:
editor_eid = eid
model.setData(eid_index, eid)
title_index = model.index(index.row(), model.COLUMN.TITLE)
model.setData(title_index, title)
if eid != editor_eid:
self._showToolTip(editor, self.TOOL_TIP_WRONG_EID)
else:
if editor.new_atom_marked:
self._clearEditorMarkers(index, editor)
else:
self._clearModelMarkers(index)
self._updateEditorText(editor, atom_name)
text = editor.text()
atoms = self._getAtoms(eid, text)
self._markAtoms(atoms, index)
editor.new_atom_marked = True
name_index = model.index(index.row(), model.COLUMN.ATOMS)
model.setData(name_index, [atom_name])
def _getEid(self, index):
"""
Get the entry id associated with a given model index
:param index: The model index to get the entry id for
:type index: `QtCore.QModelIndex`
:return: A tuple of:
- entry id
- the model index of the entry id cell
:rtype: tuple
"""
model = index.model()
eid_index = model.index(index.row(), model.COLUMN.ID)
eid = eid_index.data()
return eid, eid_index
def _getAtoms(self, eid, text):
"""
Get the atom objects for the specified atoms
:param eid: The entry id to get atoms from
:type eid: str
:param text: A comma-separated list of atoms
:type text: str
:return: A list of `schrodinger.structure._StructureAtom` objects
:rtype: list
"""
proj = maestro.project_table_get()
struc = proj[eid].getStructure()
jaginput.apply_jaguar_atom_naming(struc)
return self._getAtomSts(struc, text)
def _clearModelMarkers(self, index):
"""
Remove all markers associated with the specified model index
:param index: The model index
:type index: `QtCore.QModelIndex`
"""
atoms = index.data(ATOMS_ROLE)
if atoms:
for cur_atom in atoms:
self.removeJaguarMarker.emit([cur_atom])
def _clearEditorMarkers(self, index, editor):
"""
Remove all markers associated with the specified editor
:param index: The model index that the editor is editing
:type index: `QtCore.QModelIndex`
:param editor: The editor
:type editor: `PyQt5.QtWidgets.QLineEdit`
"""
eid, eid_index = self._getEid(index)
text = editor.text()
atoms = self._getAtoms(eid, text)
if atoms:
for cur_atom in atoms:
self.removeJaguarMarker.emit([cur_atom])
def _markAtoms(self, atoms, index):
"""
Mark the specified atoms
:param index: The model index that the editor is editing
:type index: `QtCore.QModelIndex`
"""
if atoms:
for cur_atom in atoms:
self.addJaguarMarker.emit([cur_atom], index)
[docs] def clearRowIfEmpty(self, editor, hint=None):
"""
When the editor closes, erase a row if it does not contain any atom
names. Without any atom names, the row is meaningless. By clearing the
rows immediately, we avoid the need to later warn the user about them.
:param editor: The recently-closed editor
:type editor: `PyQt5.QtWidgets.QLineEdit`
:param hint: Ignored, but present for compatibility with the Qt signal
"""
index = editor.index
atoms = index.data()
if not atoms:
# We don't want to remove the row until we've finished calling all
# of the slots connected to closeEditor. Otherwise, they'll be
# called with a blank widget instead of the editor. To avoid this,
# we use a timer to delay the removeRow call.
func = lambda: index.model().removeRow(index.row())
QtCore.QTimer.singleShot(0, func)
[docs] def createEditor(self, parent, option, index):
"""
Create an editor as in the super-class, but mark it as read-only. See
the :note in the class documentation for an explanation of why.
See Qt documentation for an explanation of arguments and return values
"""
editor = super(AtomSelectionDelegate,
self).createEditor(parent, option, index)
editor.setReadOnly(True)
# At least on Windows, read-only line edits get a gray background that
# makes them look disabled. That's not what we want here.
editor.setStyleSheet(
f"QLineEdit{{background: {LightModeColors.STANDARD_BACKGROUND};}}")
self._startPicking(editor, index, self.MAESTRO_STATUS_MESSAGE)
return editor
def _resetMarker(self, editor, hint=None):
# See parent class for documentation
if not editor.new_atom_marked:
return
index = editor.index
self._clearEditorMarkers(index, editor)
atoms = index.data(ATOMS_ROLE)
self._markAtoms(atoms, index)
editor.new_atom_marked = False
[docs] def setModelData(self, editor, model, index):
"""
Before setting the model data, clear any workspace markers we have
created and restore the markers to the model values. This avoids having
the model trying to clear a marker that no longer exists, which would
raise an exception.
See parent class and PyQt documentation for an explanation of arguments.
"""
self._resetMarker(editor)
super(AtomSelectionDelegate, self).setModelData(editor, model, index)
[docs]class ChargeSelectorDelegate(pop_up_widgets.PopUpDelegate):
"""
A delegate for selecting charge constraint atom weights.
"""
[docs] def __init__(self, parent):
super(ChargeSelectorDelegate, self).__init__(parent, None, True)
def _createEditor(self, parent, option, index):
# See parent class and Qt documentation
editor = charge_selector.ChargeSelectorLineEdit(parent)
return editor
[docs] def setEditorData(self, editor, index):
# See parent class and Qt documentation
charge = index.data()
weights = index.data(role=CHARGE_WEIGHTS_ROLE)
editor.setText(str(charge))
editor.setWeights(weights)
[docs] def setModelData(self, editor, model, index):
# See parent class and Qt documentation
if editor.hasAcceptableInput():
value = float(editor.text())
model.setData(index, value)
weights = editor.getWeights()
model.setData(index, weights, role=CHARGE_WEIGHTS_ROLE)