import decorator
import schrodinger
from schrodinger import project
from schrodinger import structure
from schrodinger.infra import mm
from schrodinger.structutils import analyze
from schrodinger.ui import maestro_ui
from schrodinger.ui.qt.appframework2 import maestro_callback
maestro = schrodinger.get_maestro()
try:
from schrodinger.maestro import markers
except ImportError:
markers = None
@decorator.decorator
def _requires_maestro(func, *args, **kwargs):
"""
A decorator that raises a MaestroNotAvailableError if the decorated function
is called outside of Maestro.
NOTE: Behavior is different from schrodinger.ui.qt.utils.maestro_required.
"""
if not maestro:
err = "Markers are not available outside of Maestro."
raise schrodinger.MaestroNotAvailableError(err)
else:
return func(*args, **kwargs)
[docs]class MarkerMixin(object):
"""
A mixin for adding markers and controlling their visibility. Note that this
Mixin requires the `maestro_callback.MaestroCallbackMixin`.
:ivar _markers: A dictionary containing all `markers._BaseMarker` derived
markers associated with this panel. Keys are generated via
`_canonicalizeAtomOrder` and `_genMarkerHash`.
:vartype _markers: dict
:ivar _marked_eid_lengths: A dictionary of {entry id: number of atoms in the
entry}. This dictionary is used to delete markers if the number of atoms in
a marked entry changes.
:vartype _marked_eid_lengths: dict
:ivar MARKER_ICONS: An object containing constants for all available marker
icons
:vartype MARKER_ICONS: `schrodinger.maestro.markers.Icons`
:ivar _multi_atom_markers: A list containing all `markers.Marker` markers
associated with this panel.
:vartype _multi_atom_markers: list
"""
[docs] def __init__(self, *args, **kwargs):
self._markers = {}
self._marked_eid_lengths = {}
self._multi_atom_markers = []
self._markername_to_marker_map = {}
if markers:
# We can't access the maestro.markers module outside of Maestro, so
# bind the constants at runtime if maestro.markers has been
# imported
self.MARKER_ICONS = markers.Icons
else:
self.MARKER_ICONS = object()
super(MarkerMixin, self).__init__(*args, **kwargs)
if hasattr(self, 'finished'):
# Hide all markers when dialog is hidden
self.finished.connect(self._hideAll)
[docs] def showEvent(self, event):
"""
Re-show all panel markers when the panel is re-shown.
"""
if not event.spontaneous():
# Ignore spontaneous events (i.e. un-minimizing the window)
self._showAll()
try:
super(MarkerMixin, self).showEvent(event)
except AttributeError:
pass # Required for QDialog subclassing.
[docs] def show(self):
"""
Re-show all panel markers when the panel is re-shown. This separate
method is needed for QDialog instances.
"""
self._showAll()
super(MarkerMixin, self).show()
[docs] def hideEvent(self, event):
self._hideAll()
try:
super(MarkerMixin, self).hideEvent(event)
except AttributeError:
pass # In case the base class has no hideEvent() method.
[docs] def closeEvent(self, event):
"""
Hide all markers when the panel is closed.
"""
self._hideAll()
try:
super(MarkerMixin, self).closeEvent(event)
except AttributeError:
pass # Required for QDialog subclassing.
[docs] @_requires_maestro
def addJaguarMarker(self,
atoms,
color=None,
icon=None,
text="",
alt_color=None,
highlight=False):
"""
Add a marker to the specified atom(s)
:param atoms: The atom or list of atoms to mark. A list may
contain between one and four atoms (inclusive).
:type atoms: list or `schrodinger.structure._StructureAtom`
:param color: The color of the marker and icon. May be an RGB tuple,
color name, color index, or `schrodinger.structutils.color` instance.
If not given, white will be used.
:type color: tuple, str, int, or `schrodinger.structutils.color`
:param icon: The icon to draw next to the marker. Should be one the
self.MARKER_ICONS constants. If not given, no icon
will be drawn.
:type icon: int
:param text: The text to display next to the marker. If not given, no
text will be displayed. Note that this argument will be ignored when
marking a single atom.
:type text: str
:param alt_color: The alternate marker color. This color is always used
for text, and is used for the marker and icon when `highlight` is True.
If not given, `color` will be used.
:type alt_color: tuple, str, int, or `schrodinger.structutils.color`
:param highlight: Whether the marker should be highlighted. A
highlighted marker is indicated with thicker lines and is colored using
`alt_color` instead of `color`.
:type highlight: bool
:return: The newly created marker
:rtype: `schrodinger.maestro.markers._BaseMarker`
:raise ValueError: If a marker already exists for the specified atoms
:note: Either an icon or text may be displayed on a marker, but not
both. If both are given, only the text will be shown.
"""
if icon is None:
icon = self.MARKER_ICONS.NONE
atoms = self._canonicalizeAtomOrder(atoms)
atoms_hashable, entry_ids = self._genMarkerHash(atoms)
if atoms_hashable in self._markers:
raise ValueError("A marker already exists for the specified atoms.")
marker = self._createJaguarMarker(atoms, color, icon, text, alt_color,
highlight)
self._setMarkerHash(marker, atoms_hashable, entry_ids)
return marker
[docs] @_requires_maestro
def addMarker(self, atoms, color=(1., 1., 1.), group_name=None):
"""
Generates a set of simple, dot-styled markers for a group of atoms.
:param atoms: List of atoms to be marked
:type atoms: list or `schrodinger.structure._StructureAtom`
:param color: The amount of red, green and blue to use, each ranging
from 0.0 to 1.0. Default is white (1., 1., 1.).
:type color: tuple of 3 floats
@group_name: Optional string to set as the name of this group of
markers in Maestro. If not set, a unique identifier will be generated.
"""
if not atoms:
raise ValueError("Specify at least one atom to mark")
atoms = self._canonicalizeAtomOrder(atoms)
marker = self._createMarker(atoms, color, group_name)
self._multi_atom_markers.append(marker)
self._markername_to_marker_map[marker.name] = marker
return marker
[docs] @_requires_maestro
def addMarkerFromAsl(self, asl, color=(1., 1., 1.), group_name=None):
"""
Generates a set of simple, dot-styled markers for group of Workspace
atoms that match the given ASL. Same atoms continue to be marked even
if the Workspace is later modified such that ASL matching changes.
:param asl: ASL for the atoms to mark.
:type atoms: str
:param color: The amount of red, green and blue to use, each ranging
from 0.0 to 1.0. Default is white (1., 1., 1.).
:type color: tuple of 3 floats
@group_name: Optional string to set as the name of this group of
markers in Maestro. If not set, a unique identifier will be generated.
:return: Marker object
:rtype: `markers.Marker`
"""
st = maestro.workspace_get()
atoms = analyze.evaluate_asl(st, asl)
atoms = [st.atom[anum] for anum in atoms]
return self.addMarker(atoms, color, group_name)
def _setMarkerHash(self, marker, atoms_hashable, entry_ids):
"""
Store the hash with the marker for use in removeJaguarMarker()
:param marker: Marker to be hashed
:type marker: `schrodinger.maestro.marker._BaseMarker` or
`schrodinger.maestro.marker.Marker` instance
:param atoms_hashable: Unique identifier for atoms being marked as
generated by _genMarkerHash
:type atoms_hashable: tuple
:param entry_ids: Set of entry IDs for the atoms being marked
:type entry_ids: set
"""
marker.hashable = atoms_hashable
self._markers[atoms_hashable] = marker
for cur_entry_id in entry_ids:
if cur_entry_id not in self._marked_eid_lengths:
num_atoms = self._calcEntryAtomTotal(cur_entry_id)
self._marked_eid_lengths[cur_entry_id] = num_atoms
[docs] @_requires_maestro
def getJaguarMarker(self, atoms):
"""
Retrieve a marker for the specified atom(s)
:param atoms: The atom or list of atoms to retrieve the marker for. A
list may contain between one and four atoms (inclusive).
:type atoms: list or `schrodinger.structure._StructureAtom`
:return: The requested marker
:rtype: `schrodinger.maestro.markers._BaseMarker`
:raise ValueError: If no marker exists for the specified atoms
:note: As indicated by the return type, this function only returns
`schrodinger.maestro.markers._BaseMarker` derived markers.
Multi atom `schrodinger.maestro.markers.Marker` type markers
are not accessible in this way.
"""
atoms = self._canonicalizeAtomOrder(atoms)
atoms_hashable, entry_ids = self._genMarkerHash(atoms)
try:
return self._markers[atoms_hashable]
except KeyError:
err = "No marker exists for the specified atoms"
raise ValueError(err)
def _canonicalizeAtomOrder(self, atoms):
"""
Make sure that `atoms` is in a standard order. In other words,
self._canonicalizeAtomOrder(atoms) is guaranteed to be equal to
self._canonicalizeAtomOrder(reversed(atoms)). This function ensures
that we don't have to worry about atom order when indexing
self._markers. Note that this function also converts an atom (i.e. not
a list) to a list of a single atom. This is necessary for input to
_createJaguarMarker().
:param atoms: An atom or list of atoms
:type atoms: list or `schrodinger.structure._StructureAtom`
:return: A list of atoms in standard order
:rtype: list
"""
if isinstance(atoms, structure._StructureAtom):
return [atoms]
elif len(atoms) == 1:
return atoms
elif atoms[0].entry_id < atoms[-1].entry_id:
return atoms
elif atoms[0].entry_id > atoms[-1].entry_id:
return list(reversed(atoms))
elif atoms[0].number_by_entry < atoms[-1].number_by_entry:
return atoms
elif atoms[0].number_by_entry > atoms[-1].number_by_entry:
return list(reversed(atoms))
else:
err = "The first and last marked atoms must be different"
raise ValueError(err)
def _createJaguarMarker(self, atoms, color, icon, text, alt_color,
highlight):
"""
Create a marker with the specified properties. See the docs to
addJaguarMarker for a description of the arguments.
:return: The newly created marker
:rtype: `schrodinger.maestro.markers._BaseMarker`
"""
if len(atoms) == 1:
return markers.AtomMarker(atoms[0], color, icon, alt_color,
highlight)
elif len(atoms) == 2:
return markers.PairMarker(atoms[0], atoms[1], color, icon, text,
alt_color, highlight)
elif len(atoms) == 3:
return markers.TripleMarker(atoms[0], atoms[1], atoms[2], color,
icon, text, alt_color, highlight)
elif len(atoms) == 4:
return markers.QuadMarker(atoms[0], atoms[1], atoms[2], atoms[3],
color, icon, text, alt_color, highlight)
else:
err = "Atom list must contain between one and four atoms."
raise ValueError(err)
def _createMarker(self, atoms, color=(1., 1., 1.), group_name=None):
"""
Creates a `schrodinger.maestromarkers.Marker` instance for the
specified group of atoms.
:param atoms: List of atoms to be marked
:type atoms: list of `schrodinger.structure._StructureAtom` instances
:param color: The amount of red, green and blue to use, each ranging
from 0.0 to 1.0. The default is white (1., 1., 1.).
:type color: tuple of 3 floats
:param group_name: Optional name to set for the group of marked atoms
in Maestro. If not set, Maestro will generate a unique name.
:type group_name: str
:return: The atom marker generated for the group of atoms
:rtype: `schrodinger.maestro.markers.Marker`
"""
marker_asl = self._createAtomAsl(atoms)
return markers.Marker(asl=marker_asl, name=group_name, color=color)
def _createAtomAsl(self, atoms):
"""
Creates unique ASL for the specified atoms that utilizes the atom's
entry ID specific data to ensure it is unique and will remain valid
regardless of workspace changes.
:param atoms: List of atoms to generate ASL for
:type atoms: list of `schrodinger.structure._StructureAtom`
:return: ASL expression to identify these atoms
:rtype: str
"""
atom_data = {}
for cur_atom in atoms:
atom_data.setdefault(cur_atom.entry_id,
set()).add(cur_atom.number_by_entry)
if None in atom_data:
raise ValueError('Can not mark atoms that are not in the Workspace')
asl_per_eid = []
for eid in sorted(atom_data.keys(), key=int):
atom_nums = sorted(atom_data[eid])
joined_nums = ",".join(map(str, atom_nums))
cur_asl = "(entry.id %s AND atom.entrynum %s)" % (eid, joined_nums)
asl_per_eid.append(cur_asl)
return " OR ".join(asl_per_eid)
def _genMarkerHash(self, atoms):
"""
Create a unique hashable id from an iterable of atoms
:param atoms: An iterable of `schrodinger.structure._StructureAtom`
:type atoms: iterable
:return: A tuple of (a unique hashable id corresponding to the
specified atoms. This hashable id consists of a tuple of entry
ids and atom numbers by entry, a set of all entry ids from the
specified atoms)
:rtype: tuple
"""
atoms_hashable = []
entry_ids = set()
for cur_atom in atoms:
eid = cur_atom.entry_id
num_by_entry = cur_atom.number_by_entry
atoms_hashable.extend((eid, num_by_entry))
entry_ids.add(eid)
atoms_hashable = tuple(atoms_hashable)
return atoms_hashable, entry_ids
def _calcEntryAtomTotal(self, eid):
"""
Determine the number of atoms in the specified entry.
:param eid: The entry id
:type eid: str
:return: The number of atoms in the specified entry. If the entry has
been deleted, returns 0.
:rtype: int
"""
if eid.isdigit():
try:
proj = maestro.project_table_get()
struc = proj[eid].getStructure(props=False)
return struc.atom_total
except KeyError as xxx_todo_changeme:
mm.MmException = xxx_todo_changeme
return 0
except project.ProjectException:
# If the project was just closed
return 0
else:
ws_struc = maestro.workspace_get()
atoms = (1 for atom in ws_struc.atom if atom.entry_id == eid)
return sum(atoms)
[docs] @_requires_maestro
def removeJaguarMarker(self, marker):
"""
Removes the specified marker
:param marker: The marker to remove
:type marker: `schrodinger.maestro.markers._BaseMarker`
:raise ValueError: If there is no marker on the specified atoms
"""
try:
hashable = marker.hashable
self._markers[hashable].hide()
del self._markers[hashable]
except KeyError:
err = "The specified marker does not exist"
raise ValueError(err)
[docs] @_requires_maestro
def removeJaguarMarkerForAtoms(self, atoms):
"""
Removes the marker for specified atom(s)
:param atoms: The atom or list of atoms to retrieve the marker for. A
list may contain between one and four atoms (inclusive).
:type atoms: list or `schrodinger.structure._StructureAtom`
:raise ValueError: If no marker exists for the specified atoms
"""
marker = self.getJaguarMarker(atoms)
self.removeJaguarMarker(marker)
[docs] @_requires_maestro
def removeMarker(self, marker):
"""
Remove the `schrodinger.maestro.markers.Marker`
:param marker: Marker to remove
:type marker: `schrodinger.maestro.markers.Marker`
:raise ValueError: If marker is the wrong type or is not associated
with the panel.
"""
if markers and not isinstance(marker, markers.Marker):
msg = ("Specified marker is not a "
"schrodinger.maestro.markers.Marker instance.")
raise ValueError(msg)
all_names = [mrk.name for mrk in self._multi_atom_markers]
idx = all_names.index(marker.name)
cur_marker = self._multi_atom_markers[idx]
cur_marker.hide()
del self._multi_atom_markers[idx]
del self._markername_to_marker_map[marker.name]
@maestro_callback.workspace_changed
def _updateMarkers(self, what_changed):
"""
Show, hide, and delete markers after the workspace has been updated
based on whether marked atoms are included in the workspace. Also sync
python and maestro marker models.
"""
if what_changed in (maestro.WORKSPACE_CHANGED_EVERYTHING,
maestro.WORKSPACE_CHANGED_APPEND,
maestro.WORKSPACE_CHANGED_CONNECTIVITY):
try:
self._clearInvalidatedJaguarMarkers()
self.showAllJaguarMarkers()
self._syncMarkers()
except project.ProjectException:
# If the project was just closed, don't do anything
pass
def _clearInvalidatedJaguarMarkers(self):
"""
Clear all markers that contain atoms from invalidated entries. An entry
is invalidated if it has been deleted or if the number of atoms it
contains has changed.
"""
# Determine all entries that are invalidated
invalidated_eids = set()
for (entry_id, prev_num_atoms) in self._marked_eid_lengths.items():
new_num_atoms = self._calcEntryAtomTotal(entry_id)
if prev_num_atoms != new_num_atoms:
invalidated_eids.add(entry_id)
# Remove all markers that refer to atoms from invalidated entries
marked_eids = set()
for cur_hashable in list(self._markers):
cur_eids = self._eidsFromHashable(cur_hashable)
if cur_eids & invalidated_eids:
self._markers[cur_hashable].hide()
del self._markers[cur_hashable]
else:
marked_eids.update(cur_eids)
# Remove the _marked_eid_lengths entry for structures that no longer
# have any markers
for cur_eid in list(self._marked_eid_lengths):
if cur_eid not in marked_eids:
del self._marked_eid_lengths[cur_eid]
def _eidsFromHashable(self, hashable):
"""
Get a set of entry ids from a key to self._markers
:param hashable: An key to self._markers
:type hashable: tuple
:return: A set of entry ids
:rtype: set
"""
return set(hashable[::2])
[docs] def showAllJaguarMarkers(self):
"""
Show all `schrodinger.maestro.markers._BaseMarker` markers for which
all marked atoms are in the workspace. Hide all other markers.
"""
if not self._markers:
# If there are no markers set, don't bother to scan the workspace
return
workspace_eids = maestro.get_included_entry_ids()
for (cur_hashable, cur_marker) in self._markers.items():
cur_eids = self._eidsFromHashable(cur_hashable)
if cur_eids.issubset(workspace_eids):
cur_marker.update()
else:
cur_marker.hide()
[docs] def showAllMarkers(self):
"""
Set all `schrodinger.maestro.markers.Marker` markers to be shown
if the relevant atoms are in the workspace. These markers are hidden
automatically by Maestro when atoms are excluded.
"""
for cur_marker in self._multi_atom_markers:
cur_marker.show()
def _showAll(self):
"""
Calls showAllMarkers() and showAllJaguarMarkers().
"""
self.showAllMarkers()
self.showAllJaguarMarkers()
def _syncMarkers(self):
"""
Sync the python and maestro marker models. Suppose a marker name is not
registered in maestro model then it should be removed from python model
as well (this will happen during the time of undo operation in maestro).
And if a marker name is present in maestro but not in python model then
the marker should be brought back in python model (this will happen
during the time of redo operation in maestro).
"""
maestro_hub = maestro_ui.MaestroHub.instance()
# marker names currently stored in maestro
maestro_markernames = maestro_hub.fetchCurrentMarkerModelNames()
for markername, marker in self._markername_to_marker_map.items():
# marker is present in python but not in maestro
if (marker in self._multi_atom_markers and
markername not in maestro_markernames):
self._multi_atom_markers.remove(marker)
# marker is present in maestro but not in python
elif (markername in maestro_markernames and
marker not in self._multi_atom_markers):
self._multi_atom_markers.append(marker)
[docs] def hideAllJaguarMarkers(self):
"""
Hide all `schrodinger.maestro.markers._BaseMarker` markers
for this panel
"""
for cur_marker in self._markers.values():
cur_marker.hide()
[docs] def hideAllMarkers(self):
"""
Hide all `schrodinger.maestro.markers.Marker` markers
for this panel.
"""
for cur_marker in self._multi_atom_markers:
cur_marker.hide()
def _hideAll(self):
"""
Calls hideAllMarkers() and hideAllJaguarMarkers().
"""
self.hideAllMarkers()
self.hideAllJaguarMarkers()
[docs] def removeAllJaguarMarkers(self):
"""
Remove all markers `schrodinger.maestro.markers._BaseMarker` markers
from this panel
"""
self.hideAllJaguarMarkers()
self._markers = {}
self._marked_eid_lengths = {}
[docs] def removeAllJaguarMarkersForEntry(self, eid):
"""
Remove all markers for the specified entry id from this panel
:param eid: The entry id to remove markers for
:type eid: str
"""
for cur_hashable in list(self._markers):
cur_eids = self._eidsFromHashable(cur_hashable)
if eid in cur_eids:
self._markers[cur_hashable].hide()
del self._markers[cur_hashable]
self._marked_eid_lengths.pop(eid, None)
[docs] def removeAllMarkers(self):
"""
Remove all `schrodinger.maestro.markers.Marker` markers
from this panel.
"""
self.hideAllMarkers()
self._multi_atom_markers = []
[docs] @_requires_maestro
def getAllJaguarMarkers(self):
"""
Get all markers._BaseMarker currently loaded into the panel
:return: An iterator of markers._BaseMarker
:rtype: iterator
"""
return self._markers.values()
[docs] def getAllMarkers(self):
"""
Get all markers.Marker loaded into the panel
:return: list(markers.Marker)
:rtype: list
"""
return self._multi_atom_markers