"""
This module contains classes to enable use of the reorder module in a GUI.
Copyright Schrodinger, LLC. All rights reserved.
"""
import argparse
import sys
import schrodinger
from schrodinger import project
from schrodinger.application.matsci import msprops
from schrodinger.application.matsci import clusterstruct
from schrodinger.application.matsci import reorder
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.Qt.QtCore import Qt
from schrodinger.structutils import build
from schrodinger.ui import picking
from schrodinger.ui.qt import structure2d
from schrodinger.ui.qt import swidgets
from schrodinger.ui.qt.appframework2 import af2
from schrodinger.ui.qt.utils import wait_cursor
from schrodinger.utils import cmdline
maestro = schrodinger.get_maestro()
_version = "0.1"
ELEMENTS = 'Monatomic elements'
SMILES = 'Unique SMILES order - conformers only'
SMARTS = 'SMARTS patterns of whole molecule or unique atoms'
SUPERPOSITION = 'Based on superimposed 3D structures'
AUTO = 'Use all methods'
SELECTED_COLOR = QtGui.QColor('orange')
MAPPED_COLOR = QtGui.QColor(QtCore.Qt.green)
GUESS_COLOR = QtGui.QColor(QtCore.Qt.cyan)
GUESS_HIGHLIGHT_COLOR = QtGui.QColor(QtCore.Qt.blue)
SUGGEST_COLOR = QtGui.QColor(QtCore.Qt.gray)
WHITE = QtGui.QColor(QtCore.Qt.white)
BLACK = QtGui.QColor(QtCore.Qt.black)
REFERENCE = 'Reference'
COMPARISON = 'Comparison'
HELP_DITTY = {
REFERENCE: 'click to select atom',
COMPARISON: 'click to map to Reference atom'
}
# SHOW_H_PREF=1 means "show hydrogens that are present in source"
SHOW_H_PREF = 1
MARKER_NAME = 'MAPPED_ATOMS'
SUGGEST_MARKER = 'SUGGEST_MARKER'
MAX_SUGGESTED_DISTANCE = 5.0
[docs]class StructureView(structure2d.structure_view):
"""
View which holds a structure_item object
"""
# Signal to emit when mouse passes over an atom - passes in the index of the
# new atom (0 if there is no atom beneath the mouse)
atom_highlighted = QtCore.pyqtSignal(int)
[docs]class StructurePic(structure2d.structure_item):
"""
The QGraphicsItem that holds a 2D image of a structure
"""
[docs] def __init__(self, scene):
"""
Create a StructurePic instance
:type scene: QGraphicsScene
:param scene: The Scene this image is placed in
"""
origin_x = origin_y = 0
width = height = 300
rect = QtCore.QRectF(origin_x, origin_y, height, width)
structure2d.structure_item.__init__(self, rect=rect)
# Show all hydrogens
self.model2d.setShowHydrogenPreference(SHOW_H_PREF)
# The renderer currently has a copy of the model, so sync it
self.renderer.getModel().setShowHydrogenPreference(SHOW_H_PREF)
scene.addItem(self)
# Create all the annotators that we'll need - each annotator is created
# only once, and we modify the atoms/colors it annotates on the fly.
numerator = structure2d.AtomNumberAnnotator()
self.add_annotator(numerator)
# Annotates the selected atom in the Reference structure
self.selection_annotator = structure2d.RedSquareAnnotator(
size=75, color=SELECTED_COLOR)
self.add_annotator(self.selection_annotator)
# Annotates the atom that the mouse passes over, or the corresponding
# atom in the other structure
self.highlight_annotator = structure2d.RedSquareAnnotator(
size=65, color=MAPPED_COLOR)
self.add_annotator(self.highlight_annotator)
# Annotates a mapped atom with a background circle
self.mapped_annotator = structure2d.CircleAnnotator(gradient=True,
radius=50.,
color=MAPPED_COLOR)
self.add_annotator(self.mapped_annotator)
# Annotates a suggested atom with a grey square
self.suggested_annotator = structure2d.RedSquareAnnotator(
size=60., color=SUGGEST_COLOR)
self.add_annotator(self.suggested_annotator)
# Necessary to track the mouse over the image
self.setAcceptHoverEvents(True)
self.last_atom_highlighted = 0
[docs] def reset(self):
"""
Clear all the annotators and remove the picture
"""
self.selection_annotator.clearAtom()
self.highlight_annotator.clearAtom()
self.mapped_annotator.clearAtoms()
self.suggested_annotator.clearAtom()
self.clear()
[docs] def hoverMoveEvent(self, event):
"""
Track when the mouse is over an atom - emit a signal when it enters a
new atom or leaves an atom. An index of 0 is emitted when the mouse is
not over any atom.
"""
if not self.pic:
return
xval = event.pos().x() - self.pic.boundingRect().left()
yval = event.pos().y() - self.pic.boundingRect().top()
width = self.pic.boundingRect().width()
height = self.pic.boundingRect().height()
index = self.renderer.getAtomAtLocation(self.chmmol, width, height,
xval, yval) + 1
if index != self.last_atom_highlighted:
# Only emit the signal for a change in atom (either entered a new
# atom or left the current atom
view = self.scene().views()[0]
view.atom_highlighted.emit(index)
[docs] def setStructure(self, struct):
"""
Set the structure - create a new image of it
:type struct: `schrodinger.structure.Structure`
:param struct: The structure object for this picture
"""
self.set_structure(struct, allowRadicals=True)
self.generate_picture()
[docs] def selectAtom(self, index):
"""
Select an atom in the image and regenerate the picture
:type index: int
:param index: The atom index to select
"""
if self.selection_annotator.getAtom() != index:
self.selection_annotator.setAtom(index)
self.generate_picture()
[docs] def markAtomMapped(self, index, guess=False, generate=True):
"""
Mark an atom as mapped (either by guess or manually) and regenerate the
picture.
:type index: int
:param index: The atom index to mark
:type guess: bool
:param guess: True if this mark results from a guess, False if not. The
background color used depends on this parameter.
:param bool generate: Whether to regenerate the 2D picture
"""
if guess:
color = GUESS_COLOR
else:
color = MAPPED_COLOR
self.mapped_annotator.addAtom(index, color=color)
if generate:
self.generate_picture()
[docs] def markAtomUnmapped(self, index):
"""
Mark an atom as unmapped and regenerate the picture.
:type index: int
:param index: The atom index to mark
"""
self.mapped_annotator.removeAtom(index)
self.generate_picture()
[docs] def highlightAtom(self, index, guess=False):
"""
Highlight an atom
:type index: int
:param index: The atom index to highlight
:type guess: bool
:param guess: True if this atom is marked by a guess, False if it was
marked by the user. The background color used depends on this
parameter.
"""
if index == 0:
return
if guess:
self.highlight_annotator.setColor(GUESS_HIGHLIGHT_COLOR)
else:
self.highlight_annotator.setColor(MAPPED_COLOR)
self.highlight_annotator.setAtom(index)
self.generate_picture()
[docs] def highlightSuggestedAtom(self, index):
"""
Highlight suggested atom
:type index: int
:param index: The atom index to highlight
"""
self.suggested_annotator.clearAtom()
self.suggested_annotator.setAtom(index)
self.generate_picture()
[docs]class StructureFrame(QtWidgets.QFrame):
"""
A QFrame that contains a 2D image of a structure. This also creates a
QListWidget that is coordinated with the 2D image. This is the base class
for both the Reference and Comparison structures.
"""
[docs] def __init__(self,
master,
label_layout,
structure_layout,
list_layout,
struct=None,
mytype=COMPARISON,
atomic_constraints=None,
enable_2d_structure=True):
"""
Create a StructureFrame instance
:type master: ReorderAtomFrame
:param master: The master Frame for this widget
:type label_layout: QBoxLayout
:param label_layout: The layout to add the label.
:type structure_layout: QBoxLayout
:param structure_layout: The layout to add this widget to.
:type list_layout: QBoxLayout
:param list_layout: The layout to add the associated ListWidget to.
:type struct: `schrodinger.structure.Structure`
:param struct: The structure for this frame
:type mytype: str
:param mytype: Either REFERENCE or COMPARISON, the type of structure
this Frame applies to.
:type atomic_constraints: list or None
:param atomic_constraints: list of atomic constraints for each atom or
None
:param bool enable_2d_structure: If True, enable/show 2D structure. Can
be slow for inorganic materials
"""
QtWidgets.QFrame.__init__(self)
self.master = master
self.enable_2d_structure = enable_2d_structure
mylayout = swidgets.SVBoxLayout(self)
structure_layout.addWidget(self)
self.mytype = mytype
self.label = swidgets.SLabel(mytype, layout=label_layout)
self.setLabel()
if self.enable_2d_structure:
# The structure image
self.scene = structure2d.structure_scene()
self.view = StructureView(self.scene)
self.pic = StructurePic(self.scene)
mylayout.addWidget(self.view)
# The atom list widget
if mytype == REFERENCE:
self.list_widget = ReferenceListWidget(
self.master, atomic_constraints=atomic_constraints)
tip = ('Click on a row to select the atom for that row in the '
'reference structure.\n\nRows for unmapped atoms show the '
'atomic symbol and atom index\nof the atom that row will '
'select.\n\nRows for mapped atoms additionally show the '
'atom index\nof the atom in the comparison structure that '
'is mapped to the atom for that row.\n\nGreen rows indicate '
'atoms that are already mapped,\nblue rows indicate atoms '
'that will be mapped if the current guess is accepted,\n'
'white rows are unmapped atoms.')
ltip = ('Green atoms are mapped, blue atoms will be mapped if the '
'guess is accepted.\nA yellow box indicates selection.\n'
'Hovering over '
'a mapped atom draws a box around that atom and\n '
'the corresponding atom in the comparison structure.')
else:
self.list_widget = ComparisonListWidget(
self.master, atomic_constraints=atomic_constraints)
tip = ('Click on a row to map the atom for that row in the '
'comparison structure\nto the'
'currently-selected reference structure atom.\n\nRows show '
'the atomic symbol and atom index\nof the comparison '
'structure atom that row will map.\n\nGreen rows indicate '
'atoms that are already mapped,\nblue rows indicate atoms '
'that will be mapped if the current guess is accepted,\n'
'white rows are unmapped atoms.')
ltip = ('Green atoms are mapped, blue atoms will be mapped if the '
'guess is accepted.\nHovering over '
'a mapped atom draws a box around that atom and\n '
'the corresponding atom in the reference structure.')
self.list_widget.setToolTip(tip)
self.label.setToolTip(ltip)
list_layout.addWidget(self.list_widget)
self.struct = struct
if self.struct:
self.setStructure(struct)
[docs] def generatePicture(self):
""" Regenerate the 2D picture to show any changes """
if self.enable_2d_structure:
self.pic.generate_picture()
[docs] def setLabel(self, guessing=False):
"""
Set the label that describes what the user should do
:type guessing: bool
:param guessing: True if there is currently a guess in place, False if
not.
"""
if guessing:
text = self.mytype + ' - accept or reject guess'
else:
text = self.mytype + ' - ' + HELP_DITTY[self.mytype]
self.label.setText(text)
[docs] def setStructure(self, struct):
"""
Set the structure for this frame. Updates the image and the list widget
:type struct: `schrodinger.structure.Structure`
:param struct: The structure for this frame
"""
self.reset()
self.struct = struct
self.dcell, tmp = clusterstruct.create_distance_cell(
struct, MAX_SUGGESTED_DISTANCE)
self.list_widget.setStructure(self.struct)
if self.enable_2d_structure:
self.pic.setStructure(self.struct)
[docs] def highlightAtom(self, index, guess=False):
"""
Highlight the given atom
:type index: int
:param index: The atom index to highlight
:type guess: bool
:param guess: True if this atom is marked by a guess, False if it was
marked by the user. The background color used depends on this
parameter.
"""
if self.enable_2d_structure:
self.pic.highlightAtom(index, guess=guess)
[docs] def unmapAtom(self, index):
"""
Mark an atom as unmapped
:type index: int
:param index: The atom index to mark
"""
self.list_widget.markAtomUnmapped(index)
if self.enable_2d_structure:
self.pic.markAtomUnmapped(index)
[docs] def reset(self, reset_structure=True):
"""
Reset to a blank frame, or remove any mapping but keep the structure
:type reset_structures: bool
:param reset_structure: True if the structure should be reset, False if
it should be kept
"""
if self.enable_2d_structure:
self.pic.reset()
self.list_widget.clear()
if reset_structure:
self.struct = None
else:
# We set the structure again so that the Picture and ListWidget are
# regenerated anew with the raw information for this structure.
self.setStructure(self.struct)
[docs] def isValidAtomIndex(self, index):
"""
Check if this index is a valid atom index for the loaded structure
:type index: int
:param index: The atom index to check
:rtype: bool
:return: True 0 < index <= atom_total, False otherwise
"""
if not self.struct:
return False
else:
return 0 < index <= self.struct.atom_total
[docs] def selectAtom(self, index):
"""
Select an atom in both the image and listwidget
:type index: int
:param index: The atom index to select
:rtype: `structure._StructureAtom`
:return: selected atom
"""
self.list_widget.selectAtom(index)
if self.enable_2d_structure:
self.pic.selectAtom(index)
return self.struct.atom[index]
[docs] def mapAtom(self, index, guess=False, generate=True):
"""
Mark an atom as mapped
:type index: int
:param index: The atom index to mark
:type guess: bool
:param guess: True if this atom is marked by a guess, False if it was
marked by the user. The background color used depends on this
parameter.
:param bool generate: Whether to regenerate the 2D picture
"""
self.list_widget.markAtomMapped(index, guess=guess)
if self.enable_2d_structure:
self.pic.markAtomMapped(index, guess=guess, generate=generate)
[docs]class ComparisonFrame(StructureFrame):
"""
The StructureFrame for the Comparison structure
"""
[docs] def suggestCloseAtom(self, xyz):
"""
Highlight and return index of the closest atom from the XYZ location
taking into account PBCs.
:type xyz: list(float)
:param: Coordinates
:rtype: int
:return: Closest atom index
"""
if not self.dcell:
return
min_distsq = float('Inf')
atom_idx = None
for atom in self.dcell.query_atoms(*xyz):
distsq = atom.getDistanceSquared()
if distsq < min_distsq:
min_distsq = distsq
atom_idx = atom.getIndex()
if atom_idx is None:
return
self.list_widget.markSuggestedAtom(atom_idx)
if self.enable_2d_structure:
self.pic.highlightSuggestedAtom(atom_idx)
return atom_idx
[docs] def suggestedIndex(self):
"""
Get last suggested index.
:rtype: int or None
:return: last suggested index
"""
return self.list_widget.suggested_atom
[docs]class ReferenceFrame(StructureFrame):
"""
The StructureFrame for the Reference structure
"""
[docs] def __init__(self, *args, **kwargs):
"""
Create a ReferenceFrame object. Parameters are described in the
StructureFrame class
"""
kwargs['mytype'] = REFERENCE
StructureFrame.__init__(self, *args, **kwargs)
[docs] def mapAtom(self, ref_index, advance=True, **kwargs):
"""
Map the given atom index in the comparison structure to the given atom
index in the reference structure.
:type ref_index: int
:param ref_index: The atom index of the reference molecule
:type advance: bool
:param advance: Whether to advance the selected atom or not. The
selection will automatically be advanced if ref_index is the currently
selected atom regardless of this setting.
"""
super().mapAtom(ref_index, **kwargs)
if advance or self.selectedIndex() == ref_index:
self.advanceSelection(ref_index)
[docs] def advanceSelection(self, index):
"""
Move the selection to the next unmapped atom with a higher index. If
none exist, start checking again at atom 1.
:type index: int
:param index: Start search for unmapped atoms with indexes higher than
this number
"""
indexes = list(range(index + 1, self.struct.atom_total + 1))
indexes += list(range(1, index))
for ind in indexes:
if not self.master.compAtomFromRefAtom(ind):
self.master.selectRefAtom(ind)
break
[docs] def selectedIndex(self):
"""
Get the index of the currently selected atom
:rtype: int
:return: The index of the selected Reference atom
"""
return self.list_widget.currentRow() + 1
[docs]class ReorderAtomFrame(swidgets.SFrame):
"""
A frame which allows the user to reorder the atoms of a comparison structure
based on the atom order of a reference structure. The widgets provide the
ability to guess at atom order based on multiple algorithms and also
manually specify order using 2D images of the structures.
"""
MAP_COUNTER_LABEL = '{num_mapped_atoms} of {atom_total}\natoms mapped'
[docs] def __init__(self,
parent=None,
by_element=True,
by_smiles=True,
by_smarts=True,
by_superposition=True,
default_method=SMARTS,
layout=None,
reference_structure=None,
comparison_structure=None,
reference_structure_constraints=None,
enable_2d_structure=True):
"""
Create a ReorderAtomFrame widget
:type parent: QWidget
:param parent: The parent widget with a .warning(str) method
:type by_element: bool
:param by_element: Provide widgets to allow a guess based on lone
elements
:type by_smiles: bool
:param by_smiles: Provide widgets to allow a guess based on unique
SMILES order
:type by_smarts: bool
:param by_smarts: Provide widgets to allow a guess based on SMARTS
patterns
:type by_superposition: bool
:param by_superposition: Provide widgets to allow a guess based on
3D superposition of the structures
:type default_method: str
:param default_method: Default guess method
:type layout: QBoxLayout
:param layout: The layout to place this frame into
:type reference_structure: `schrodinger.structure.Structure`
:param reference_structure: The structure to use as the reference
:type comparison_structure: `schrodinger.structure.Structure`
:param comparison_structure: The structure to use as the comparison
:param bool enable_2d_structure: If True, enable/show 2D structure. Can
be slow for inorganic materials
"""
swidgets.SFrame.__init__(self, layout=layout)
self.guess = {}
self.mapping = {}
self.reverse_mapping = {}
self.master = parent
self.enable_2d_structure = enable_2d_structure
if maestro:
# Initialize markers
for marker in (MARKER_NAME, SUGGEST_MARKER):
maestro.command('markers %s not all' % marker)
# This layout is accessible externally via the .layout() method
mylayout = self.mylayout
# Method to guess at atom ordering
self.guess_gb = swidgets.SGroupBox('Guess method',
parent_layout=mylayout)
glayout = self.guess_gb.layout
items = [ELEMENTS, SMILES, SMARTS, SUPERPOSITION, AUTO]
defindex = items.index(default_method)
self.guess_rbg = swidgets.SRadioButtonGroup(labels=items,
layout=glayout,
default_index=defindex)
rbtips = {
ELEMENTS: 'Maps atoms that are the only atom of an element'
' in each structure.',
SMILES: 'Creates a unique SMILES string for each structure.\n'
'If the SMILES strings match, the product atoms then\n'
'by mapped to the reactant atoms.\n\nThis method only '
'works for conformers and may not be able to \n'
'determine the mapping for hydrogen atoms.',
SMARTS: 'Maps those atoms with unique SMARTS patterns',
SUPERPOSITION:
'The currently mapped atoms are superimposed\n'
'on each other to bring the two structures into\n'
'alignment. Unmapped reference atoms are then mapped\n'
'to the closest comparison atom of the same element.\n'
'Atoms may not be mapped if no atom of the same\n'
'element is nearby.',
AUTO: 'Map all atoms that can be mapped by any of the above '
'methods.'
}
for button in self.guess_rbg.buttons():
button.setToolTip(rbtips.get(str(button.text()), ""))
blayout = swidgets.SHBoxLayout(layout=glayout)
blayout.addStretch()
self.guess_btn = swidgets.SPushButton('Guess',
command=self.guessOrder,
layout=blayout)
tip = ('Guess at a mapping using the method selected above.\nGuessed '
'atom mappings will be shown in blue and must be\n'
'accepted/rejected before manual mapping can continue')
self.guess_btn.setToolTip(tip)
self.guess_accept_btn = swidgets.SPushButton('Accept Guess',
command=self.guessAccepted,
layout=blayout)
tip = ('Accept all the guessed atom mappings')
self.guess_accept_btn.setToolTip(tip)
self.guess_reject_btn = swidgets.SPushButton('Reject Guess',
command=self.guessRejected,
layout=blayout)
tip = ('Reject all the guessed atom mappings')
self.guess_reject_btn.setToolTip(tip)
self.guess_accept_btn.setEnabled(False)
self.guess_reject_btn.setEnabled(False)
self.automap_cb = swidgets.SCheckBox(
'Automatically map hydrogens and '
'monatomic elements',
layout=mylayout,
checked=True)
tip = ('After each time a comparison atom is mapped to a reference atom'
',\nan attempt will be made to automatically map any unmapped\n'
'protons bound to mapped atoms or unmapped atoms that are the\n'
'last remaining atom of that element.')
self.automap_cb.setToolTip(tip)
# pick
text = ('Pick reference and comparison atom pairs in Workspace')
self.pick_cb = swidgets.SCheckBoxToggle(text,
checked=False,
layout=mylayout,
command=self._pickStateChanged)
self.pick_cb.setToolTip('Pick atom pairs for superposition')
self.pick_toggle = PickPairToggleIntermediate(
self.pick_cb,
self._processPickedPair,
pick_text='Pick reference and comparison atom pairs')
# Interactive atom re-numbering
label_layout = swidgets.SHBoxLayout(layout=mylayout)
self.structure_frame = QtWidgets.QFrame()
slayout = swidgets.SHBoxLayout(self.structure_frame)
self.list_frame = QtWidgets.QFrame()
list_layout = swidgets.SHBoxLayout(self.list_frame)
list_layout.addStretch()
if self.enable_2d_structure:
self.splitter = QtWidgets.QSplitter(QtCore.Qt.Vertical)
mylayout.addWidget(self.splitter)
self.splitter.addWidget(self.structure_frame)
self.splitter.addWidget(self.list_frame)
else:
mylayout.addWidget(self.list_frame)
# Left/reactant/reference side list of atoms
self.reference = ReferenceFrame(
self,
label_layout,
slayout,
list_layout,
struct=reference_structure,
atomic_constraints=reference_structure_constraints,
enable_2d_structure=enable_2d_structure)
list_layout.addStretch()
# Add an indicator for number of atoms mapped
rmap_layout = swidgets.SVBoxLayout(layout=list_layout)
rmap_layout.addStretch()
self.rmap_count_label = swidgets.SLabel(layout=rmap_layout)
self.rmap_count_label.setAlignment(Qt.AlignHCenter)
# Create a button centered in this layout but at the bottom of it
self.reset_map_btn = swidgets.SPushButton('Clear Mapping',
layout=rmap_layout,
command=self.resetMap)
tip = 'Reset all mapping information for these two structures'
self.reset_map_btn.setToolTip(tip)
rmap_layout.setAlignment(self.reset_map_btn, Qt.AlignHCenter)
# Left/reactant/comparison side list of atoms
list_layout.addStretch()
self.comparison = ComparisonFrame(
self,
label_layout,
slayout,
list_layout,
struct=comparison_structure,
atomic_constraints=reference_structure_constraints,
enable_2d_structure=enable_2d_structure)
list_layout.addStretch()
# Connect signals to react to user interaction
if self.enable_2d_structure:
self.reference.view.atom_clicked.connect(self.selectRefAtom)
self.reference.view.atom_highlighted.connect(self.highlightRefAtom)
self.comparison.view.atom_clicked.connect(self.mapCompAtom)
self.comparison.view.atom_highlighted.connect(
self.highlightCompAtom)
self.reference.list_widget.atom_clicked.connect(self.selectRefAtom)
self.comparison.list_widget.atom_clicked.connect(self.mapCompAtom)
# Call reset to select default reference list item
if reference_structure and comparison_structure:
self.resetMap()
self.pick_toggle.setStructures(reference_structure,
comparison_structure)
def _pickStateChanged(self, state):
"""
React to a change in self.pick_cb state.
:type state: bool
:param state: True if self.pick_cb is checked
"""
pt = maestro.project_table_get()
ref_id = self.reference.struct.property[msprops.ENTRY_ID_PROP]
comp_id = self.comparison.struct.property[msprops.ENTRY_ID_PROP]
maestro.command('beginundoblock')
maestro.command('tilemode tile=%s' % str(bool(state)).lower())
# Include ref entry first so that it's on the left side in the tile mode
# Cannot include both entries at once
maestro.command('entrywsincludeonly entry "%s"' % ref_id)
if state:
maestro.command('entrywsinclude entry "%s"' % comp_id)
self.markAtomsInWS()
maestro.command('endundoblock')
[docs] def updateMapCounter(self):
"""
Updates the label indicating the number of atoms mapped
:rtype: None
"""
num_mapped_atoms = len(self.mapping)
# If the reference structure does not exist, report 0 total atoms.
try:
atom_total = self.reference.struct.atom_total
except AttributeError:
atom_total = 0
counter = self.MAP_COUNTER_LABEL.format(
num_mapped_atoms=num_mapped_atoms, atom_total=atom_total)
self.rmap_count_label.setText(counter)
[docs] def show2DStructures(self, state):
"""
Show or hide 2D structure pictures
:type state: bool
:param state: If True, show 2D structure picture, otherwise hide
"""
if state:
# To show both size, just set values different than 0
self.splitter.setSizes([100, 100])
else:
# To hide one side, set 0 value for it
self.splitter.setSizes([0, 1])
[docs] def setStructures(self, reference, comparison):
"""
Set the reference and comparison structures
:type reference: `schrodinger.structure.Structure`
:param reference: The structure to use as the reference
:type comparison: `schrodinger.structure.Structure`
:param comparison: The structure to use as the comparison
"""
self.reset()
self.reference.setStructure(reference)
self.comparison.setStructure(comparison)
self.pick_toggle.setStructures(reference, comparison)
[docs] def guessMode(self, guessmode):
"""
Set the current guess mode - affects label texts and enabling/disabling
widgets and user interaction signals.
:type guessmode: bool
:param guessmode: True if a guess is currently taking place and has not
yet been Accepted/Rejected, False if any other situation
"""
if guessmode:
if self.enable_2d_structure:
self.reference.view.atom_clicked.disconnect(self.selectRefAtom)
self.comparison.view.atom_clicked.disconnect(self.mapCompAtom)
self.reference.list_widget.atom_clicked.disconnect(
self.selectRefAtom)
self.comparison.list_widget.atom_clicked.disconnect(
self.mapCompAtom)
else:
if self.enable_2d_structure:
self.reference.view.atom_clicked.connect(self.selectRefAtom)
self.comparison.view.atom_clicked.connect(self.mapCompAtom)
self.reference.list_widget.atom_clicked.connect(self.selectRefAtom)
self.comparison.list_widget.atom_clicked.connect(self.mapCompAtom)
self.guess_btn.setEnabled(not guessmode)
self.guess_accept_btn.setEnabled(guessmode)
self.guess_reject_btn.setEnabled(guessmode)
self.reference.setLabel(guessing=guessmode)
self.comparison.setLabel(guessing=guessmode)
[docs] def selectRefAtom(self, index):
"""
Select the current atom in the reference structure
:type index: int
:param index: The index of the atom to select
"""
if not self.reference.isValidAtomIndex(index):
# Workaround for MATSCI-991
return
ref_atom = self.reference.selectAtom(index)
comp_idx = self.comparison.suggestCloseAtom(ref_atom.xyz)
self.markSuggestedAtomsInWS(index, comp_idx)
[docs] def mapPair(self,
ref_index,
comp_index,
advance=True,
guess=False,
generate=True):
"""
Create a mapping between an atom in the Comparison structure to an atom
in the Reference structure.
:type ref_index: int
:param ref_index: The atom index of the reference atom to map
:type comp_index: int
:param comp_index: The atom index of the comparison atom to map
:type advance: bool
:param advance: True if the Reference atom selection should be advanced,
False if not
:type guess: bool
:param guess: True if this mark results from a guess, False if not. The
background color used depends on this parameter.
:param bool generate: Whether to regenerate the 2D picture
"""
self.mapping[ref_index] = comp_index
self.reverse_mapping[comp_index] = ref_index
self.comparison.mapAtom(comp_index, guess=guess, generate=generate)
self.reference.mapAtom(ref_index,
guess=guess,
advance=advance,
generate=generate)
self.updateMapCounter()
[docs] def unmapPair(self, ref_index, comp_index):
"""
Create a mapping between an atom in the Comparison structure to an atom
in the Reference structure.
:type ref_index: int
:param ref_index: The atom index of the reference atom to map
:type comp_index: int
:param comp_index: The atom index of the comparison atom to map
"""
del self.mapping[ref_index]
del self.reverse_mapping[comp_index]
self.reference.unmapAtom(ref_index)
self.comparison.unmapAtom(comp_index)
self.updateMapCounter()
[docs] def mapCompAtom(self, comp_index, ref_index=None, auto=True, advance=True):
"""
Map an atom in the Comparison structure to an atom in the Reference
structure
:type comp_index: int
:param comp_index: The atom index of the comparison atom to map
:type ref_index: int or None
:param ref_index: If an int, the index of the reference atom to map to.
If not given, the selected reference atom will be used
:type auto: bool
:param auto: Whether to Auto guess additional atoms based on this
mapping. Default is True - but autoguessing will only occur if the user
has checked the Auto guess box.
:type advance: bool
:param advance: True if the Reference atom selection should be advanced,
False if not
"""
if not self.comparison.isValidAtomIndex(comp_index):
# Workaround for MATSCI-991
return
if ref_index is None:
ref_index = self.reference.selectedIndex()
# Remove any previous mapping for the reference atom
previous_comp_index = self.compAtomFromRefAtom(ref_index)
if previous_comp_index:
self.unmapPair(ref_index, previous_comp_index)
# Remove any previous mapping for the comparison atom
previous_ref_index = self.refAtomFromCompAtom(comp_index)
if previous_ref_index:
self.unmapPair(previous_ref_index, comp_index)
self.mapPair(ref_index, comp_index, advance=advance)
if auto and self.automap_cb.isChecked():
self.autoMap()
self.markAtomsInWS()
[docs] def markAtomsInWS(self):
"""
Mark selected atoms for guess in the WS.
"""
if not self.pick_cb.isChecked():
return
maestro.command('hidemarkers ' + MARKER_NAME)
ref_atoms = self.reference.list_widget.getMarkedAtoms()
comp_atoms = self.comparison.list_widget.getMarkedAtoms()
comp_atoms = [x + self.reference.struct.atom_total for x in comp_atoms]
indexes = ref_atoms + comp_atoms
if indexes:
# Maestro expects RGB normalized to 1
rgb = [x / 255 for x in MAPPED_COLOR.getRgb()[:3]]
asl = 'atom.n ' + ','.join([str(x) for x in indexes])
color = 'r=%f g=%f b=%f' % tuple(rgb)
maestro.command('markers %s %s %s' % (color, MARKER_NAME, asl))
maestro.command('showmarkers %s' % MARKER_NAME)
self.markSuggestedAtomsInWS(self.reference.selectedIndex(),
self.comparison.suggestedIndex())
[docs] def markSuggestedAtomsInWS(self, ref_idx, comp_idx):
"""
Pick reference atom in the WS and mark suggested comparison atom.
:type ref_idx: int
:param ref_idx: Index of the atom to pick in the reference WS
:type comp_idx: int or None
:param comp_idx: If int, mark atom in the comparison WS structure
"""
if not maestro:
return
maestro.command('hidemarkers ' + SUGGEST_MARKER)
if not self.pick_cb.isChecked():
return
# Prevent from calling the pick function, not to enter in the
# infinite recursion
self.pick_toggle._atomPicked(ref_idx, call_pick_function=False)
if comp_idx is None:
return
comp_idx += self.reference.struct.atom_total
asl = 'atom.n %d' % comp_idx
# Maestro expects RGB normalized to 1
rgb = [x / 255 for x in GUESS_COLOR.getRgb()[:3]]
color = 'r=%f g=%f b=%f' % tuple(rgb)
maestro.command('markers %s %s %s' % (color, SUGGEST_MARKER, asl))
maestro.command('showmarkers %s' % SUGGEST_MARKER)
[docs] def autoMap(self):
"""
Automatically map no-brainer atoms based on the current mapping.
No-brainer atoms are atoms that are the only atoms remaining of a given
element, or hydrogens bound to mapped atoms that can be uniquely
identified.
"""
element_map = reorder.map_by_lone_element(self.reference.struct,
self.comparison.struct,
atom_map=self.mapping)
for ref, comp in element_map.items():
self.mapCompAtom(comp, ref_index=ref, auto=False, advance=False)
proton_map = reorder.map_hydrogens(self.reference.struct,
self.comparison.struct, self.mapping)
for ref, comp in proton_map.items():
self.mapCompAtom(comp, ref_index=ref, auto=False, advance=False)
[docs] def compAtomFromRefAtom(self, index):
"""
Get the atom in the Comparison structure that is mapped to the given
index in the Reference structure
:type index: int
:param index: The index of the Reference atom to check
:rtype: int
:return: The index of the Comparison atom mapped to the given Reference
atom
"""
return self.mapping.get(index, 0)
[docs] def refAtomFromCompAtom(self, index):
"""
Get the atom in the Reference structure that is mapped to the given
index in the Comparison structure
:type index: int
:param index: The index of the Comparison atom to check
:rtype: int
:return: The index of the Reference atom mapped to the given Comparison
atom
"""
return self.reverse_mapping.get(index, 0)
[docs] def highlightCompAtom(self, comp_index):
"""
Highlight an atom in the Comparision structure
:type comp_index: int
:param comp_index: The index of the comparison atom to highlight
"""
ref_index = self.refAtomFromCompAtom(comp_index)
if not ref_index:
# We only highlight mapped atoms
comp_index = 0
self.highlightMapPair(ref_index, comp_index)
[docs] def highlightRefAtom(self, ref_index):
"""
Highlight an atom in the Reference structure
:type ref_index: int
:param ref_index: The index of the reference atom to highlight
"""
comp_index = self.compAtomFromRefAtom(ref_index)
if not comp_index:
# We only highlight mapped atoms
ref_index = 0
self.highlightMapPair(ref_index, comp_index)
[docs] def highlightMapPair(self, ref_index, comp_index):
"""
Highlight the mapped pair of atoms
:type ref_index: int
:param ref_index: The index of the reference atom to highlight
:type comp_index: int
:param comp_index: The index of the comparison atom to highlight
"""
guess = ref_index in self.guess
self.reference.highlightAtom(ref_index, guess=guess)
self.comparison.highlightAtom(comp_index, guess=guess)
[docs] def guessOrder(self):
"""
Make a guess at the new order of the Comparison atoms
"""
if not self.reference.struct:
self.master.warning('Structures must be loaded first')
return
self.guessMode(True)
self.guess = {}
do_elements = do_smiles = do_smarts = do_superposition = False
choice = self.guess_rbg.checkedText()
if choice == AUTO:
do_elements = do_smiles = do_smarts = do_superposition = True
elif choice == ELEMENTS:
do_elements = True
elif choice == SMILES:
do_smiles = True
elif choice == SMARTS:
do_smarts = True
elif choice == SUPERPOSITION:
do_superposition = True
all_guesses = {}
amap = self.mapping.copy()
rstruct = self.reference.struct
cstruct = self.comparison.struct
# reorder.map_by functions can be slow, wrap them with wait cursor
# (MATSCI-6722)
if do_elements:
with wait_cursor:
aguess = reorder.map_by_lone_element(rstruct,
cstruct,
atom_map=amap)
all_guesses.update(aguess)
if do_smiles:
supermap = amap.copy()
supermap.update(all_guesses)
with wait_cursor:
aguess = reorder.map_by_smiles(rstruct,
cstruct,
atom_map=supermap)
all_guesses.update(aguess)
if do_smarts:
supermap = amap.copy()
supermap.update(all_guesses)
with wait_cursor:
aguess = reorder.map_by_smarts(rstruct,
cstruct,
atom_map=supermap)
all_guesses.update(aguess)
superposed_warn = False
if do_superposition:
supermap = amap.copy()
supermap.update(all_guesses)
if len(supermap) < 3:
msg = ('At least 3 atoms must be mapped to use a superimposed '
'guess. These atoms are used to superimpose the '
'structures. Currently only %d atoms are mapped so '
'skipping superimposed guess.' % len(supermap))
self.master.warning(msg)
superposed_warn = True
else:
with wait_cursor:
aguess = reorder.map_by_superposition(rstruct,
cstruct,
atom_map=supermap)
all_guesses.update(aguess)
if not all_guesses:
if not superposed_warn or any([do_elements, do_smiles, do_smarts]):
# Don't give this warning if superposition was the only guess
# tried and we already said we couldn't do it
self.master.warning('No atoms could be mapped')
self.finishGuess()
else:
with wait_cursor:
for ref, comp in all_guesses.items():
if ref not in self.mapping:
self.guess[ref] = comp
# Do not regenerate the picture after marking each
# pair, wait until the very end - saves much time
self.mapPair(ref,
comp,
guess=True,
advance=False,
generate=False)
# Now regenerate the final picture
if all_guesses:
self.reference.generatePicture()
self.comparison.generatePicture()
[docs] def guessAccepted(self):
"""
User has accepted the guess, make the atoms truly mapped
"""
# Mapping from the guess can be slow (MATSCI-6722)
with wait_cursor:
for ref, comp in self.guess.items():
# Do not regenerate the picture after each atom - a final
# picture will be generated by highlightMapPair below - this
# saves a lot of time for larger structures
self.mapPair(ref,
comp,
guess=False,
advance=False,
generate=False)
self.highlightMapPair(0, 0)
self.markAtomsInWS()
self.finishGuess()
[docs] def guessRejected(self):
"""
User has rejected the guess remove it
"""
# Unmapping from the guess can be slow (MATSCI-6722)
with wait_cursor:
for ref, comp in self.guess.items():
self.unmapPair(ref, comp)
self.finishGuess()
[docs] def isInGuessMode(self):
"""
Is the panel currently in guess mode?
:rtype: bool
:return: True if the panel is in guess mode, False if not
"""
return not self.guess_btn.isEnabled()
[docs] def finishGuess(self):
"""
Finish guess mode after the user has accepted or rejected the guess.
Note - this method should be callable even if not currently in a guess.
"""
if self.isInGuessMode():
# We are currently in the middle of a guess, finish it
self.guess = {}
self.guessMode(False)
self.reference.advanceSelection(0)
[docs] def doneReordering(self):
"""
Check if the reordering process is done
:rtype: bool
:return: True if all the atoms in one of the structures are mapped
"""
num_mapped = len(self.mapping)
return (num_mapped == self.reference.struct.atom_total or
num_mapped == self.comparison.struct.atom_total)
[docs] def getReorderedStructure(self):
"""
Get a reordered version of the comparison structure
:rtype: `schrodinger.structure.Structure`, list or (None, None)
:return: The comparison structure with the atoms reordered as specified
in the mapping, and the list used to reorder the structure. The first
element of the list is the original index of the atom that is first in
the reordered structure, the second is the original index of the atom
that is second in the reordered structure, etc. (None, None) is
returned if the structure is not ready for reordering.
"""
if not self.guess_btn.isEnabled():
self.master.warning('The guess must be accepted or rejected first')
return None, None
if not self.doneReordering():
return None, None
order = []
for index in range(1, self.reference.struct.atom_total + 1):
order.append(self.mapping[index])
if self.comparison.struct.atom_total > len(order):
for index in range(1, self.comparison.struct.atom_total + 1):
if index not in self.reverse_mapping:
order.append(index)
return build.reorder_atoms(self.comparison.struct, order), order
def _processPickedPair(self, indices):
"""
Process a picked atoms pair.
:param indices: (ref_idx, comp_idx) pair
:type indices: list(int)
"""
# MATSCI-6560
if not indices:
return
# Intermediate picking state, only one atom is picked
if len(indices) == 1:
if self.reference.isValidAtomIndex(indices[0]):
self.selectRefAtom(indices[0])
else:
self.master.error('Select reference atom first')
self.markSuggestedAtomsInWS(self.reference.selectedIndex(),
self.comparison.suggestedIndex())
return
# since picking from two separate entries included in
# the Maestro Workspace the given atom indices are
# aggregated with reference first, comparison second, therefore
# offset the comparison index by the number of reference atoms
ref_idx, tmp_idx = indices
comp_idx = tmp_idx - self.reference.struct.atom_total
if not all([
self.reference.isValidAtomIndex(ref_idx),
self.comparison.isValidAtomIndex(comp_idx)
]):
self.master.error(
'Atom pairs must be picked in (reference, '
'comparison) order where reference is on the left of the '
'Workspace and comparison is on the right.')
return
self.selectRefAtom(ref_idx)
self.mapCompAtom(comp_idx)
[docs] def resetMap(self):
"""
Reset only the map, keeping the structures
"""
self.reset(reset_structures=False)
if self.reference.struct:
self.reference.selectAtom(1)
self.markAtomsInWS()
self.updateMapCounter()
[docs] def reset(self, reset_structures=True):
"""
Reset the frame
:type reset_structures: bool
:param reset_structures: True if structures should be reset, False if
they should be kept
"""
self.finishGuess()
self.mapping = {}
self.reverse_mapping = {}
self.reference.reset(reset_structure=reset_structures)
self.comparison.reset(reset_structure=reset_structures)
Super = af2.JobApp
[docs]class ReorderAtomPanel(Super):
"""
A simple panel that takes two selected entries and allows one of them to be
reordered - this is mainly for testing and example purposes.
"""
[docs] def setPanelOptions(self):
"""
Override the generic parent class to set panel options
"""
Super.setPanelOptions(self)
self.title = 'Reorder Atoms'
self.program_name = 'Reorderatoms'
self.input_selector_options = {
'file': True,
'selected_entries': False,
'included_entries': True,
'included_entry': False,
'workspace': False
}
[docs] def layOut(self):
"""
Lay out the widgets for this panel
"""
Super.layOut(self)
layout = self.main_layout
swidgets.SPushButton('Load', command=self.loadStructures, layout=layout)
self.reorder_frame = ReorderAtomFrame(parent=self, layout=layout)
size = self.size()
size.setHeight(800)
size.setWidth(800)
self.resize(size)
[docs] def loadStructures(self):
"""
Load in the two structures. The first selected entry will be the
reference.
"""
structs = list(self.input_selector.structures())
if len(structs) == 2:
self.reorder_frame.setStructures(*structs)
else:
self.warning('Exactly 2 structures must be included in the '
'workspace')
[docs] @af2.appmethods.custom('Create Reordered Entry', 1,
'Create a new, reordered entry')
def createReorderedEntry(self):
"""
Create a new, reordered version of the second selected entry. If run
from maestro, this will be created as a new project entry and included
in the workspace.
"""
done = self.reorder_frame.doneReordering()
if not done:
self.warning('Not all reference structure atoms are mapped to the '
'comparison structure')
else:
struct, order = self.reorder_frame.getReorderedStructure()
if not struct:
return
if not maestro:
for ind, atom in enumerate(struct.atom):
print(atom.element, atom.index, order[ind])
else:
struct.title = struct.title + '-reordered'
ptable = maestro.project_table_get()
row = ptable.importStructure(struct)
row.in_workspace = project.IN_WORKSPACE
[docs]class ReorderAtomsDialog(swidgets.SDialog):
"""
A Window-Modal dialog that allows the user to reorder atoms in a comparison
structure based on a reference structure.
"""
[docs] def __init__(self,
master,
struct1,
struct2,
modality=QtCore.Qt.NonModal,
struct1_atomic_constraints=None,
default_guess_method=SMARTS,
enable_2d_structure=True,
**kwargs):
"""
Create a ReorderAtomDialog instance. See parent for `**kwargs` docs.
:param QWidget master: The parent widget
:type struct1: `schrodinger.structure.Structure`
:param struct1: The reference structure. This structure will not be
reordered.
:type struct2: `schrodinger.structure.Structure`
:param struct2: The comparison structure. A copy of this structure will
be reordered when the user clicks OK
:type modality: int
:param modality: The modality of the dialog. Default is Qt.NonModal,
other options may be Qt.WindowModal or Qt.ApplicationModal.
:param list struct1_atomic_constraints: List of constraints in text
format
:type default_guess_method: str
:param default_guess_method: Default guess method
"""
self.struct1 = struct1
self.struct2 = struct2
self.struct1_atomic_constraints = struct1_atomic_constraints
self.default_guess_method = default_guess_method
self.enable_2d_structure = enable_2d_structure
super().__init__(master, **kwargs)
self.setWindowModality(modality)
self.reordered_struct = None
self.reordered_order = None
[docs] def layOut(self):
"""
Lay out the custom widgets in this dialog.
"""
layout = self.mylayout
self.reorder_frame = ReorderAtomFrame(
parent=self,
reference_structure=self.struct1,
reference_structure_constraints=self.struct1_atomic_constraints,
comparison_structure=self.struct2,
layout=layout,
default_method=self.default_guess_method,
enable_2d_structure=self.enable_2d_structure)
size = self.size()
size.setHeight(600)
size.setWidth(800)
self.resize(size)
[docs] def accept(self):
"""
User has clicked accept, get the reordered structure, call the
accept_command with it and close the dialog.
:rtype: int or None
:return: QDialog.Accepted if a structure has been reordered, or None if
the reordering was incomplete and the dialog not closed
"""
done = self.reorder_frame.doneReordering()
if not done:
self.warning('Not all reference structure atoms are mapped to the '
'comparison structure')
return
else:
self.reordered_struct, self.reordered_order = \
self.reorder_frame.getReorderedStructure()
if not self.reordered_struct:
return
if maestro:
maestro.command('hidemarkers ' + MARKER_NAME)
maestro.invoke_picking_loss_callback()
# This must go last (MATSCI-9898)
if self.user_accept_function:
self.user_accept_function(self.reordered_struct,
self.reordered_order)
return swidgets.SDialog.accept(self)
[docs] def reject(self):
"""
User has clicked reject - do nothing but call custom reject method
:rtype: int
:return: QDialog.Reject
"""
if maestro:
maestro.command('hidemarkers ' + MARKER_NAME)
maestro.invoke_picking_loss_callback()
return swidgets.SDialog.reject(self)
mypanel = None
[docs]def panel():
"""Top-level function for bringing up the panel."""
global mypanel
# Only create the panel once
if not mypanel:
mypanel = ReorderAtomPanel()
mypanel.show()
# Bring the panel to the front
mypanel.raise_()
description = ('Description of script')
[docs]def main(*args):
parser = argparse.ArgumentParser(description=description, add_help=False)
parser.add_argument('-h',
'-help',
action='help',
default=argparse.SUPPRESS,
help='Show this help message and exit.')
parser.add_argument('-v',
'-version',
action='version',
default=argparse.SUPPRESS,
version=_version,
help="Show the program's version number and exit.")
options = parser.parse_args(args)
panel()
mypanel.run()
# For running outside of Maestro:
if __name__ == '__main__':
cmdline.main_wrapper(main, *sys.argv[1:])