"""
A collection of numerous useful functions and classes for working with
Maestro.
Copyright Schrodinger, LLC. All rights reserved.
"""
# Contributors: Yujie Wu
import schrodinger.infra.mm as mm
import schrodinger.project as project
from schrodinger import structure
from schrodinger.structutils import analyze
from schrodinger.structutils import color as color_module
from . import util
_version = "$Revision: 1.29 $"
try:
import schrodinger.maestro as mae
except:
mae = None # for pychecker
UNIQUE_ID_PROP = "i_desmond_internal_entry_id"
[docs]def get_atoms_from_uid(workspace_st, uid):
extract_atoms = set()
molecules = set()
for atom in workspace_st.atom:
if atom.property.get(UNIQUE_ID_PROP) == uid:
extract_atoms.add(atom.index)
molecules.add(atom.molecule_number)
#print "atoms before adding mol atoms:", extract_atoms
#print "mols:", molecules
# Add any other atoms from this molecule (which the user might have added):
for molnum in molecules:
extract_atoms.update(workspace_st.molecule[molnum].getAtomIndices())
#print "atoms after adding mol atoms:", extract_atoms
return extract_atoms
[docs]def create_atom_tag(st, tag, value=None):
"""
Creates one or more atom properties for all atoms in CT.
:param tag: The property name as a string or a list of property
names in a list of strings.
:param value: The value for the corresponding atom property. If
'tag' is a list, then 'value' should also be a list.
"""
handle = st.handle
if (tag.__class__ == str):
tag = [
tag,
]
if (value is not None):
value = [
value,
]
if value is None:
for t in tag:
mm.mmct_atom_property_set(handle, t, mm.MMCT_ATOMS_ALL)
else:
for t, v in zip(tag, value):
mm.mmct_atom_property_set(handle, t, mm.MMCT_ATOMS_ALL, v)
[docs]def mark_atom(st, tag, value, atom=[], asl=""): # noqa: M511
"""
Sets value (given by the 'value' argument) to atom property (given
by the 'tag' argument) for selected atoms in the CT.
The 'tag' and 'value' arguments can be lists, but they must be both
lists or both not.
The selected atoms are given by the 'atom' argument or by an ASL
expression ('asl'). Either 'atom' or 'asl' should be provided; if
both are provided, 'asl' will be ignored.
"""
at_ = st.atom
if (atom == []):
if not asl:
return
mae.command("workspaceselectionreplace " + asl)
atom = mae.selected_atoms_get()
mae.command("workspaceselectionreplace atom.num 0")
if (tag.__class__ == str):
tag = [
tag,
]
value = [
value,
]
for i_atom in atom:
for t, v in zip(tag, value):
at_[int(i_atom)].property[t] = v
[docs]def get_entry_id(entry_name):
"""
Returns the entry ID of the entry whose entry name being 'entry_name',
or None if there is no such an entry.
Note that if there are more than one entries have the same name of
'entry_name', the entry id of the first one will be returned.
"""
pt = mae.project_table_get()
pt.selectRows("entry_name " + entry_name)
for row in pt.selected_rows:
return row["s_m_entry_id"]
return None
[docs]def duplicate_entry(entry_id):
"""
Makes a copy of an entry, and returns the ID of the copy entry (i.e.,
the new entry).
:param entry_id: The ID of the entry to be copied. If this entry does
not exist, nothing will be copied, and None will be returned.
"""
pt = mae.project_table_get()
old = pt.last_added_entry["s_m_entry_id"]
mae.command("entryduplicate entry_id " + str(entry_id))
pt = mae.project_table_get()
new = pt.last_added_entry["s_m_entry_id"]
if (new != old):
return new
return None
[docs]def get_entryname_prefix(suggested_prefix=""):
"""
Returns a string to be used as a prefix of entry names so that the
resulted entry name is guaranteed to be different from those existing at
the time of getting this prefix.
Note a side-effect of this function: the row selection on the project
table will be changed.
:param suggested_prefix: A string from which the returned prefix will be
derived.
"""
while (True):
if ("" == suggested_prefix):
suggested_prefix = util.random_string(3)
pt = mae.project_table_get()
for i in range(128):
pt.selectRows(project.REPLACE,
esl="entry_re " + suggested_prefix + "*")
if (0 == pt.getSelectedRowTotal()):
return suggested_prefix
else:
suggested_prefix += "_"
else:
suggested_prefix = ""
[docs]def noautofit_call(callback):
"""
Supresses autofit preference when calling to 'callback'. The original
preference is recovered after calling to 'callback'.
TODO: This function should be removed, if its functionality is really
needed, it should use the new autofit preferences, and it should be a
context manager instead of a callback function
"""
return callback()
[docs]def autoprojectsync_call(callback):
"""
Temporarily sets 'projectsync' preference to 'auto' before calling to
'callback'. The original preference is recovered after calling to
'callback'.
"""
projectsync = mae.get_command_option("prefer", "projectsync")
if (projectsync == "auto"):
return callback()
mae.command("prefer projectsync=auto")
ret = callback()
mae.command("prefer projectsync=" + projectsync)
return ret
[docs]def nofitgrow_call(callback):
"""
Temporarily sets 'fitgrow' preference to 'false' before calling to
'callback'. The original preference is recovered after calling to
'callback'.
"""
fitgrow = mae.get_command_option("prefer", "fitgrow")
if (fitgrow == "false"):
return callback()
mae.command("prefer fitgrow=false")
ret = callback()
mae.command("prefer fitgrow=" + fitgrow)
return ret
[docs]def depreference_call(callback):
"""
"""
return noautofit_call(
lambda: nofitgrow_call(lambda: autoprojectsync_call(callback)))
[docs]class NewEntry(object):
"""
A class for representing a structure that can be included/excluded
in the Workspace.
"""
# This is used to track the last unique identifier that was assigned to an entry:
last_entry_unique_id = 0
[docs] def __init__(self, ct):
self.ct = ct.copy()
NewEntry.last_entry_unique_id += 1
self.uid = NewEntry.last_entry_unique_id
self._updateProp()
self.name = self.ct.title
# __init__
def _updateProp(self):
for atom in self.ct.atom:
atom.property[UNIQUE_ID_PROP] = self.uid
[docs] def asl(self):
return "(atom.%s %i)" % (UNIQUE_ID_PROP, self.uid)
[docs] def set_readonly(self, readonly=True):
"""
Sets this entry to be readonly if 'readonly' is True.
If 'readonly' is False, sets this entry to be not readonly.
"""
[docs] def set_deletable(self, deletable=True):
"""
Sets this entry to be deletable if 'deletable' is True.
If 'deletable' is False, sets this entry to be undeletable.
"""
[docs] def updateToWorkspace(self):
self._updateProp()
workspace_st = mae.workspace_get()
delete_atoms = get_atoms_from_uid(workspace_st, self.uid)
workspace_st.deleteAtoms(delete_atoms)
workspace_st.extend(self.ct)
mae.workspace_set(workspace_st)
[docs] def show(self):
"""
Shows the CT in Maestro workspace.
"""
self._updateProp()
workspace_st = mae.workspace_get()
delete_atoms = get_atoms_from_uid(workspace_st, self.uid)
workspace_st.deleteAtoms(delete_atoms)
workspace_st.extend(self.ct)
mae.workspace_set(workspace_st)
[docs] def show_only(self):
"""
Shows the CT in Maestro workspace and hides other CTs.
"""
self._updateProp()
mae.workspace_set(self.ct)
mae.redraw()
[docs] def hide(self):
workspace_st = mae.workspace_get()
delete_atoms = get_atoms_from_uid(workspace_st, self.uid)
workspace_st.deleteAtoms(delete_atoms)
mae.workspace_set(workspace_st)
[docs] def is_shown(self):
"""
Returns True if this entry is already included into the Workspace.
"""
workspace_st = mae.workspace_get()
for atom in workspace_st.atom:
if atom.property.get(UNIQUE_ID_PROP) == self.uid:
return True
return False
[docs] def commit(self):
"""
Commits any belonging change in the current workspace to this entry.
"""
workspace_st = mae.workspace_get()
extract_atoms = get_atoms_from_uid(workspace_st, self.uid)
#print "COMITTING ENTRY; atoms:", extract_atoms
if extract_atoms:
self.ct = workspace_st.extract(extract_atoms)
self._updateProp()
[docs] def rename(self, name):
"""
Changes the entry name to 'name'.
"""
self.name = name
#pt = mae.project_table_get()
#pt.selectRows( esl = "entry_id " + self.id )
#for row in pt.selected_rows :
# row["Title"] = self.name
# break
[docs] def get_ct(self, copy=False):
"""
Returns a 'Structure' object of the CT of this entry.
"""
# FIXME incorporate changes from the Workspace?
if copy:
ct = self.ct.copy()
else:
ct = self.ct
return ct
[docs] def write(self, fname, mode=""):
"""
Writes the CT to a file, whose name is given by 'fname'. The format
of the file is determined by the extension name of 'fname'.
:param mode: If set to 'a', then instead of overwriting, this
function will append the structure to the end of the file.
"""
ct = self.get_ct()
if ("a" == mode or "A" == mode):
ct.append(fname)
else:
ct.write(fname)
[docs] def deleteAtomsByAsl(self, asl):
atoms = analyze.evaluate_asl(self.ct, asl)
self.ct.deleteAtoms(atoms)
[docs]class NewEntryGroup(object):
"""
A class for representing a group of Structures that can be
included in or excluded from the Workspace as requested.
"""
[docs] def __init__(self, structures, group_name=None):
"""
Creates an entry group with a given list of CTs.
:param struc: Must be list of `schrodinger.structure.Structure`
objects.
"""
self.name = group_name
self.entry = []
self._index = 0
if None == self.name:
self.name = get_entryname_prefix()
for ct in structures:
self.entry.append(NewEntry(ct))
self._index += 1
# __init__
[docs] def dtor(self):
"""
Deletes all entries in this entry-group.
"""
self.entry = []
self._index = 0
[docs] def duplicate(self, entry):
"""
Duplicates an entry, and the new copy will be put in
this entry-group.
"""
self.entry.append(NewEntry(entry.ct))
self._index += 1
return self.entry[-1]
[docs] def delete(self, entry):
"""
Deletes one or more entries from the group.
Deletion happens only when the entry is in this entry group.
:param entry: Can be either a single `Entry` object or a list of
`Entry` objects.
"""
try:
e0 = entry[0]
except TypeError:
entry = [
entry,
]
except IndexError:
return
for e in entry:
self.entry.remove(e)
del e
[docs] def is_shown(self):
"""
Returns True if all entries in this group are included into the
Workspace, False otherwise.
"""
workspace_st = mae.workspace_get()
included_entries = set()
for atom in workspace_st.atom:
uid = atom.property.get(UNIQUE_ID_PROP)
included_entries.add(uid)
for e in self.entry:
if e.uid not in included_entries:
return False
return True
[docs] def is_shown_only(self):
"""
Returns True if all entries in this group and only entries in this
group are included into the Workspace.
"""
workspace_st = mae.workspace_get()
included_entries = set()
for atom in workspace_st.atom:
uid = atom.property.get(UNIQUE_ID_PROP)
if uid is None:
return False
return True
[docs] def show(self):
"""
Shows all entries in this group onto Maestro workspace if
'should_show' is True, hides them if it is False.
"""
workspace_st = mae.workspace_get()
uids_in_workspace = set()
for atom in workspace_st.atom:
uid = atom.property.get(UNIQUE_ID_PROP)
if uid is not None:
uids_in_workspace.add(uid)
# FIXME temp code, show_only instead:
workspace_st = None
for e in self.entry:
e._updateProp()
#print " showing entry", e.uid, "atoms:", e.ct.atom_total
if workspace_st is None:
workspace_st = e.ct.copy()
else:
workspace_st.extend(e.ct)
mae.workspace_set(workspace_st)
[docs] def hide(self):
workspace_st = mae.workspace_get()
delete_atoms = set()
for e in self.entry:
delete_atoms.update(get_atoms_from_uid(workspace_st, e.uid))
workspace_st.deleteAtoms(delete_atoms)
[docs] def show_only(self):
"""
Shows all entries in this group onto Maestro workspace and hides
other entries.
"""
combined_st = None
for e in self.entry:
if combined_st is None:
combined_st = e.ct.copy()
else:
combined_st.extend(e.ct)
mae.workspace_set(combined_st)
[docs] def commit(self):
"""
Commits any belonging changes in the Workspace to their
corresponding entries that are in this group.
"""
workspace_st = mae.workspace_get()
for e in self.entry:
extract_atoms = get_atoms_from_uid(workspace_st, e.uid)
if extract_atoms:
e.ct = workspace_st.extract(extract_atoms)
[docs] def get_ct(self, copy=False):
"""
Returns a list of 'Structure' objects of the CTs of this group.
"""
self.commit()
struc = []
for e in self.entry:
struc.append(e.get_ct(copy))
return struc
[docs] def getWorkspaceEntry(self):
"""
Return the entry that is currently in the Workspace. Or None.
"""
workspace_st = mae.workspace_get()
uids_in_workspace = set()
for atom in workspace_st.atom:
uid = atom.property.get(UNIQUE_ID_PROP)
if uid is not None:
uids_in_workspace.add(uid)
# NOTE: sometimes len(uids_in_workspace) will be >1. This is fine.
for e in self.entry:
if e.uid in uids_in_workspace:
return e
return None
[docs] def write(self, fname, mode=""):
"""
Writes all CTs to a file, whose name is given by 'fname'.
The format of the file is determined by the extensiion name of
'fname'.
:param mode: If set to 'a', then instead of overwriting this
function will append the structure to the end of the file.
"""
for e in self.entry:
e.write(fname, mode)
mode = "a"
[docs] def printStatus(self):
workspace_st = mae.workspace_get()
atom_count_by_uid = {}
for atom in workspace_st.atom:
uid = atom.property.get(UNIQUE_ID_PROP)
#if uid != None:
if True:
try:
atom_count_by_uid[uid] += 1
except KeyError:
atom_count_by_uid[uid] = 1
[docs] def getAtomEntry(self, atom_obj):
"""
Return the entry object (or None) that the given
atom object is associated with.
"""
uid = atom_obj.property.get(UNIQUE_ID_PROP)
if uid is None:
return None
for e in self.entry:
if e.uid == uid:
return e
return None
[docs]class Workspace(object):
"""
A class for representing the Workspace.
"""
_ct = None
[docs] def __init__(self):
"""
Does nothing.
"""
def _show_entry(self, eid):
"""
Displays the entry or entries as given by 'eid' in the Workspace.
Note that 'eid' can be either an integer that is the ID of a single
entry or a list of integers that are the IDs of a few entries.
"""
try:
s = ""
for i in eid:
s += str(i) + " "
except TypeError:
s = str(i)
mae.command("entrywsinclude entry_id " + s)
[docs] def show_entry(self, eid):
autoprojectsync_call(lambda: self._show_entry(eid))
def _show_only_entry(self, eid):
"""
Displays only the entry or entries as given by 'eid' in the
Workspace. The meaning of 'eid' is similar to that in `_show_enty`.
"""
try:
s = ""
for i in eid:
s += str(i) + " "
except TypeError:
s = str(i)
mae.command("entrywsincludeonly entry_id " + s)
[docs] def show_only_entry(self, eid):
autoprojectsync_call(lambda: self._show_only_entry(eid))
def _grab(self):
"""
Gets the CTs and entry IDs for the entries currently displayed in
the Workspace.
This function returns a tuple of two elements. The 1st element is a
list of `schrodinger.structure.Structure` objects, the 2nd one is a
list of entry IDs.
"""
struc = []
eid = []
pt = mae.project_table_get()
for row in pt.included_rows:
eid.append(row["s_m_entry_id"])
# Is there a scratch CT?
eid_str = ""
if (eid != []):
for i in eid:
eid_str += str(i) + " "
mae.command("entrywsexclude entry_id " + eid_str)
ct = self.ct
if (ct.atom_total > 0):
eid.insert(0, "Scratch")
struc.insert(0, ct.copy())
if (eid_str != ""):
mae.command("entrywsinclude entry_id " + eid_str)
for row in pt.included_rows:
st = row.getStructure()
struc.append(st)
return struc, eid
[docs] def grab(self):
return autoprojectsync_call(lambda: self._grab())
def _paint_asl(self, asl, color):
"""
Colors the selected atoms as given by 'asl'.
"""
if ("scheme." == color[:7]):
mae.command("colorscheme scheme=\"" + color[7:] + "\" " + asl)
else:
mae.command("coloratom color=\"" + color + "\" " + asl)
[docs] def paint_asl(self, asl, color):
autoprojectsync_call(lambda: self._paint_asl(asl, color))
[docs] def paint_mol(self, i_mol, color):
"""
Colors a molecule as given by the index of the molecule ('i_mol').
"""
self.paint_asl("mol.num " + str(i_mol), color)
def _ball_atom(self, atom=[], asl="", should_ball=True): # noqa: M511
"""
Represents selected atoms in the ball-and-stick mode.
The selected atoms are chosen by either 'atom' or 'asl' argument. If
both are given, 'asl' will be ignored.
:param atom: A list of atom indices.
:param asl: An ASL expression (as a string).
:param should_ball: If True, then the selected atoms will be drawn
in the wire mode.
"""
if ([] == atom):
if ("" == asl):
return
else:
asl = "atom.num 0"
for i in atom:
asl += " " + str(i)
try:
if (should_ball):
mae.command("repatom rep = ballnstick " + asl)
else:
mae.command("repatombonds " + asl)
except mae.MaestroCommand:
pass
[docs] def ball_atom(self, atom=[], asl="", should_ball=True): # noqa: M511
autoprojectsync_call(lambda: self._ball_atom(atom, asl, should_ball))
[docs] def display_asl(self, asl, should_display=True):
"""
Display the selected atoms (as given by the ASL expression 'asl') in
the Workspace.
If 'should_display' is False, then the selected atoms will be
undisplayed.
"""
if (should_display):
mae.command("displayatom " + asl)
else:
mae.command("undisplayatom " + asl)
[docs] def display_only_asl(self, asl, should_display=True):
"""
Similar to `display_asl`, but non-selected atoms will be hidden.
If 'should_display' is False, then the selected atoms will be
undisplayed without any effect on the non-selected atoms.
"""
if (should_display):
mae.command("displayonlyatom " + asl)
else:
mae.command("undisplayatom " + asl)
def _get_ct(self):
"""
"""
Workspace._ct = mae.workspace_get()
return Workspace._ct
ct = property(
_get_ct,
None,
doc="Readonly property: returns the workspace as a 'Structure' object.")
def _commit(self):
"""
Commits changes on 'self.ct' to the Workspace.
"""
if (Workspace._ct is not None):
mae.workspace_set(Workspace._ct, True, False)
Workspace._ct = None
[docs] def commit(self):
autoprojectsync_call(lambda: self._commit())
[docs] def get_atom(self, asl):
"""
Returns a list of indices of atoms as selected by the 'asl'.
"""
mae.command("workspaceselectionreplace " + asl)
atom = mae.selected_atoms_get()
mae.command("workspaceselectionreplace atom.num 0")
return atom
def _create_atom_tag(self, tag, value=None):
"""
Creates one or more atom properties for all atoms included into the
Workspace.
The change will be committed to the included entries.
:param tag: The property name as a string or a list of property
names in a list of strings.
:param value: The value for the corresponding atom property. If
'tag' is a list, then 'value' should also be a list.
"""
ct = self.ct
create_atom_tag(ct, tag, value)
self.commit()
[docs] def create_atom_tag(self, tag, value=None):
autoprojectsync_call(lambda: self._create_atom_tag(tag, value))
def _delete_atom_tag(self, tag):
"""
Deletes one or more atom properties for all atoms included into the
Workspace.
The change will be committed to the included entries.
:param tag: The property name as a string or a list of property
names in a list of strings.
"""
if (tag.__class__ == str):
tag = [
tag,
]
for t in tag:
self.ct.deletePropertyFromAllAtoms(t)
self.commit()
[docs] def delete_atom_tag(self, tag):
autoprojectsync_call(lambda: self._delete_atom_tag(tag))
def _mark_atom(self, tag, value, atom=[], asl=""): # noqa: M511
"""
Sets value (given by the 'value' argument) to atom property (given
by the 'tag' argument) for selected atoms in the Workspace.
The 'tag' and 'value' arguments can be lists, but they must be both
lists or both not.
The selected atoms are given by the 'atom' argument or by an ASL
expression ('asl'). Either 'atom' or 'asl' should be provided; if
both are provided, 'asl' will be ignored.
The change will be committed to the included entries.
"""
ct = self.ct
mark_atom(ct, tag, value, atom, asl)
self.commit()
[docs] def mark_atom(self, tag, value, atom=[], asl=""): # noqa: M511
autoprojectsync_call(lambda: self._mark_atom(tag, value, atom, asl))
def _run():
tsuite = util.TestSuite("chorus.maestro")
prefix = get_entryname_prefix("testentryname")
tsuite << (prefix == "testentryname", "get_entryname_prefix 1", prefix)
# New: works on a CT instead of the Workspace:
[docs]def ball_atom(st, atoms, should_ball=True):
if type(atoms) != type([]):
# Assume an ASL was given
atoms = analyze.evaluate_asl(st, atoms)
if should_ball:
st.applyStyle(structure.ATOM_BALLNSTICK, structure.BOND_BALLNSTICK,
atoms)
else:
st.applyStyle(structure.ATOM_NOSTYLE, structure.BOND_WIRE, atoms)
# constants ATOM_NOSTYLE, ATOM_CIRCLE, ATOM_CPK, ATOM_BALLNSTICK. Default
# is ATOM_BALLNSTICK.
# constants BOND_NOSTYLE, BOND_WIRE, BOND_TUBE, BOND_BALLNSTICK. Default
# is BOND_BALLNSTICK.
[docs]def paint_asl(st, asl, color):
"""
Colors the selected atoms as given by 'asl'.
"""
if ("scheme." == color[:7]):
element_green_scheme = color_module.get_color_scheme(color[7:])
element_green_scheme.apply(st, asl)
else:
atoms = analyze.evaluate_asl(st, asl)
for ai in atoms:
st.atom[ai].color = color
#EOF