Source code for schrodinger.application.matsci.sliceplot
"""
A plot that shows a 2D slice in a matplotlib plot and in the WS
Copyright Schrodinger, LLC. All rights reserved.
"""
import schrodinger
from schrodinger.graphics3d import common as graphics_common
from schrodinger.graphics3d import polygon
from schrodinger.ui.qt import smatplotlib
from schrodinger.ui.qt import swidgets
maestro = schrodinger.get_maestro()
[docs]class SlicePlot(swidgets.SFrame):
"""
A frame that plots a 2D slice of information about a structure and shows a
plane in the workspace that indicates where the 2D slice is coming from.
Widgets allow the user to specify what 2D slice to plot.
"""
LABEL_TEXT = 'Currently showing layer %d out of %d layers'
[docs] def __init__(self, nav_stretch=True):
"""
Create a SlicePlot instance
:type nav_stretch: bool
:param nav_stretch: Whether to add a stretch to the navigation widget
layout
"""
swidgets.SFrame.__init__(self)
layout = self.mylayout
# Navigation
self.nav_layout = swidgets.SHBoxLayout(layout=layout)
style = self.style()
self.back_btn = swidgets.SPushButton("",
layout=self.nav_layout,
command=self.moveBack)
self.back_btn.setIcon(style.standardIcon(style.SP_ArrowLeft))
self.layer_label = swidgets.SLabel(self.LABEL_TEXT % (0, 0),
layout=self.nav_layout)
self.forward_btn = swidgets.SPushButton("",
layout=self.nav_layout,
command=self.moveForward)
self.forward_btn.setIcon(style.standardIcon(style.SP_ArrowRight))
self.layer_dator = swidgets.SNonNegativeIntValidator(bottom=1)
self.layer_edit = swidgets.SLineEdit(validator=self.layer_dator,
always_valid=True,
layout=self.nav_layout,
width=60)
self.layer_edit.editingFinished.connect(self.changeLayer)
if nav_stretch:
self.nav_layout.addStretch()
# The actual plot
self.plot = smatplotlib.SmatplotlibCanvas(layout=layout)
self.axes = self.plot.fig.add_subplot(111)
self.decorateAxes()
self.legend_layout = swidgets.SVBoxLayout(layout=layout)
self.redisplay_btn = swidgets.SPushButton('Re-display All Atoms',
layout=layout,
command=self.redisplayAtoms)
# The slice indicator in the WS
self.ws_slab = None
self.group = graphics_common.Group()
self.hideSlab()
self.current_layer = 0
[docs] def redisplayAtoms(self):
"""
Show all the atoms in the workspace
"""
maestro.command('displayatom all')
[docs] def hideSlab(self):
"""
Hide the image of the slab in the WS
"""
self.group.hide()
[docs] def hideEvent(self, *args):
"""
Overwrite the parent method to ensure the Workspace slab gets hidden
"""
self.hideSlab()
return swidgets.SFrame.hideEvent(self, *args)
[docs] def closeEvent(self, *args):
"""
Overwrite the parent method to ensure the Workspace slab gets hidden
"""
self.hideSlab()
return swidgets.SFrame.closeEvent(self, *args)
[docs] def showEvent(self, *args):
"""
Overwrite the parent method to ensure the Workspace slab gets shown when
this tab is shown
"""
self.group.show()
return swidgets.SFrame.showEvent(self, *args)
[docs] def decorateAxes(self):
"""
Add labels, etc to the plot
"""
self.refreshPlot()
[docs] def reset(self):
"""
Reset this tab
"""
self.hideSlab()
self.group.clear()
self.current_layer = 0
self.max_layer = 0
self.setLayerWidgetStates()
[docs] def getAxisNum(self):
"""
Subclasses must implement this to return the index of the current axis
:rtype: int
:return: The index of the current axis
"""
raise NotImplementedError
def _get2DOrigins(self):
"""
Subclasses must implement this to return the XYZ coordinates of the
system origin and the length of the system in each direction
:rtype: list, list
:return: First list is the XYZ coordinates of the system origin, second
list is the length of the system in all three directions
"""
raise NotImplementedError
def _getSliceLocation(self):
"""
Get the location of the cross section along the axis
:rtype: float
:return: The axis coordinate of the center of the cross section
"""
raise NotImplementedError
[docs] def createSlab(self):
"""
Create a rectangle indicating the current layer in the workspace
"""
self.group.clear()
axis_num = self.getAxisNum()
origins, lengths = self._get2DOriginsAndLengths()
# Form the 4 vertices
vertices = []
vertices.append([origins[0], origins[1]])
vertices.append([origins[0] + lengths[0], origins[1]])
vertices.append([origins[0] + lengths[0], origins[1] + lengths[1]])
vertices.append([origins[0], origins[1] + lengths[1]])
# Offset the slice to the proper location along the axis
axis_num = self.getAxisNum()
axis_location = self._getSliceLocation()
for vertex in vertices:
vertex.insert(axis_num, axis_location)
# Create the polygon in the workspace
self.ws_slab = polygon.MaestroPolygon(vertices,
color='pink',
opacity=0.67)
self.group.add(self.ws_slab)
if self.isVisible():
self.group.show()
[docs] def setLayerLabelSize(self):
"""
Set the size of the label that gives the current layer so that it always
stays at its maximum width, avoiding widgets jittering left and right as
the label content changes
"""
label = swidgets.SLabel("")
maxhint = 0
for val in range(self.max_layer):
label.setText(self.LABEL_TEXT % (val, self.max_layer))
maxhint = max(maxhint, label.sizeHint().width())
self.layer_label.setFixedWidth(maxhint)
[docs] def plotData(self):
"""
Plot the data for the current layer and axis
"""
# Draw the plot on the GUI
self.refreshPlot()
self.createSlab()
[docs] def refreshPlot(self):
self.plot.fig.tight_layout()
self.plot.draw()
[docs] def moveBack(self):
"""
Plot the layer previous to the current one
"""
self.current_layer -= 1
self.setLayerWidgetStates()
self.plotData()
[docs] def moveForward(self):
"""
Plot the layer after the current one
"""
self.current_layer += 1
self.setLayerWidgetStates()
self.plotData()
[docs] def changeLayer(self):
"""
Plot the layer indicated in the layer line edit
"""
new_layer = self.layer_edit.int() - 1
if new_layer == self.current_layer:
return
self.current_layer = self.layer_edit.int() - 1
self.setLayerWidgetStates()
self.plotData()
[docs] def displayLocalAtoms(self, indexes):
"""
Set visibility and representation of atoms based on proximity
:type indexes: list
:param indexes: The atom indexes for the central atoms to display
"""
if not indexes:
maestro.command('displayonlyatom not all')
return
asl = 'atom.n ' + ','.join([str(x) for x in indexes])
maestro.command('workspaceselectionreplace ' + asl)
maestro.command('repatom rep=cpk atom.selected')
maestro.command('repbond rep=wire')
try:
maestro.command('repatombonds not atom.selected')
except maestro.MaestroCommand:
# There are no atoms not selected
pass
maestro.command(
'displayonlyatom within 5 atom.selected or atom.selected')