Source code for schrodinger.ui.qt.atomselector

#Name: Atom Selector
#Command: pythonrun atomselector.panel
"""
PyQt version of the Maestro's ASL frame.

Designed to be used within Maestro, but possible to be used
outside of Maestro as well. Some options will disabled if run
outside of Maestro.

Copyright Schrodinger, LLC. All rights reserved.
"""

from collections import OrderedDict
from functools import partial

import schrodinger
from schrodinger.models import mappers
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.Qt.QtCore import Qt
from schrodinger.ui import maestro_ui
from schrodinger.ui import picking
from schrodinger.ui.qt import swidgets
from schrodinger.ui.qt import utils as qt_utils

from . import stylesheet

maestro = schrodinger.get_maestro()

try:
    from schrodinger.maestro import markers
except ImportError:
    markers = None

PICK_ATOMS = 0
PICK_RESIDUES = 1
PICK_CHAINS = 2
PICK_MOLECULES = 3
PICK_ENTRIES = 4
PLUS_ICON = ":/images/toolbuttons/asl_gizmo_plus.png"
PINK = (1.0, 0.8, 0.8)


[docs]class ASLPopupMenu(QtWidgets.QMenu): """ ASL menu intended to be shown off of a button. :cvar hidden: A signal emitted when menu is hidden. :vartype hidden: `QtCore.pyqtSignal` """ hidden = QtCore.pyqtSignal()
[docs] def hideEvent(self, event): """ Emit hidden() signal when popup is hidden """ super().hideEvent(event) self.hidden.emit()
[docs]class ASLItem(QtWidgets.QPushButton): """ Asl item to display text along with the plus button :cvar itemClicked: A signal emitted when this asl item is clicked. - Asl of the item :vartype itemClicked: `QtCore.pyqtSignal` """ itemClicked = QtCore.pyqtSignal(str)
[docs] def __init__(self, text, asl, parent): """ :param text: Display text of the asl item :type text: str :param asl: Asl of the item :type asl: str :param parent: Parent of the item :type parent: QtWidgets.QWidget """ super().__init__(QtGui.QIcon(PLUS_ICON), text, parent) self.setFlat(True) self.setStyleSheet("text-align:left") self.clicked.connect(partial(self.itemClicked.emit, asl))
[docs]class AtomSelector(QtWidgets.QGroupBox): """ :cvar aslModified: Emitted when a new atom is picked or the ASL is manually edited by the user. - New asl :vartype aslModified: `QtCore.pyqtSignal` :cvar atomSelectionDialogAboutToBeShown: A signal emitted when 'Atom Selection' dialog is about to be shown. :vartype atomSelectionDialogAboutToBeShown: `QtCore.pyqtSignal` :cvar atomSelectionDialogDismissed: A signal emitted when 'Atom Selection' dialog is dismissed. :vartype atomSelectionDialogDismissed: `QtCore.pyqtSignal` :cvar aslTextModified: Emitted when there is ANY change in the asl text field. Use this in case any action button needs to be enabled/disabled based on asl text field. - New asl :vartype aslTextModified: `QtCore.pyqtSignal` """ aslModified = QtCore.pyqtSignal(str) atomSelectionDialogAboutToBeShown = QtCore.pyqtSignal() atomSelectionDialogDismissed = QtCore.pyqtSignal() aslTextModified = QtCore.pyqtSignal(str) # Dictionary for display name to asl ASL_ITEMS = OrderedDict(( ('Workspace Selection', 'workspace_selection'), ('Displayed Atoms', 'displayed_atoms'), ('Protein', '(protein) and not ligand'), ('Protein Backbone', '(backbone) and not ligand'), ('Protein Side Chains', '(withinbonds 1 sidechain) and not ligand'), ('Protein Near Ligand', 'protein_near_ligand'), ('Ligands', 'ligand'), ('Nucleic Acids', 'nucleic_acids'), ('Waters', 'water'), ('Ions', 'ions'), ('Metal Atoms', 'metals'), ('Heavy Atoms', 'heavy_atoms'), ('Hydrogens-All', 'all_hydrogens'), ('Hydrogens-Nonpolar', 'non_polar_hydrogens'), ('Hydrogens-Nonpolar Ligand', 'non_polar_ligand_hydrogens'), ('Hydrogens-Polar', 'polar_hydrogens'), ('Membrane', 'membrane'))) # yapf: disable
[docs] def __init__(self, parent, label="", pick_text="Pick atom in Workspace", show_asl=True, show_all=True, show_select=True, show_previous=True, show_selection=True, append_mode=True, show_markers=False, show_pick=True, default_pick_mode=PICK_ATOMS, show_plus=False, selection_button_text="Selection"): """ AtomSelector requires <parent> argument, which should be a Qt Widget into which this frame should be added. This widget acts like any other PyQt widget. parent - The parent widget :type label: str :param label: Label to show above the widget :type pick_text: str :param pick_text: Text that will be displayed on the bottom of the main Maestro window when the pick button is checked. :type show_asl: bool :param show_asl: Whether to show the QLineEdit field for the ASL. :type show_all: bool :param show_all: Whether to show the "All" button. Clicking it will select the "all" ASL> :type show_select: bool :param show_select: Whether to show the "Select..." button. Clicking it would show an ASL selection dialog. :type show_previous: bool :param show_previous: Whether to show the "Previous" button. Clicking it would select the previous ASL that was used. :type show_selection: bool :param show_selection: Whether to show the "Selection" button. Clicking it would use ASL derived from the atoms selected in the Workspace. :type append_mode: bool :param append_mode: When a new atom is picked, whether to append to an existing ASL instead of replacing it. :type show_markers: bool :param show_markers: Whther to show the "Markers" checkbox. :type show_pick: bool :param show_pick: Whether to show the "Pick" checkbox along with the pick menu. :type default_pick_mode: int :param default_pick_mode: What the default pick mode should be. One of: PICK_ATOMS, PICK_RESIDUES, PICK_CHAINS, PICK_MOLECULES, PICK_ENTRIES. Default is PICK_ATOMS. :type show_plus: bool :param show_plus: Whther to show the "+" button. :type selection_button_text: str :param selection_button_text: This will be honored only if show_selection is True. By default it is "Selection", if show_asl & show_plus are True, then selection button text will be "Load Selection", otherwise given text will be set to the selection button. Usage example: atom_selector = AtomSelector(parent) layout.addWidget(atom_selector) To get the selected ASL, simply call <AtomSelector>.getAsl() """ QtWidgets.QGroupBox.__init__(self, parent) self.parent_window = self.parent() self.setTitle(label) self.setStyleSheet(stylesheet.ATOMSELECTORX_STYLESHEET) # 'Select...' & 'Selection' buttons can be shown only when maestro is available show_select = show_select and maestro show_selection = show_selection and maestro button_height = parent.fontMetrics().height() + 4 self.main_layout = QtWidgets.QVBoxLayout(self) self.main_layout.setContentsMargins(8, 8, 8, 8) self.main_layout.setSpacing(5) self._append_mode = append_mode self._previous_value = "" # ASL row: self.asl_layout = QtWidgets.QHBoxLayout() # Use MM_QLineEdit for inbuilt clear button self.asl_ef = maestro_ui.MM_QLineEdit(self) self.asl_ef.setToolTip("Enter ASL to define a group of atoms") self.asl_ef.textChanged.connect(self.aslTextModified) if show_asl: self.asl_layout.addWidget(self.asl_ef) self.asl_ef.addClearButton() self.asl_ef.clearButtonClicked.connect(self.reset) else: self.asl_ef.setVisible(False) self.pick_toggle = None self.pick_layout = None self.default_pick_mode = default_pick_mode if show_pick: self.setupPickToggle(pick_text) if show_markers and maestro: if self.pick_layout is None: self.pick_layout = QtWidgets.QHBoxLayout() # If the markers checkbox is disabled, hide the markers. self.markers_toggle = swidgets.SCheckBox(self, disabled_checkstate=False) self.d_markers_toggle = self.markers_toggle.isChecked() self.markers_toggle.setText("Markers") self.pick_layout.addWidget(self.markers_toggle) self.markers_toggle.setToolTip("Show markers in the Workspace " "on the selected atoms") self.marker = markers.Marker(color=PINK) self.marker.hide() self.markers_toggle.toggled.connect(self.marker.setVisible) else: self.marker = None # If only 1 or 2 buttons are to be shown, append them to the pick_layout, # otherwise create a separate row for them above the pick layout. # If show_plus then select & previous buttons will be added to '+' # button menu, otherwise they should be added to the outside button layout button_bools = (show_all, show_selection, show_markers, (not show_plus and show_previous), (not show_plus and show_select), (show_all and not show_asl)) show_button_row = len([button for button in button_bools if button]) > 2 if (self.pick_layout is None) or show_button_row: button_layout = QtWidgets.QHBoxLayout() else: button_layout = self.pick_layout # Whether to show the "All" button. It will be shown next to the ASL # field (if shown) or in the "button row" if show_all: self.all_button = QtWidgets.QToolButton(self) self.all_button.setObjectName('all_button') self.all_button.setToolTip( "Click to select all atoms in the Workspace") if maestro: self.all_button.setFixedWidth(button_height) self.all_button.setFixedHeight(button_height) else: self.all_button.setText("All") # If the ASL field is shown, add the "all" button to the right of it, # if it's hidden, add it to the button layout. if show_asl: self.asl_layout.addWidget(self.all_button) else: button_layout.addWidget(self.all_button) self.all_button.clicked.connect(self._allClicked) # Whether to show the "Selection" button. It places the WS selection into # the ASL field. if show_selection: self.selection_button = QtWidgets.QPushButton(self) # If show_asl and show_plus are True, then set the selection button # text to "Load Selection" and right align it. if show_asl and show_plus: button_layout.addStretch(1) selection_button_text = "Load Selection" # Set the object name to apply stylesheet self.selection_button.setObjectName("selection_button") self.selection_button.setText(selection_button_text) button_layout.addWidget(self.selection_button) self.selection_button.clicked.connect(self._selectionClicked) self.previous_button = None # Setup plus button if show_plus: plus_layout = self.asl_layout if ( show_asl or not show_button_row) else button_layout self.setupPlusButton(plus_layout, button_height, show_select, show_previous) else: if show_previous: self._addPreviousButton(button_layout) # Whether to show the "Select..." button, which opens a Maestro Selection dialog. if show_select: self._addSelectButton(button_layout) # Add layouts and add stretches to them: self.main_layout.addLayout(self.asl_layout) if show_button_row: button_layout.addStretch() self.main_layout.addLayout(button_layout) if self.pick_layout: self.pick_layout.addStretch() self.main_layout.addLayout(self.pick_layout) if not show_button_row and not self.pick_layout and show_asl and show_plus: button_layout.addStretch() self.main_layout.addLayout(button_layout) # Make the connections: self.asl_ef.editingFinished.connect(self._aslChanged)
[docs] def setupPickToggle(self, pick_text): """ Set up pick toggle by creating pick layout, and adding pick toggle & combo menu to it. :type pick_text: str :param pick_text: Text that will be displayed on the bottom of the main Maestro window when the pick button is checked. """ self.pick_layout = QtWidgets.QHBoxLayout() self.pick_toggle = QtWidgets.QCheckBox(self) self.pick_toggle.setText("Pick:") self.pick_layout.addWidget(self.pick_toggle) self.pick_menu = QtWidgets.QComboBox(self) # If altered, must change the constants at the top of the module as well: self.pick_menu.addItem("Atoms") self.pick_menu.addItem("Residues") self.pick_menu.addItem("Chains") self.pick_menu.addItem("Molecules") self.pick_menu.addItem("Entries") self.pick_menu.setCurrentIndex(self.default_pick_mode) self.pick_layout.addWidget(self.pick_menu) self.picksite_wrapper = picking.PickAtomToggle( self.pick_toggle, pick_text=pick_text, pick_function=self._atomPicked, ) if not maestro: self.pick_toggle.setEnabled(False) self.pick_menu.setEnabled(False)
[docs] def setupPlusButton(self, layout, button_height, show_select, show_previous): """ Setups the plus button by adding the button and its popup widget :type layout: QtWidgets.QLayout :param layout: Layout to which the plus button should be added. :type button_height: int :param button_height: Height and width of the plus button. :type show_previous: bool :param show_previous: Whether to show the "Previous" button. Clicking it would select the previous ASL that was used. :type show_select: bool :param show_select: Whether to show the "Select..." button. Clicking it would show an ASL selection dialog. """ self.plus_button = QtWidgets.QToolButton(self) self.plus_button.setObjectName('plus_button') self.plus_button.setFixedWidth(button_height) self.plus_button.setFixedHeight(button_height) self.plus_button.setToolTip( "Click to choose from the Workspace " "selection or predefined\natom sets, or to open the Atom " "Selection dialog") layout.addWidget(self.plus_button) self.popup_widget = ASLPopupMenu(self.plus_button) self.plus_button.setMenu(self.popup_widget) self.plus_button.setPopupMode(self.plus_button.InstantPopup) self.plus_button.setArrowType(Qt.NoArrow) self.popup_widget.hidden.connect(self.updatePlusButtonStyle) self._addPopupMenuItems() if show_select or show_previous: self.popup_widget.addSeparator() button_widget = QtWidgets.QWidget() button_layout = QtWidgets.QHBoxLayout(button_widget) button_layout.setContentsMargins(5, 2, 5, 2) button_layout.setSpacing(10) qt_utils.add_widget_to_menu(button_widget, self.popup_widget) # Whether to show the "Select..." button, which opens a Maestro Selection # dialog. if show_select: self._addSelectButton(button_layout) self.select_button.clicked.connect(self.popup_widget.hide) if show_previous: self._addPreviousButton(button_layout, "Previous ASL")
def _addPopupMenuItems(self): """ Add the asl menu items to the popup menu. Can be overwritten in child classes to customize the menu """ for name, asl in self.ASL_ITEMS.items(): asl_item = ASLItem(name, asl, self.plus_button) asl_item.itemClicked.connect(self.aslItemClicked) qt_utils.add_widget_to_menu(asl_item, self.popup_widget) def _addSelectButton(self, layout): """ Add "Select..." button to the given layout :type layout: QtWidgets.QLayout :param layout: Layout to which the select button should be added. """ self.select_button = QtWidgets.QPushButton(self) self.select_button.setText("Select...") self.select_button.setToolTip("Click to select atoms") layout.addWidget(self.select_button) self.select_button.clicked.connect(self._selectClicked) def _addPreviousButton(self, layout, text="Previous"): """ Add previous button to the given layout :type layout: QtWidgets.QLayout :param layout: Layout to which the previous button should be added. :type text: str :param text: Text of the previous button. """ self.previous_button = QtWidgets.QPushButton(self) self.previous_button.setText(text) self.previous_button.setToolTip("Click to use previous ASL " "defined in the Atom Selection dialog") layout.addWidget(self.previous_button) self.previous_button.clicked.connect(self._previousClicked) self.previous_button.setEnabled(False)
[docs] def updatePlusButtonStyle(self): """ Update plus_button style """ qt_utils.update_widget_style(self.plus_button)
def _aslChanged(self): asl = str(self.asl_ef.text()) if self.marker and asl: self.marker.setAsl(asl) elif self.marker: # If the previous ASL is now deleted don't show any markers. self.marker.setAsl('NOT all') self.aslModified.emit(asl) def _setAsl(self, asl): # For backwards compatibility self.setAsl(asl)
[docs] def setAsl(self, asl): # FIXME: the previous value should be global with Maestro's previous value self._previous_value = str(self.asl_ef.text()) if self.previous_button: self.previous_button.setEnabled(True) self.asl_ef.setText(asl) self._aslChanged()
[docs] def reset(self): """ Reset the widget to the defaults """ self.setAsl("") if self.previous_button: self.previous_button.setEnabled(False) if self.pick_toggle: self.picksite_wrapper.stop() self.pick_menu.setCurrentIndex(self.default_pick_mode) if self.marker: self.markers_toggle.setChecked(self.d_markers_toggle)
[docs] def aslItemClicked(self, asl): """ Update the asl according to the clicked asl item. """ self.popup_widget.hide() current_asl = self.getAsl() final_asl = asl if self._append_mode and current_asl: final_asl = '%s or (%s)' % (current_asl, asl) self.setAsl(final_asl)
[docs] def loadWorkspaceSelection(self): """ Loads workspace selection to the AtomSelector """ if maestro: self._selectionClicked()
def _allClicked(self): self.setAsl("all") def _selectionClicked(self): asl = maestro.selected_atoms_get_asl() if not asl: asl = "" self.setAsl(asl) def _previousClicked(self): self.setAsl(self._previous_value) def _selectClicked(self): prev_asl = str(self.asl_ef.text()) self.atomSelectionDialogAboutToBeShown.emit() asl = maestro.atom_selection_dialog("Atom selection", current_asl=prev_asl) if asl != "": # User pressed OK self.setAsl(asl) if self.parent_window: self.parent_window.raise_() self.parent_window.activateWindow() self.atomSelectionDialogDismissed.emit() def _atomPicked(self, atomnum): index = self.pick_menu.currentIndex() if index == PICK_ATOMS: picked_asl = "atom.num %i" % atomnum elif index == PICK_RESIDUES: workspace_st = maestro.workspace_get(copy=False) atom = workspace_st.atom[atomnum] picked_asl = '(mol.num %i and chain.name "%s" and res.inscode "%s" and res.num %i)' % \ (atom.molecule_number, atom.chain, atom.inscode, atom.resnum) elif index == PICK_CHAINS: workspace_st = maestro.workspace_get(copy=False) atom = workspace_st.atom[atomnum] picked_asl = 'chain.name "%s"' % atom.chain # FIXME probably need to further qualify with "molecule" or you may # get multiple chains selected (comment from code review). Though # this will not work correctly if there are gaps in the chain. # Perhaps something like this would work as expected? # atom_list = atom.getChain().getAtomIndices() # picked_asl = analyze.generate_asl(workspace_st, atom_list) elif index == PICK_MOLECULES: workspace_st = maestro.workspace_get(copy=False) atom = workspace_st.atom[atomnum] picked_asl = 'mol.num %i' % atom.molecule_number elif index == PICK_ENTRIES: workspace_st = maestro.workspace_get(copy=False) atom = workspace_st.atom[atomnum] picked_asl = 'entry.id %s' % atom.entry_id else: raise ValueError("Invalid value for the pick_menu") prev_asl = self.asl_ef.text() if prev_asl and self._append_mode: self.setAsl("%s | %s" % (prev_asl, picked_asl)) else: self.setAsl(picked_asl)
[docs] def getAsl(self): """ Return the selected ASL string """ return str(self.asl_ef.text())
[docs] def setDarkStyle(self): """ Set dark style to 'All' and '+' buttons """ for button in (self.all_button, self.plus_button): button.setProperty("dark", "true")
[docs] def setMarkersEnabled(self, enable: bool): """ Set enabled/disabled state of the `Markers` checkbox. Note: Disabling the `Markers` checkbox hides the workspace markers. :param enable: Whether to enable or disable the `Markers` checkbox. """ if self.marker is None: # Not showing markers. return self.markers_toggle.setEnabled(enable)
[docs]class MappableAtomSelector(mappers.TargetMixin, AtomSelector):
[docs] def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.aslTextModified.connect(self.targetValueChanged)
[docs] def targetGetValue(self): # @overrides: mappers.TargetMixin return self.getAsl()
[docs] def targetSetValue(self, asl): # @overrides: mappers.TargetMixin self.setAsl(asl)
[docs]def panel(): frame = QtWidgets.QFrame() atom_selector = AtomSelector(frame) layout = QtWidgets.QVBoxLayout() layout.addWidget(atom_selector) frame.setLayout(layout) frame.show()
if __name__ == "__main__": # Test for this widget (Run outside of Maestro) # Shows how to use the AtomSelector widget import sys app = QtWidgets.QApplication(sys.argv) frame = QtWidgets.QFrame() atom_selector = AtomSelector(frame) def print_asl(asl): print(asl) atom_selector.aslModified.connect(print_asl) layout = QtWidgets.QVBoxLayout() layout.addWidget(atom_selector) frame.setLayout(layout) frame.show() frame.raise_() sys.exit(app.exec())