Source code for schrodinger.ui.qt.atom_weights_table_widget
from typing import List
from typing import Optional
import schrodinger
from schrodinger import comparison
from schrodinger import structure
from schrodinger.graphics3d import sphere
from schrodinger.models import mappers
from schrodinger.models import parameters
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.ui.qt import atom_weights_table_ui
from schrodinger.ui.qt import basewidgets
from schrodinger.ui.qt.mapperwidgets import MappableComboBox
from schrodinger.ui.qt.mapperwidgets import plptable
from schrodinger.ui.qt.mapperwidgets.plptable import FieldColumn
from schrodinger.ui.qt.standard import colors
maestro = schrodinger.get_maestro()
INITIAL_WEIGHT = 0.5
COMBO_WIDTH = 100
INFO_TOOLTIP = ('Assign weights from 0 to 1 to indicate'
'\ntheir relative importance')
SORT_ROLE = QtCore.Qt.UserRole + 1
[docs]class AtomWeight(parameters.CompoundParam):
"""
Represent weight parameters for a single atom.
"""
atom_number: int
element: str
pdb_name: str
weight: float
[docs]class AtomWeightsSpec(plptable.TableSpec):
"""
Table spec for the custom weights.
"""
atom_number = FieldColumn(AtomWeight.atom_number, title='Atom #')
element = FieldColumn(AtomWeight.element, title='Element')
pdb_name = FieldColumn(AtomWeight.pdb_name, title='PDB Name')
weight = FieldColumn(AtomWeight.weight, title='Weight')
[docs] @atom_number.data_method(SORT_ROLE)
@element.data_method(SORT_ROLE)
@pdb_name.data_method(SORT_ROLE)
@weight.data_method(SORT_ROLE)
def sort_role(self, field):
return field
[docs] @weight.data_method(role=QtCore.Qt.FontRole)
def weight_font(self, weight):
font = QtGui.QFont()
if weight == INITIAL_WEIGHT:
font.setItalic(True)
return font
[docs] @weight.data_method(role=QtCore.Qt.ForegroundRole)
def weight_color(self, weight):
if weight == INITIAL_WEIGHT:
color = QtGui.QColor(colors.NativeColors.GRAY)
else:
color = QtGui.QColor(colors.NativeColors.BLACK)
brush = QtGui.QBrush(color)
return brush
[docs]class AtomWeightsTableWidget(plptable.PLPTableWidget):
"""
Table to display atom weights data.
"""
[docs] def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.setSpec(AtomWeightsSpec())
sort_proxy_model = QtCore.QSortFilterProxyModel()
sort_proxy_model.setSortRole(SORT_ROLE)
self.addProxy(sort_proxy_model)
self._setUpTableView()
def _setUpTableView(self):
self.view.setSortingEnabled(True)
self.view.setAlternatingRowColors(True)
self.view.horizontalHeader().setStretchLastSection(True)
self.view.verticalHeader().hide()
[docs]class AtomWeightsSettings(parameters.CompoundParam):
"""
Class for storing Custom Weight Options. Contains all logic for updating
the table data and spheres.
"""
allow_custom_weights: bool = False
use_custom_weights: bool = False
default_weight: float = INITIAL_WEIGHT
custom_weight: float = INITIAL_WEIGHT
atom_weights: List[AtomWeight]
selected_weights: List[AtomWeight]
_template_sts: List[structure.Structure]
_weight_property: str = None
template_stsChanged = QtCore.pyqtSignal(object)
[docs] def initConcrete(self):
self._prev_default_weight = self.default_weight
self._template_stsChanged.connect(self.template_stsChanged)
@property
def template_sts(self) -> List[structure.Structure]:
return self._template_sts
@template_sts.setter
def template_sts(self, template_sts: List[structure.Structure]):
"""
Set template structures as `template_sts`.
:param template_sts: Structures to apply custom atom weights to.
"""
self._template_sts = template_sts
self.updateAtomWeights()
@property
def template_st(self) -> Optional[structure.Structure]:
"""
Return the first template structure or None.
"""
if not self.template_sts:
return
return self.template_sts[0]
[docs] def updateAtomWeights(self):
"""
Clear the table model and then add a row for each atom in the current
template_st.
"""
self._prev_default_weight = self.default_weight
self.atom_weights.clear()
self.selected_weights.clear()
if not self.template_st or not self.use_custom_weights:
return
for atom in self.template_st.atom:
weight = self.default_weight
if self._weight_property:
weight = atom.property.get(self._weight_property,
self.default_weight)
self.atom_weights.append(
AtomWeight(atom_number=atom.index,
element=atom.element,
pdb_name=atom.pdbname,
weight=weight))
[docs] def applyDefaultWeight(self):
"""
Apply the default weights to all rows that haven't been assigned a
custom weight, and update all spheres.
"""
for atom in self.atom_weights:
if atom.weight == self._prev_default_weight:
atom.weight = self.default_weight
self._prev_default_weight = self.default_weight
[docs] def applyCustomWeight(self):
"""
Apply the custom weight to currently selected rows and update all
spheres.
"""
for atom in self.selected_weights:
atom.weight = self.custom_weight
[docs] def addWeightPropsToTemplateStructures(self, weight_property: str):
"""
Add `weight_property` as a property to the atoms in all template
structures.
:param weight_property: the name of the associated weight property
"""
if not self.use_custom_weights:
return
if self.atom_weights:
self._weight_property = weight_property
for weight in self.atom_weights:
try:
atom = self.template_st.atom[weight.atom_number]
except IndexError:
msg = ('ERROR: WEIGHTS: Template structure has'
'no atom %i in it.' % weight.atom_number)
print(msg)
continue
atom.property[weight_property] = weight.weight
else:
for atom in self.template_st.atom:
atom.property[weight_property] = self.default_weight
[docs]class WeightSphereManager(mappers.MapperMixin):
"""
Manages 3D spheres representing atom weights of each individual atoms in
AtomWeightsSettings.
"""
model_class = AtomWeightsSettings
OPACITY_DEFAULT = 0.4
OPACITY_SELECTED = 1.0
COLOR_DEFAULT = (0.5, 0.5, 1.0)
COLOR_SELECTED = (1.0, 1.0, 0.0)
RESOLUTION_DEFAULT = 50
[docs] def __init__(self):
super().__init__()
self._sphere_group = sphere.SphereGroup()
self._setupMapperMixin()
[docs] def defineMappings(self):
M = self.model_class
if not maestro:
return []
return [
(self.updateSpheres, M.atom_weights),
(self.updateSpheres, M.selected_weights),
(self.updateSpheres, M.use_custom_weights),
(self.updateSpheres, M.allow_custom_weights),
]
[docs] def getSignalsAndSlots(self, model):
return [
(model.template_stsChanged, self.updateSpheres),
]
[docs] def updateSpheres(self):
"""
Clear all spheres and, if necessary, create a sphere for each row in the
atom weights table with appropriate colors.
"""
if not maestro:
return
self.clearSpheres()
template_in_ws = False
if self.model.template_st:
template_in_ws = comparison.are_same_geometry(
maestro.workspace_get(), self.model.template_st)
if all((self.model.allow_custom_weights, self.model.use_custom_weights,
template_in_ws)):
self._createSpheres()
[docs] def clearSpheres(self):
"""
Delete all spheres from the sphere group. Clears all spheres in the WS.
"""
self._sphere_group.clear()
def _createSpheres(self):
"""
Create a new sphere for each atom in the table with appropriate size
and colors.
"""
template_st = self.model.template_st
for atom_weight in self.model.atom_weights:
atom = template_st.atom[atom_weight.atom_number]
radius = 0.5 * atom_weight.weight
selected = False
if atom_weight in self.model.selected_weights:
selected = True
self._createSphere(x=atom.x,
y=atom.y,
z=atom.z,
radius=radius,
selected=selected)
def _createSphere(self, x: float, y: float, z: float, radius: int,
selected: bool):
"""
Create and add a new sphere to the sphere group.
:param x: sphere x coordinate
:param y: sphere y coordinate
:param z: sphere z coordinate
:param radius: sphere radius
:param selected: whether or not the sphere is currently selected.
"""
color = self.COLOR_DEFAULT
opacity = self.OPACITY_DEFAULT
if selected:
color = self.COLOR_SELECTED
opacity = self.OPACITY_SELECTED
sphere_marker = sphere.MaestroSphere(x=x,
y=y,
z=z,
radius=radius,
resolution=self.RESOLUTION_DEFAULT,
color=color,
opacity=opacity)
self._sphere_group.add(sphere_marker)
class _StyledDefaultWeightCombo(MappableComboBox):
"""
An MappableComboBox which makes the initial weight text italicized and gray.
"""
def __init__(self, parent):
super().__init__(parent)
# Same width as other combo defined in atom_weights_table_ui
self.setFixedWidth(COMBO_WIDTH)
def paintEvent(self, event):
"""
Re-implement QComboBox.paintEvent to use italic font and gray
color for the text of the default atom weight.
"""
painter = QtWidgets.QStylePainter(self)
painter.setPen(self.palette().color(QtGui.QPalette.Text))
opt = QtWidgets.QStyleOptionComboBox()
self.initStyleOption(opt)
if self.currentText() == str(INITIAL_WEIGHT):
font = painter.font()
font.setItalic(True)
painter.setFont(font)
pen = QtGui.QPen(QtGui.QColor(colors.NativeColors.GRAY))
painter.setPen(pen)
# draw the combobox frame, focusrect, selected, etc.
painter.drawComplexControl(QtWidgets.QStyle.CC_ComboBox, opt)
# draw the icon and text
painter.drawControl(QtWidgets.QStyle.CE_ComboBoxLabel, opt)
[docs]class AtomWeightsWidget(mappers.MapperMixin, basewidgets.BaseWidget):
"""
A widget for the custom weights table. User can specify default and custom
weights for the atoms and then use the values to run their custom job with
an additional WEIGHT_PROPERTY.
Initially, set the template structures using setTemplateStructures() method.
To add weights to the template structures, call
addWeightPropsToTemplateStructures() with the WEIGHT_PROPERTY.
Atom weights can be visualized with sphere markers in the workspace by
instantiating this alongside `WeightSphereManager` and mapping both views
to the same model.
All the params of `AtomWeightsSettings` need to be synced up with the parent
widget for an ideal arrangement.
"""
model_class = AtomWeightsSettings
ui_module = atom_weights_table_ui
ATOM_WEIGHTS = [0, 0.25, 0.5, 0.75, 1]
[docs] def initSetUp(self):
super().initSetUp()
self.weights_table = AtomWeightsTableWidget(self)
self.default_weight_combo = _StyledDefaultWeightCombo(self)
self.ui.apply_default_weight_btn.clicked.connect(
self._applyDefaultWeight)
self.ui.apply_custom_weight_btn.clicked.connect(self._applyCustomWeight)
self.ui.get_ws_selection_btn.clicked.connect(self._getWSSelection)
self.ui.info_btn.setToolTip(INFO_TOOLTIP)
_atom_weights_items = {
str(atom_weight): atom_weight for atom_weight in self.ATOM_WEIGHTS
}
self.ui.custom_weight_combo.addItems(_atom_weights_items)
self.default_weight_combo.addItems(_atom_weights_items)
[docs] def initLayOut(self):
super().initLayOut()
self.ui.default_weight_layout.insertWidget(1, self.default_weight_combo)
self.ui.weights_table_layout.addWidget(self.weights_table)
[docs] def defineMappings(self):
M = self.model_class
return [
(self.weights_table, M.atom_weights),
(self.weights_table.selection_target, M.selected_weights),
(self.ui.customize_weights_cb, M.use_custom_weights),
(self.ui.custom_weight_combo, M.custom_weight),
(self.default_weight_combo, M.default_weight),
(self._onAllowCustomWeightsChanged, M.allow_custom_weights),
(self._onUseCustomWeightsChanged, M.use_custom_weights),
(self._onWeightSelectionChanged, M.selected_weights),
]
@property
def template_sts(self) -> List[structure.Structure]:
"""
Return the list of template structures currently held in the model.
"""
return self.model.template_sts
[docs] def setTemplateStructures(self, template_sts: List[structure.Structure]):
"""
Set the model's template structures as `template_sts` or [] if query
is an `*.ivol` file, and determine whether the get_ws_selection_btn should
be enabled.
:param template_sts: Structures to apply custom atom weights to.
"""
self.model.template_sts = template_sts
enabled = False
if self.model.template_st and maestro:
enabled = comparison.are_same_geometry(maestro.workspace_get(),
self.model.template_st)
self.ui.get_ws_selection_btn.setEnabled(enabled)
[docs] def addWeightPropsToTemplateStructures(self, weight_property):
"""
Add weight_property as a property to the atoms in all template
structures.
:type weight_property: str
:param weight_property: the name of the associated weight property
"""
self.model.addWeightPropsToTemplateStructures(weight_property)
def _getWSSelection(self):
"""
Clear existing table selection and select rows corresponding to current
WS atom selection.
"""
self.model.selected_weights.clear()
atom_idxs = maestro.selected_atoms_get()
selected_weights = [
atom_weight for atom_weight in self.model.atom_weights
if atom_weight.atom_number in atom_idxs
]
self.weights_table.setSelectedParams(selected_weights)
def _onAllowCustomWeightsChanged(self):
"""
Enable/disable the "Customize atom weights" check box.
"""
enabled = self.model.allow_custom_weights
self.ui.customize_weights_cb.setEnabled(enabled)
def _onUseCustomWeightsChanged(self):
"""
Update the table when this option is toggled.
"""
self.model.updateAtomWeights()
enabled = self.model.use_custom_weights
self.ui.custom_weights_box.setEnabled(enabled)
def _onWeightSelectionChanged(self):
"""
Determine whether to enable the apply_custom_weight button.
"""
enable = bool(self.model.selected_weights)
self.ui.apply_custom_weight_btn.setEnabled(enable)
def _applyDefaultWeight(self):
self.model.applyDefaultWeight()
# Won't immediately show changed weights without this
self.weights_table.repaint()
def _applyCustomWeight(self):
self.model.applyCustomWeight()
# Won't immediately show changed weights without this
self.weights_table.repaint()