from collections import Counter
from collections import OrderedDict
import schrodinger
from schrodinger.application.jaguar import input as jaginput
from schrodinger.infra import mm
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt.QtCore import Qt
from schrodinger.structutils import analyze
from schrodinger.ui import picking
from .base_tab import BaseTab
maestro = schrodinger.get_maestro()
try:
from schrodinger.maestro import markers
except ImportError:
markers = None
PICK_COMBO_ATOMS, PICK_COMBO_BONDS = list(range(2))
PICK_ATOMS = ["Atoms"]
PICK_ATOMS_AND_BONDS = ["Atoms", "Bonds"]
PICK_TYPES = OrderedDict(
((mm.MMJAG_COORD_CART_X, ("x", PICK_ATOMS, 1)), (mm.MMJAG_COORD_CART_Y,
("y", PICK_ATOMS, 1)),
(mm.MMJAG_COORD_CART_Z, ("z", PICK_ATOMS, 1)), (mm.MMJAG_COORD_CART_XYZ,
("xyz", PICK_ATOMS, 1)),
(mm.MMJAG_COORD_DISTANCE, ("distance", PICK_ATOMS_AND_BONDS, 2)),
(mm.MMJAG_COORD_ANGLE, ("angle", PICK_ATOMS_AND_BONDS, 3)),
(mm.MMJAG_COORD_TORSION, ("dihedral", PICK_ATOMS_AND_BONDS, 4))))
COORD_TYPES = {
mm.MMJAG_COORD_CART_X, mm.MMJAG_COORD_CART_Y, mm.MMJAG_COORD_CART_Z,
mm.MMJAG_COORD_CART_XYZ
}
[docs]class CoordinateTab(BaseTab):
"""
A parent class for the Scan and Optimization tabs
:cvar coordinateAdded: A signal emitted when user adds new
coordinate. The signal is emitted with:
- a list of atom numbers
- coordinate type
:vartype coordinateAdded: `PyQt5.QtCore.pyqtSignal`
:cvar coordinateDeleted: A signal emitted when user deletes a coordinate.
Emitted with a list of atom numbers.
:vartype coordinateDeleted: `PyQt5.QtCore.pyqtSignal`
:cvar allCoordinatesDeleted: A signal emitted when all coordinates are
deleted. Emitted with no arguments.
:vartype allCoordinatesDeleted: `PyQt5.QtCore.pyqtSignal`
:cvar coordinateSelected: A signal emitted when user selects a coordinate in
the table. Emitted with a list of atom numbers for the selected coordinate.
:vartype coordinateSelected: `PyQt5.QtCore.pyqtSignal`
:cvar coordinateDeselected: A signal emitted when user deselects a
coordinate in the table. Emitted with a list of atom numbers for the
deselected coordinate.
:vartype coordinateDeselected: `PyQt5.QtCore.pyqtSignal`
:cvar refreshMarkers: A signal emitted when the workspace markers should be
refreshed, i.e., when we should make sure that only markers for the
currently selected tab are displayed.
:vartype coordinateDeselected: `PyQt5.QtCore.pyqtSignal`
"""
coordinateAdded = QtCore.pyqtSignal(list, int)
coordinateDeleted = QtCore.pyqtSignal(list)
allCoordinatesDeleted = QtCore.pyqtSignal()
coordinateSelected = QtCore.pyqtSignal(list)
coordinateDeselected = QtCore.pyqtSignal(list)
refreshMarkers = QtCore.pyqtSignal()
[docs] def setup(self):
super(CoordinateTab, self).setup()
self._acceptable_constraint_eids = set()
self._picking_err = None
self._marker_count = Counter()
def _getAtomsForRow(self, row_num):
"""
Get the atom indices for the specified row
:param row_num: A row number
:type row_num: int
:return: A list of atom indices
:rtype: list
"""
indices_column = self.model.COLUMN.INDICES
indices_index = self.model.index(row_num, indices_column)
atoms = indices_index.data()
return atoms
def _highlightSelectedMarkers(self, current_sel, previous_sel):
"""
Respond to a change in the selected table rows by changing the currently
selected workspace marker
:param current_sel: The new table selection
:type current_sel: `PyQt5.QtCore.QItemSelection`
:param previous_sel: The previous table selection
:type previous_sel: `PyQt5.QtCore.QItemSelection`
"""
# Make sure we deselect the old marker before we select the new one in
# case both rows refer to the same marker
self._checkSelection(previous_sel, self.coordinateDeselected)
self._checkSelection(current_sel, self.coordinateSelected)
def _checkSelection(self, sel, signal):
"""
If the specified table selection is not empty, emit the given signal
with the atom indices from the selected row.
:param sel: The table selection
:type sel: `PyQt5.QtCore.QItemSelection`
:param signal: The singal to emit
:type signal: `PyQt5.QtCore.pyqtSignal`
"""
indices = sel.indexes()
if not indices:
return
row = indices[0].row()
atoms = self._getAtomsForRow(row)
if atoms is not None:
signal.emit(atoms)
def _determineIfConstraintsAddable(self):
"""
Determine if the panel, workspace, and project are in a state where we
can add constraints.
:return: If we can add constraints, return None. Otherwise, return the
text of the error message that should be displayed to the user.
:rtype: NoneType or str
"""
if self._picking_err:
return self._picking_err
ws_struc = maestro.workspace_get()
included = {atom.entry_id for atom in ws_struc.atom}
if len(included) != 1:
return ("Only one entry may be included in the workspace when "
"adding constraints")
eid = included.pop()
if eid in self._acceptable_constraint_eids:
return None
else:
return ("The workspace entry is not currently selected")
[docs] def setAcceptableContraintEids(self, eids, picking_err):
"""
Set the constraint picking restrictions
:param eids: The entry ids for which coordinate picking is acceptable.
:type eids: set
:param picking_err: If picking should not be allowed at all, this is the
text of the error that will displayed to the user. If picking is
allowed, should be None.
:type picking_err: str or NoneType
"""
self._acceptable_constraint_eids = eids
self._picking_err = picking_err
self.deleteAllRows()
[docs] def stopPicking(self):
"""
Stop constraint picking
"""
self.ui.pick_cb.setChecked(False)
def _resetDefaults(self):
"""
This function resets coordinates table and sets coordinate picking to
its default state. Note that this function is not called reset() since
it does not need to be called from the panel class.
"""
self.deleteAllRows()
self.ui.pick_cb.setChecked(False)
self.ui.coord_type_combo.setCurrentIndex(0)
self.ui.pick_combo.setCurrentIndex(0)
def _emitCoordinateAdded(self, atoms, coordinate_type):
"""
If a marker doesn't yet exist for the specified atoms, emit the
coordinateAdded signal to request that one be created. For single atom
markers, We use _marker_count to keep track of the the number of
constraints per marker. Note that We can only get multiple constraints
per markers for single atom markers (i.e. if there's a Coordinate - X
constraint and a Coordinate - Y constraint for the same atom)
:param atoms: A list of atom indices for the atoms to mark
:type atoms: list
:param coordinate_type: The coordinate type
:type coordinate_type: int
"""
if len(atoms) == 1:
atom_num = atoms[0]
if self._marker_count[atom_num] == 0:
self.coordinateAdded.emit(atoms, coordinate_type)
self._marker_count[atom_num] += 1
else:
self.coordinateAdded.emit(atoms, coordinate_type)
def _emitCoordinateDeleted(self, atoms):
"""
After a constraint has been deleted for the specified atoms, delete the
corresponding marker if there are no other constraints for that group of
atoms.
:param atoms: A list of atom indices for the deleted constraint
:type atoms: list
"""
if len(atoms) == 1:
atom_num = atoms[0]
self._marker_count[atom_num] -= 1
if self._marker_count[atom_num] == 0:
self.coordinateDeleted.emit(atoms)
else:
self.coordinateDeselected.emit(atoms)
else:
self.coordinateDeleted.emit(atoms)
[docs]class CoordinatePicker(QtCore.QObject):
"""
This class is responsible for atom and bond picking. Depending on the type
of coordinate it will fill up the list of picked atoms up to a max size for
the current coordinate type before emitting a signal.
:cvar pickCompleted: A signal emitted when the user picks required
number of atoms for current coordinate type. This signal is emitted
with a list of picked atoms as an argument.
:vartype pickCompleted: `PyQt5.QtCore.pyqtSignal`
:cvar PICK_MAX_ATOMS: A dictionary that maps mmjag coordinate type to max
number of atoms needed to define this coordinate. It should include all
coordinate types used in Jaguar Scan and Optimization tabs, where picker
is used.
:vartype PICK_MAX_ATOMS: dict
"""
pickCompleted = QtCore.pyqtSignal(list)
PINK = (1.0, 0.8, 0.8)
PICK_MAX_ATOMS = OrderedDict(
((mm.MMJAG_COORD_CART_X, 1), (mm.MMJAG_COORD_CART_Y, 1),
(mm.MMJAG_COORD_CART_Z, 1), (mm.MMJAG_COORD_CART_XYZ, 1),
(mm.MMJAG_COORD_DISTANCE, 2), (mm.MMJAG_COORD_ANGLE,
3), (mm.MMJAG_COORD_TORSION, 4)))
[docs] def __init__(self,
coordinate_types,
pick_cb,
coord_type_combo,
pick_combo,
parent=None):
"""
Picker class initializer.
:param coordinate_types: ordered dictionary that contains
coordinate types that should be made available in the picker.
:type coordinate_types: `collections.OrderedDict`
:param pick_cb: check box used to pick atoms or bonds
:type pick_cb: `QtWidgets.QCheckBox`
:param coord_type_combo: combo box that allows to select
coordinate type such as distance, angle etc.
:type coord_type_combo: `QtWidgets.QComboBox`
:param pick_combo: combo box that allows to select pick
type: atom or bond.
:type pick_combo: `QtWidgets.QComboBox`
"""
super(CoordinatePicker, self).__init__(parent)
self.coord_type_combo = coord_type_combo
self.picked_atoms = []
self.coordinate_types = coordinate_types
self.pick_cb = pick_cb
self.coord_type_combo = coord_type_combo
self.pick_combo = pick_combo
self.picker = picking.PickMixedToggle(
self.pick_cb,
enable_lasso=True,
pick_atom_function=self._pickAtom,
pick_bond_function=self._pickBond,
pick_atom_text="Pick atom to define scan coordinate.",
pick_bond_text="Pick bond to define scan coordinate.")
if maestro:
self.marker = markers.Marker(color=self.PINK)
# populate pick widgets with coordinate types
self.populateTypeCombo()
self.populatePickCombo()
# make connections
coord_type_combo.currentIndexChanged.connect(self.coordinateTypeChanged)
pick_combo.currentIndexChanged.connect(self._pickChanged)
pick_cb.stateChanged.connect(self._clearPicked)
[docs] def populateTypeCombo(self):
"""
This function is used to initialize coordinate type combox box.
"""
self.coord_type_combo.addItemsFromDict(self.coordinate_types)
[docs] def populatePickCombo(self):
"""
This function repopulates pick combo box depending on the
current selection of coordinate type. It also attempts to
preserve current pick type selection if possible.
"""
index = self.pick_combo.currentIndex()
data = self.coord_type_combo.currentData()
coord_type, items, max_atoms = PICK_TYPES[data]
self.pick_combo.clear()
self.pick_combo.addItems(items)
if index < 0 or index >= self.pick_combo.count():
index = 0
self.pick_combo.setCurrentIndex(index)
# Enable lasso only for Cartesian picking (not for dist/angle/dihedral):
self.picker.enable_lasso = data in COORD_TYPES
[docs] def coordinateTypeChanged(self):
"""
This slot is called when coordinate type is changed.
"""
self.populatePickCombo()
self.pick_cb.setChecked(True)
self._clearPicked()
def _pickAtom(self, atom_or_asl):
"""
This function is called when atom or atoms are picked.
:param atom_or_asl: atom number of the picked atom or ASL string.
:type atom_or_asl: int or str
"""
if type(atom_or_asl) is int:
self.picked_atoms.append(atom_or_asl)
else:
st = maestro.workspace_get()
self.picked_atoms += analyze.evaluate_asl(st, atom_or_asl)
self._checkPicked()
def _pickBond(self, atom1, atom2):
"""
This function is called when bond is picked.
:param atom1: atom number of the first bond atom
:type atom1: int
:param atom2: atom number of the second bond atom
:type atom2: int
"""
if len(self.picked_atoms) == 0:
self.picked_atoms.extend([atom1, atom2])
else:
if atom1 in self.picked_atoms:
self.picked_atoms.append(atom2)
elif atom2 in self.picked_atoms:
self.picked_atoms.append(atom1)
else:
self._clearPicked()
return
self._checkPicked()
def _checkPicked(self):
"""
This function checks the number of currently picked atoms. If we have
enough atoms to define coordinate of the current type pickCompleted
signal is emitted and list of currently picked atoms is cleared.
"""
data = self.coord_type_combo.currentData()
max_atoms = self.PICK_MAX_ATOMS[data]
if data in COORD_TYPES:
# In Cartesian(coordinate) constraints, each atom is a unique constraint
for anum in self.picked_atoms:
self.pickCompleted.emit([anum])
self._clearPicked()
elif len(self.picked_atoms) == max_atoms:
# Completed picking the distance/angle/torsion constraint:
self.pickCompleted.emit(self.picked_atoms)
self._clearPicked()
else:
self._setASL()
def _setASL(self):
"""
This function is used to show markers in the workspace for a
set of currently picked atoms.
"""
atom_strings = list(map(str, self.picked_atoms))
asl = "atom.num %s" % ",".join(atom_strings)
if maestro:
self.marker.show(asl=asl)
def _pickChanged(self, index):
"""
This function is called when we switch between atom and bond picking.
"""
if index == PICK_COMBO_ATOMS:
self.picker.setPickAtom(True)
elif index == PICK_COMBO_BONDS:
self.picker.setPickBond(True)
def _clearPicked(self):
"""
This function clears the list of picked atoms and remove markers from
the workspace.
"""
self.picked_atoms = []
if maestro:
self.marker.hide()
[docs]class CoordinateData(object):
"""
This class is a base class for constraint and scan coordinate
classes. It should not(!) be initialized by itself.
:ivar st: ct structure for which coordinates are defined
:vartype st: `schrodinger.structure.Structure`
:ivar atom_indices: indices of atoms, which define this coordinate
:vartype atom_indices: list
:ivar coordinate_name: name of this coordinate based on atom indices
:vartype coordinate_name: str
:ivar coordinate_type: coordinate type
:vartype coordinate_type: int
:cvar COLUMN: class that contains information about columns in
which coordinates data is displayed. It should contain NAMES
variable for column names and indices of columns. This object
needs to be initialize in derived classes.
:vartype COLUMN: object
"""
[docs] def __init__(self, st, atoms, coordinate_type):
"""
Initialize coordinates data given a structure, set of atom indices and
coordinate type. We apply the jaguar naming scheme to the
structure.
:param st: structure
:type st: `schrodinger.structure.Structure`
:param atoms: atom indices
:type atoms: list
:param coordinate_type: coordinate type
:type coordinate_type: int
"""
self.st = st
jaginput.apply_jaguar_atom_naming(self.st)
self.atom_indices = atoms
self.coordinate_type = coordinate_type
self.validate()
[docs] def validate(self):
"""
This function checks that atom indices contain correct number of
elements for a given coordinate type. If thats not the case ValueError
exception is raised.
"""
# check atom indices in the list
num_atoms = len(self.st.atom)
for idx in self.atom_indices:
if idx < 1 or idx > num_atoms:
raise ValueError("Incorrect atom index")
# check that atom list has correct length
try:
c_type, items, max_atoms = PICK_TYPES[self.coordinate_type]
except KeyError:
raise ValueError("Incorrect coordinate type %s" %
self.coordinate_type)
if len(self.atom_indices) != max_atoms:
raise ValueError("Coordinate atom indices list length does not"
" match coordinate type")
def _getCoordinateName(self):
"""
This function returns coordinate name, which is constructed
from the atom names.
:return: coordinate name
:rtype: str
"""
return " ".join([self._getAtomName(x) for x in self.atom_indices])
def _getAtomName(self, atom):
"""
This function converts an atom index into an atom name.
:param atom: atom index
:type atom: int
:return: atom name
:rtype: str
"""
return self.st.atom[atom].name
[docs]class CoordinatesModel(QtCore.QAbstractTableModel):
"""
A base class for cordinates models used for constraint and scan
coordinates in Scan and Optimization tabs. This class should not(!)
be initialized on its own. This model is used with Qt view.
:cvar COLUMN: class that contains information about columns in
which coordinates data is displayed. It should contain NAMES
variable for column names and indices of columns. This object
needs to be initialize in derived classes.
:vartype COLUMN: object
"""
COLUMN = None
[docs] def __init__(self, parent=None):
super(CoordinatesModel, self).__init__(parent)
self.coords = []
[docs] def rowCount(self, parent=None):
"""
Return the number of rows in the model
:param parent: Unused, but preset for PyQt compatibility
:return: The number of rows in the model.
:rtype: int
"""
return len(self.coords)
[docs] def columnCount(self, parent=None):
"""
Return the number of columns in the model
:param parent: Unused, but preset for PyQt compatibility
:return: The number of columns in the model.
:rtype: int
"""
return self.COLUMN.NUM_COLS
[docs] def checkNewCoordinate(self, atoms, coordinate_type):
"""
This function check whether this coordinate is already present
in this model.
:param atoms: atom indices
:type atoms: list
:param coordinate_type: coordinate type
:type coordinate_type: int
:return: True if this coordinate has not been found and False otherwise.
:rtype: bool
"""
is_new = True
if self.findCoordinate(atoms, coordinate_type) is not None:
is_new = False
return is_new
[docs] def reset(self):
"""
Remove any existing data
"""
self.beginResetModel()
self.coords = []
self.endResetModel()
[docs] def removeRow(self, row, parent=QtCore.QModelIndex()): # noqa: M511
"""
Removes the given row from the child items of the parent specified.
Returns true if the row is removed; otherwise returns false.
:param row: row index
:type row: int
:param index: parent index
:type index: `QtCore.QModelIndex`
:return: True or False
:rtype: bool
"""
if row < 0 or row >= self.rowCount() or self.rowCount == 0:
return False
self.beginRemoveRows(parent, row, row)
self.coords.pop(row)
self.endRemoveRows()
return True
[docs] def removeCoordinate(self, atoms, coordinate_type):
"""
This function searches for a given coordinate. If match is found
coordinate is removed.
:param atoms: atom indices
:type atoms: list
:param coordinate_type: coordinate type
:type coordinate_type: int
:return: True if this coordinate was found and removed, False otherwise.
:rtype: bool
"""
idx = self.findCoordinate(atoms, coordinate_type)
if idx is not None:
return self.removeRow(idx)
return False
[docs] def findCoordinate(self, atoms, coordinate_type):
"""
This function searches for coordinate defined by atoms list and
coordinate type. If match is found this function returns row index
and None otherwise.
:param atoms: atom indices
:type atoms: list
:param coordinate_type: coordinate type
:type coordinate_type: int
:return: row index if this coordinate has been found and None otherwise.
:rtype: int or None
"""
for idx, coord in enumerate(self.coords):
same_type = (coordinate_type == coord.coordinate_type)
same_coord = (set(coord.atom_indices) == set(atoms))
if same_type and same_coord:
return idx
return None