"""
Functions and classes specific to energy trajectory plots.
"""
import os
import numpy as np
from schrodinger import get_maestro
from schrodinger.Qt.QtCore import Qt
from schrodinger.Qt import QtCore
from schrodinger.trajectory.trajectory_gui_dir import traj_plot_models
from schrodinger.ui.qt import basewidgets
from schrodinger.ui.qt import table_helper
from schrodinger.ui.qt import filedialog
from schrodinger.ui.qt import atomselector
from schrodinger.ui.qt.appframework2 import jobnames
from schrodinger.Qt import QtWidgets
from schrodinger.utils import sea
from schrodinger.trajectory.trajectory_gui_dir import job_settings_dialog_ui
from schrodinger.trajectory.trajectory_gui_dir import custom_asl_sets_dialog_ui
from schrodinger.trajectory.trajectory_gui_dir import custom_substructure_sets_dialog_ui
from schrodinger.structutils import analyze
# NOTE: These ASLs are temporary, and will be removed when Structure Hierarchy
# is exposed to Python (MAE-45958):
METAL_ION_ASL = '(metals OR ions)'
OTHER_SOLVENT_ASL = '(solvent AND NOT water)'
OTHER_HETS_ASL = '(NOT (ligand OR protein OR solvent OR metals OR ions OR nucleic_acids))'
# TODO replace with aliases once SHARED-8150 is implemented
DNA_ASL = '(res. " DA " OR res. " DC " OR res. " DG " OR res. " DT ")'
RNA_ASL = '(res. " C " OR res. " U " OR res. " G " OR res. " A ")'
# ASLs for "Substructure Sets (Grouped)" menu item:
GROUP_ASLS = {
'Proteins': 'protein',
'Ligands': 'ligand',
'Water': 'water',
'Other solvent': OTHER_SOLVENT_ASL,
'Metals and ions': METAL_ION_ASL,
'Other hets': OTHER_HETS_ASL,
'Membranes': 'membrane',
'DNA': DNA_ASL,
'RNA': RNA_ASL,
}
# TODO: Separate out DNA and RNA into separate sets (currently not easy
# until until SHARED-8150 is implemented).
maestro = get_maestro()
[docs]def iter_children(item):
"""
Iterate over all children of a QTreeWidgetItem.
"""
for i in range(item.childCount()):
yield item.child(i)
[docs]def iter_tree_items(tree_widget):
"""
Iterate over all items (both top level and their children) of a QTreeWidget.
"""
for parent in iter_children(tree_widget.invisibleRootItem()):
yield parent
for child in iter_children(parent):
yield child
[docs]def find_all_grouped_sets(st):
"""
Return a dict of ASLs, one for each group, for all substructure set groups
found in the structure.
"""
sets = {}
for name, asl in GROUP_ASLS.items():
if analyze.evaluate_asl(st, asl):
sets[name] = asl
return sets
[docs]def find_all_individual_sets(st):
"""
Return a list of ASLs, one for each substructure set (ungrouped). Each
ligand molecule and protein chain will be assigned its own set. Other
substructures are still grouped together by type.
"""
groups = [find_ligands(st), find_proteins(st), find_other(st)]
sets = {name: asl for items in groups for name, asl in items}
return sets
[docs]def find_sets_for_all_molecules(st):
"""
Assign each molecule to its own set. Useful only for structure with few
molecules (no solvent)
"""
return {
f'Molecule {mol.number}': f'mol.n {mol.number}' for mol in st.molecule
}
[docs]def find_ligands(st):
"""
Return a list of (ligand name, ASL) for all ligands found in the structure.
"""
ligands_list = []
for i, ligand in enumerate(analyze.find_ligands(st), start=1):
atoms = ligand.atom_indexes
asl = analyze.generate_asl(st, atoms)
residues = {res for res in ligand.st.residue}
name = f'Ligand {i}'
if len(residues) == 1:
res = residues.pop()
chain = res.chain.replace(' ', '_')
name = f'{chain}:{res.pdbres}{res.resnum}{res.inscode.strip()} (Ligand {i})'
else:
# Multi-residue ligand
name = f'Ligand {i}'
ligands_list.append((name, asl))
return ligands_list
[docs]def find_proteins(st):
"""
Find protein chains in the given structure.
"""
non_protein_atoms = set(analyze.evaluate_asl(st, 'NOT protein'))
# Remove peptide ligands also:
for ligand in analyze.find_ligands(st):
non_protein_atoms.update(ligand.atom_indexes)
protein_chain_list = []
for chain in st.chain:
atoms = set(chain.getAtomIndices()) - non_protein_atoms
if atoms:
asl = analyze.generate_asl(st, atoms)
_name = chain.name.replace(' ', '_')
name = f'Chain {_name} (Protein)'
protein_chain_list.append((name, asl))
return protein_chain_list
[docs]def find_solvents(st):
"""
Find solvents (including water) in the given structure.
"""
water_present = bool(analyze.evaluate_asl(st, 'water'))
other_solvents_present = bool(analyze.evaluate_asl(st, OTHER_SOLVENT_ASL))
solvents = []
if water_present:
solvents.append(('Water', 'water'))
if other_solvents_present:
solvents.append(('Other Solvents', OTHER_SOLVENT_ASL))
# TODO remove grouping if only one of above is present
return solvents
[docs]def find_other(st):
"""
Find metals, ions, membranes, DNA, and RNA in the given structure.
"""
other_list = []
for name, asl in {
'Metals / Ions': METAL_ION_ASL,
'Other Hets': OTHER_HETS_ASL,
'Membranes': 'membrane',
'DNA': DNA_ASL,
'RNA': RNA_ASL,
}.items():
atoms = analyze.evaluate_asl(st, asl)
if atoms:
other_list.append((name, asl))
return other_list
[docs]class CustomAslSetRow:
"""
Class representing a row in the custom ASL dialog's sets table.
"""
[docs] def __init__(self, name, asl, num_atoms):
self.name = name
self.asl = asl
self.num_atoms = num_atoms
class CustomAslSetsColumns(table_helper.TableColumns):
"""
Columns of the CustomAslSetsTableModel table.
"""
SetName = table_helper.Column('Set Name')
ASL = table_helper.Column('ASL')
NumAtoms = table_helper.Column('# Atoms')
[docs]class CustomAslSetsTableModel(table_helper.RowBasedTableModel):
"""
Model for the table at the bottom of the "Define Custom ASL Sets" dialog.
"""
Column = CustomAslSetsColumns
ROW_CLASS = CustomAslSetRow
@table_helper.data_method(Qt.DisplayRole)
def _displayData(self, col_idx, row_obj, role):
# See table_helper.data_method for documentation
if col_idx == self.Column.SetName:
return row_obj.name
elif col_idx == self.Column.ASL:
return row_obj.asl
elif col_idx == self.Column.NumAtoms:
return row_obj.num_atoms
@table_helper.data_method(Qt.ToolTipRole)
def _displayToolTip(self, col_idx, row_obj, role):
# See table_helper.data_method for documentation
if col_idx == self.Column.ASL:
return row_obj.asl
[docs]class CustomAslSetsDialog(basewidgets.BaseDialog):
"""
Dialog for defining sets via custom ASLs.
:cvar setsDefined: A signal emitted when OK button is pressed after
defining new sets.
"""
ui_module = custom_asl_sets_dialog_ui
setsDefined = QtCore.pyqtSignal(dict)
[docs] def __init__(self, st, parent=None):
super().__init__(parent)
self.st = st
self.setWindowTitle('Define Custom ASL Sets')
[docs] def initSetOptions(self):
super().initSetOptions()
self.std_btn_specs[self.StdBtn.Ok] = self.okPressed
self.std_btn_specs[self.StdBtn.Cancel] = None
[docs] def initSetUp(self):
super().initSetUp()
self.atom_selector = atomselector.AtomSelector(
self,
label="ASL:",
pick_text="Pick atom in Workspace",
show_asl=True,
show_all=True,
show_select=True,
show_previous=False,
show_selection=False,
append_mode=True,
show_markers=False,
show_pick=True,
default_pick_mode=atomselector.PICK_MOLECULES,
show_plus=True,
)
self.ui.atom_selector_layout.addWidget(self.atom_selector)
self.atom_selector.picksite_wrapper.start()
self.model = CustomAslSetsTableModel()
self.ui.sets_table_view.setModel(self.model)
self.ui.sets_table_view.setWordWrap(False)
self.ui.sets_table_view.setColumnWidth(CustomAslSetsColumns.SetName,
100)
self.ui.sets_table_view.setColumnWidth(CustomAslSetsColumns.ASL, 200)
self.ui.sets_table_view.setColumnWidth(CustomAslSetsColumns.NumAtoms,
60)
self.ui.sets_table_view.setSelectionBehavior(
QtWidgets.QAbstractItemView.SelectRows)
self.ui.sets_table_view.selectionModel().selectionChanged.connect(
self.setTableSelectionChanged)
self.ui.add_set_btn.clicked.connect(self._addSet)
self.ui.delete_set_btn.clicked.connect(self._removeSelectedSet)
self.createUniqueSetName()
self.resize(400, 500)
[docs] def createUniqueSetName(self):
used_names = {row.name for row in self.model.rows}
prev_name = self.ui.set_name_le.text()
new_name = jobnames.update_jobname(prev_name,
"Set",
name_list=used_names)
self.ui.set_name_le.setText(new_name)
self.ui.set_name_le.repaint()
def _addSet(self):
"""
Add set defined by the atom selector frame to the bottom table.
"""
set_name = self.ui.set_name_le.text().strip()
if not set_name:
self.warning("Please enter a set name")
return
asl = self.atom_selector.getAsl()
if not asl.strip():
self.warning('Please define an ASL')
return
if not analyze.validate_asl(asl):
self.warning('Specified ASL is invalid')
return
atoms = analyze.evaluate_asl(self.st, asl)
if not atoms:
self.warning('Specified ASL did not match any atoms')
return
if self.model.rowCount():
existing_asl = ' OR '.join(
('%s' % row.asl for row in self.model.rows))
matched_atoms = set(analyze.evaluate_asl(self.st, existing_asl))
if matched_atoms.intersection(atoms):
self.warning(
"Specified ASL matches some of the atoms that were already assigned to another set."
)
return
row = CustomAslSetRow(set_name, asl, len(atoms))
self.model.appendRowObject(row)
self.atom_selector.setAsl('')
self.createUniqueSetName()
[docs] def setTableSelectionChanged(self):
enable = bool(self.ui.sets_table_view.selectionModel().selectedRows())
self.ui.delete_set_btn.setEnabled(enable)
def _removeSelectedSet(self):
"""
Remove selected set from the bottom table.
"""
self.model.removeRowsByIndices(
self.ui.sets_table_view.selectionModel().selectedRows())
[docs] def okPressed(self):
count = self.model.rowCount()
if count == 0:
self.warning('Please define sets to use')
return False
if count > 10:
self.warning('Maximum of 10 sets are supported')
return False
sets = {row.name: row.asl for row in self.model.rows}
self.close()
self.setsDefined.emit(sets)
[docs]class CustomSubstructureSetsDialog(basewidgets.BaseDialog):
"""
Dialog for defining which substructure sets to use for generating energy
plots. Substructure sets are automatically detected in a manner similar to
what is used by Structure Hierarchy. User is able to specify up to 10 of
those sets to use for energy interaction calculations.
:cvar setsDefined: A signal emitted when OK button is pressed after
defining new sets.
"""
ui_module = custom_substructure_sets_dialog_ui
setsDefined = QtCore.pyqtSignal(dict)
NAME_COL = 0
ASL_ROLE = Qt.UserRole + 1
IN_USE_ROLE = Qt.UserRole + 2
[docs] def addItem(self, parent, name, asl):
"""
Create a new QTreeWidgetItem name and ASL, and add it to the parent.
"""
item = QtWidgets.QTreeWidgetItem(parent)
item.setText(self.NAME_COL, name)
item.setData(self.NAME_COL, self.ASL_ROLE, asl)
item.setData(self.NAME_COL, self.IN_USE_ROLE, False)
return item
[docs] def addGroup(self, group_name, sets):
"""
Add given sets grouped together under a top-level parent. If only
one item is present in the set, no parent is created, and that set
is added as a top-level item. If no sets are present, nothing is
done.
"""
if not sets:
return
if len(sets) == 1:
name, asl = sets[0]
self.addItem(self.ui.tree_widget, name, asl)
return
# Multiple items present; group them under a parent item:
group_asl = ' OR '.join([asl for name, asl in sets])
tl_item = self.addItem(self.ui.tree_widget, group_name, group_asl)
for name, asl in sets:
self.addItem(tl_item, name, asl)
[docs] def __init__(self, st, parent=None):
super().__init__(parent)
self.setWindowTitle("Define Custom Substructure Sets")
self.ui.tree_widget.setItemsExpandable(False)
self.ui.tree_widget.setSelectionMode(self.ui.tree_widget.MultiSelection)
self.ui.use_sets_widget.setSelectionMode(
self.ui.use_sets_widget.MultiSelection)
self.addGroup('All Ligands', find_ligands(st))
self.addGroup('All Protein Chains', find_proteins(st))
self.addGroup('All Solvents', find_solvents(st))
# These have no top-level parent groups:
for name, asl in find_other(st):
self.addItem(self.ui.tree_widget, name, asl)
self.ui.tree_widget.expandAll()
self.ui.tree_widget.setHeaderHidden(True)
self.ui.tree_widget.setRootIsDecorated(False)
[docs] def initSetOptions(self):
super().initSetOptions()
self.std_btn_specs[self.StdBtn.Ok] = self.okPressed
self.std_btn_specs[self.StdBtn.Cancel] = None
[docs] def initSetUp(self):
super().initSetUp()
self.ui.tree_widget.itemSelectionChanged.connect(
self.treeSelectionChanged)
self.ui.add_btn.clicked.connect(self.addSelectedSets)
self.ui.use_sets_widget.itemSelectionChanged.connect(
self.listSelectionChanged)
self.ui.remove_btn.clicked.connect(self.removeSelectedSets)
[docs] def treeSelectionChanged(self):
tree = self.ui.tree_widget
sel_items = tree.selectedItems()
sel_model = tree.selectionModel()
# Enable all items:
for parent in iter_children(tree.invisibleRootItem()):
for child in iter_children(parent):
child.setDisabled(False)
# De-select children of selected top-level items
for item in sel_items:
for child in iter_children(item):
child.setDisabled(True)
sel_model.select(tree.indexFromItem(child), sel_model.Deselect)
self.ui.add_btn.setEnabled(bool(sel_items))
[docs] def listSelectionChanged(self):
sel_items = self.ui.use_sets_widget.selectedItems()
self.ui.remove_btn.setEnabled(bool(sel_items))
[docs] def addSelectedSets(self):
"""
Called when the "Add >" button is clicked. An item will be added to
the right table for each item selected in the left tree table. Items
in the tree table will be hidden. When a group is added, all of its
children are also hidden.
"""
tree = self.ui.tree_widget
sel_items = tree.selectedItems()
for item in sel_items:
name = item.text(self.NAME_COL)
asl = item.data(self.NAME_COL, self.ASL_ROLE)
list_item = QtWidgets.QListWidgetItem(name)
list_item.setData(self.ASL_ROLE, asl)
self.ui.use_sets_widget.addItem(list_item)
item.setData(self.NAME_COL, self.IN_USE_ROLE, True)
# For any added parent group, un-add any of it's children to avoid
# having same atoms be present in multiple sets. This must be done
# as a separate loop, after children were updated above.
for item in sel_items:
for subitem in iter_children(item):
name = subitem.text(self.NAME_COL)
subitem.setData(self.NAME_COL, self.IN_USE_ROLE, False)
items = self.ui.use_sets_widget.findItems(name, Qt.MatchExactly)
if items:
self.ui.use_sets_widget.takeItem(
self.ui.use_sets_widget.row(items[0]))
self._showOrHideItems()
def _showOrHideItems(self):
"""
Show/hide items in tree widget depending on whether they are in
use or not. Parent groups are also hidden if all their children are
in use.
"""
tree = self.ui.tree_widget
for toplevel in iter_children(tree.invisibleRootItem()):
in_use = toplevel.data(self.NAME_COL, self.IN_USE_ROLE)
if toplevel.childCount() > 0:
all_children_in_use = True
for child in iter_children(toplevel):
child_in_use = child.data(self.NAME_COL, self.IN_USE_ROLE)
if not child_in_use:
all_children_in_use = False
child.setHidden(child_in_use)
in_use = in_use or all_children_in_use
toplevel.setHidden(in_use)
[docs] def removeSelectedSets(self):
sel_items = self.ui.use_sets_widget.selectedItems()
sel_texts = [item.text() for item in sel_items]
for item in sel_items:
self.ui.use_sets_widget.takeItem(self.ui.use_sets_widget.row(item))
tree = self.ui.tree_widget
for item in iter_tree_items(tree):
if item.text(self.NAME_COL) in sel_texts:
item.setData(self.NAME_COL, self.IN_USE_ROLE, False)
self._showOrHideItems()
tree.clearSelection()
[docs] def okPressed(self):
count = self.ui.use_sets_widget.count()
if count == 0:
self.warning('Please define sets to use')
return False
if count > 10:
self.warning('Maximum of 10 sets are supported')
return False
sets = {}
for item in iter_tree_items(self.ui.tree_widget):
if item.data(self.NAME_COL, self.IN_USE_ROLE):
name = item.data(self.NAME_COL, Qt.DisplayRole)
asl = item.data(self.NAME_COL, self.ASL_ROLE)
sets[name] = asl
self.close()
self.setsDefined.emit(sets)
[docs]class EnergyJobSettingsDialog(basewidgets.BaseOptionsDialog):
"""
Dialog for letting the user select the host to run the energy analysis
jobs on. Must be a remote host on Mac & Windows machines, as Desmond
requires Linux.
"""
model_class = traj_plot_models.TrajectoryEnergyJobTask.JobConfig
ui_module = job_settings_dialog_ui
[docs] def initSetOptions(self):
super().initSetOptions()
self.setWindowTitle("Energy Job Settings")
[docs] def defineMappings(self):
M = self.model_class
return [
(self.ui.host_selector, M.host_settings),
]
[docs] def isDoNotShowChecked(self):
return self.ui.do_not_show_again_cb.isChecked()
[docs] def showDialog(self, show_cb=False):
"""
Sets the visibility of the "do not show again" checkbox, then
execs the dialog
:param show_cb: Whether to show "do not show again" checkbox
:type show_cb: bool
"""
self.ui.do_not_show_again_cb.setVisible(show_cb)
self.exec()
[docs]class BrowseCfgFileDialog(QtWidgets.QDialog):
"""
Dialog box for asking the user to select a CFG file.
"""
[docs] def __init__(self, cfg_file):
super().__init__()
self.layout = QtWidgets.QVBoxLayout(self)
info_label = QtWidgets.QLabel(
f'The trajectory config file "{os.path.basename(cfg_file)}" could not be found.\nIt is required for the requested energy analysis.'
)
browse_layout = QtWidgets.QHBoxLayout()
self.layout.addWidget(info_label)
self.layout.addSpacing(10)
browse_label = QtWidgets.QLabel('Locate file:')
browse_btn = QtWidgets.QPushButton('Browse...')
browse_btn.clicked.connect(self.browseCfgFile)
self.initial_dir = os.path.dirname(cfg_file)
self.new_cfg_file = None
browse_layout.addWidget(browse_label)
browse_layout.addWidget(browse_btn)
browse_layout.addStretch()
self.layout.addLayout(browse_layout)
self.buttonBox = QtWidgets.QDialogButtonBox(self)
self.buttonBox.setStandardButtons(QtWidgets.QDialogButtonBox.Cancel |
QtWidgets.QDialogButtonBox.Ok)
self.buttonBox.accepted.connect(self.accept)
self.buttonBox.rejected.connect(self.reject)
self.okButton = self.buttonBox.button(QtWidgets.QDialogButtonBox.Ok)
self.okButton.setEnabled(False)
self.layout.addWidget(self.buttonBox)
[docs] def browseCfgFile(self):
filename = filedialog.get_open_file_name(
parent=self,
caption='Open File',
dir=self.initial_dir,
filter='Configuration Files (*.cfg)')
if filename:
self.okButton.setEnabled(True)
self.new_cfg_file = filename
[docs] def askForNewCfgFile(self):
"""
Show the dialog, and when it's dismissed, return the path to the new
CFG file, or None if dialog was cancelled.
"""
if self.exec():
return self.new_cfg_file
else:
return None
[docs]def parse_results_file(results_file):
"""
Parse the given results file (generated by analyze_simulation.py) and
and return data needed for generating a plot.
:param results_file: Result text file
:type results_file: str
:return: (Energy data, List of set ASLs, List of trajectory frame times in
picoseconds).
:rtype: (sea.Map, list, list)
"""
with open(results_file) as fh:
txt = fh.read()
cfg = sea.Map(txt)
keywords = cfg.Keywords
assert len(keywords) == 1
sem = keywords[0].SelectionEnergyMatrix
results = sem.Result
set_asls = sem.Asl_Selections.val
# Calculate the time of each frame (for x-axis of the plots):
frame0_time = cfg.TrajectoryFirstTime.val
interval = cfg.TrajectoryInterval_ps.val
num_frames = cfg.TrajectoryNumFrames.val
frame_times = list(np.arange(frame0_time, interval * num_frames, interval))
return results, set_asls, frame_times
[docs]def sum_results(results, use_sets, use_terms, include_self):
"""
Return a list of energies, one for each frame, by summing up the
terms specified by "use_terms" for interactions between sets
specified via "use_sets".
:param results: Energy data returned by parse_results_file().
:type results: sea.Map
:param use_sets: Which set names to sum energy for.
:type use_sets: list(str)
:param use_terms: Which terms to sum up. One or more of: elec, vdw,
stretch, angle, dihedral.
:type use_terms: list(str)
:param include_self: Whether to include terms within each set. If
False, include energies for interactions between terms are included.
:type include_self: bool
:return: Energy plot data, as a numpy array. Each item represents a frame,
and value is the total energy.
:rtype: numpy.array or None
"""
cross_ene = results.cross_ene
self_ene = results.self_ene
# Sum of these energies will produce the final values for the graph:
energy_arrays = []
for edict in (cross_ene, self_ene):
for term in use_terms:
val = edict[term]
for key in val.keys():
energies = val[key].val
if edict is cross_ene:
keys = key.split('-')
assert len(keys) == 2
if keys[0] not in use_sets or keys[1] not in use_sets:
continue
else:
# Self term
if not include_self or not key in use_sets:
continue
energy_arrays.append(np.asarray(energies))
if not energy_arrays:
return None
sum_energies = np.asarray(energy_arrays).sum(axis=0)
return sum_energies
if __name__ == "__main__":
# For testing the CustomSubstructureSetsDialog or CustomAslSetsDialog
# class outside of Maestro
import sys
from schrodinger import structure
app = QtWidgets.QApplication(sys.argv)
def sets_defined(sets):
print('selected sets:', sets)
st = structure.Structure.read(sys.argv[1])
dialog = CustomSubstructureSetsDialog(st, None)
# dialog = CustomAslSetsDialog(st, None)
dialog.setsDefined.connect(sets_defined)
sys.exit(dialog.exec())