#Author: Pat Lorton
import math
from schrodinger.infra import canvas2d
from schrodinger.Qt.QtCore import QEvent
from schrodinger.Qt.QtCore import QLineF
from schrodinger.Qt.QtCore import QPointF
from schrodinger.Qt.QtCore import QRectF
from schrodinger.Qt.QtCore import Qt
from schrodinger.Qt.QtCore import QTimeLine
from schrodinger.Qt.QtCore import QTimer
from schrodinger.Qt.QtCore import pyqtSignal
from schrodinger.Qt.QtGui import QBrush
from schrodinger.Qt.QtGui import QColor
from schrodinger.Qt.QtGui import QCursor
from schrodinger.Qt.QtGui import QMouseEvent
from schrodinger.Qt.QtGui import QPainter
from schrodinger.Qt.QtGui import QPainterPath
from schrodinger.Qt.QtGui import QPen
from schrodinger.Qt.QtGui import QPolygonF
from schrodinger.Qt.QtWidgets import QGraphicsItem
from schrodinger.Qt.QtWidgets import QGraphicsPathItem
from schrodinger.Qt.QtWidgets import QGraphicsPolygonItem
from schrodinger.Qt.QtWidgets import QGraphicsRectItem
from schrodinger.Qt.QtWidgets import QGraphicsScene
from schrodinger.Qt.QtWidgets import QGraphicsSimpleTextItem
from schrodinger.Qt.QtWidgets import QGraphicsTextItem
from schrodinger.Qt.QtWidgets import QGraphicsView
from schrodinger.Qt.QtWidgets import QMessageBox
from schrodinger.Qt.QtWidgets import QStyle
from schrodinger.Qt.QtWidgets import QStyleOptionGraphicsItem
from schrodinger.Qt.QtWidgets import QVBoxLayout
from schrodinger.Qt.QtWidgets import QWidget
from schrodinger.ui.qt import network_visualizer
from . import structure2d
# Pixel scaling factor, so that raster graphics will be rendered with adequate
# resolution
PSF = 10
#### Spring layout parameters ####
STARTING_TEMPERATURE = .1
ITERATIONS = 100
SCALE = 3.0 # Scale of distances to node size. Higher number is greater separation
PUSH = 1 # Multiplier for node-node repulsion
PUSHEXP = 6 # Exponential dependence of node-node repulsion
DEFAULT_EDGE_COLOR_HEX = '#000000'
SELECTED_EDGE_COLOR_HEX = '#ffd419'
# Show 'bad' edges using orange color because it is easier to destinguish from
# default color for colorblind individuals (PANEL-13630)
BAD_EDGE_COLOR_HEX = '#ef8635'
NO_PEN = QPen()
NO_PEN.setStyle(Qt.NoPen)
[docs]def calculateArrow(end_point, arrow_end_point, mag=4, skew=4):
"""
Calculate how to draw an arrow at the end_point of a line,
this requires the line's other endpoint so that we know the angle at
which the arrow is to be drawn.
:return: a QPolygonF containing the arrow
:param mag: controls the magnitude of the arrow head smaller value smaller
smaller head
:type mag: int
"""
assert mag > 0
l_orig = QLineF(end_point, arrow_end_point)
dx = arrow_end_point.x() - end_point.x()
dy = (arrow_end_point.y() - end_point.y())
# Magnitude of the line i.e. the length of the line between the points
orig_mag = math.sqrt(dx**2 + dy**2)
#theta = math.acos(dx/orig_mag)
#if dy >= 0:
# theta = (math.pi * 2.0) - theta
#arrowSize = 20.0
#rp1 = QPointF(arrow_end_point) + QPointF(math.sin(theta + math.pi / 3.0) * arrowSize, math.cos(theta + math.pi /3) * arrowSize)
#rp2 = QPointF(arrow_end_point) + QPointF(math.sin(theta + math.pi - math.pi / 3.0) * arrowSize, math.cos(theta + math.pi - math.pi / 3.0) * arrowSize)
mag_ratio = orig_mag / mag
if mag_ratio == 0:
#was .001
mag_ratio = 0.001
#Lower the distance of dx/dy to be on the magnitude of the arrows length
dx /= mag_ratio
dy /= mag_ratio
rp1 = QPointF(arrow_end_point)
rp2 = QPointF(arrow_end_point)
rp1 += QPointF(-(dx + dy), -(dy - dx))
rp2 += QPointF(-(dx - dy), -(dy + dx))
p = QPointF(0, 0)
for i in range(skew):
arrow_head_base = QLineF(rp1, rp2)
if l_orig.intersect(arrow_head_base, p) == QLineF.BoundedIntersection:
rp1 = QPointF(p)
rp2 = QPointF(p)
rp1 += QPointF(-(dx + dy), -(dy - dx))
rp2 += QPointF(-(dx - dy), -(dy + dx))
return QPolygonF([arrow_end_point, rp1, rp2, arrow_end_point])
[docs]class NetworkNode(QGraphicsRectItem):
"""
This is an abstract class for the Node's of the graph, you can subclass
from this class to render the data "val", whatever it may be.
"""
model_scale = 120 * PSF
default_size_factor = 50
[docs] def __init__(self, model, network):
"""
:param model: the model node corresponding to this view node
:type mode: network_visualizer.Node
:param network: the graph view to which this node view belongs
:type network: NetworkViewer
"""
super().__init__()
self._setupOptions()
self._setupDraw()
self.hovering = False
self.creatingConnection = False
self.connectionSelected = False
self.network = network
self.model = model
self.syncModel()
def _setupOptions(self):
flags = QGraphicsItem.ItemIsSelectable | QGraphicsItem.ItemIsMovable
self.setFlags(flags)
self.setZValue(1) #Make sure nodes are in front of connections
self.setAcceptHoverEvents(True)
def _setupDraw(self):
default_length = self.default_size_factor * PSF
self.setRect(0, 0, default_length, default_length)
self.qp = QPen()
self.setPen(self.qp)
self.qb = QBrush(QColor(255, 255, 255))
self.setBrush(self.qb)
# Make this more like the color of the project table
self.qb_selected = QBrush(QColor(255, 253, 175))
self.qb_hovering = QBrush(QColor(255, 255, 255))
self.qb_hover_ok = QBrush(QColor(60, 230, 60)) # MOVE THESE TO NETWORK
self.qb_hover_bad = QBrush(QColor(230, 60, 60))
[docs] def getEdges(self):
"""
Retrieve all the edges connected to this node.
"""
return self.network.getEdges(self.model)
#===========================================================================
# Model-view synchronization
#===========================================================================
[docs] def syncModel(self):
pos = self.model.pos()
if pos is not None:
pos = self.scalePosFromModel(pos)
self.setPos(pos)
[docs] def scalePosToModel(self, pos):
x = pos.x()
y = pos.y()
return x / self.model_scale, y / self.model_scale
[docs] def scalePosFromModel(self, pos):
x, y = pos
return QPointF(x * self.model_scale, y * self.model_scale)
#===========================================================================
# Events
#===========================================================================
[docs] def hoverEnterEvent(self, e):
self.hovering = True
self.network.nodeHovered(self)
# Cache to prevent constant repainting while hovering (and dragging).
# This causes bad anti-aliasing of the rounded rectangle, so only turn
# the cache on when hovering
self.setCacheMode(QGraphicsItem.DeviceCoordinateCache)
self.update()
[docs] def hoverLeaveEvent(self, e):
self.hovering = False
self.network.nodeUnhovered(self)
self.setCacheMode(QGraphicsItem.NoCache)
self.update()
[docs] def mousePressEvent(self, e):
"""
Accept mouse events so that selection doesnt get lost (QTBUG-10138)
Without this fix, when context menu is up, mouse events get
propagated to the parent and it behaves as if you did not click on a
node. This causes th
"""
e.accept()
super().mousePressEvent(e)
[docs] def mouseMoveEvent(self, e):
"""
Move the item and tell the network to redraw connection lines
"""
super().mouseMoveEvent(e)
nodes = set()
for item in self.scene().selectedItems():
if isinstance(item, NetworkNode):
x, y = self.scalePosToModel(item.scenePos())
item.model.setPos(x, y)
nodes.add(item.model)
edges = self.network.getEdges(nodes)
for edge in edges:
edge.recalculateShape()
[docs] def mouseDoubleClickEvent(self, e):
self.network.zoomItem(self)
[docs] def mouseReleaseEvent(self, e):
QGraphicsRectItem.mouseReleaseEvent(self, e)
self.network.expandToNodes()
#===========================================================================
# Draw operations
#===========================================================================
[docs] def centerPos(self):
return self.boundingRect().center() + self.pos()
[docs] def paint(self, painter, option, widget):
if self.hovering:
self.setBrush(self.qb_hovering)
if self.isSelected():
self.setBrush(self.qb_selected)
else:
self.setBrush(self.qb)
painter.setBrush(self.brush())
painter.drawRoundedRect(self.rect(), 8 * PSF, 8 * PSF)
def __str__(self):
return f'<{self.__class__.__name__}("{self.model.name}")>'
def __repr__(self):
return self.__str__()
[docs]class TextNetworkNode(NetworkNode):
"""
Network node that optionally displays the string representation of a node
attribute as text.
"""
text_y_pos = 15 * PSF
text_size = 10 * PSF
text_offset = 5
[docs] def __init__(self, model, network, text_attr=None):
"""
:param model: the model node corresponding to this view node
:type mode: Node
:param network: the graph view to which this node view belongs
:type network: NetworkViewer
:param text_attr: the name of the model node attribute which will be
used as the source of the text for this view node (optional)
:type text_attr: str or NoneType
"""
self.text_attr = text_attr
self.text_item = None
super().__init__(model, network)
def _addLabel(self):
self.text_item = QGraphicsTextItem('', self)
self._setTextSize(self.text_size)
def _setTextSize(self, size):
"""
Sets the font size for the text item.
:param size: text size
:type size: int
"""
self.text_size = size
if self.text_item:
font = self.text_item.font()
font.setPointSize(self.text_size)
self.text_item.setFont(font)
self._updateNodeSize()
[docs] def syncModel(self):
super().syncModel()
if self.text_attr:
text = getattr(self.model, self.text_attr)
else:
text = self.model.name
if text:
self._setText(text)
def _updateNodeSize(self):
"""
Update node width to fit the size of the text in the node.
Note that this operation should preserve the position of the center of
the node.
"""
text_width = self.text_item.boundingRect().width()
new_height = self.default_size_factor * PSF
new_width = max(text_width + 2 * self.text_offset * PSF, new_height)
nrect = self.rect()
width, height = nrect.width(), nrect.height()
if (width, height) == (new_width, new_height):
return
self.setRect(0, 0, new_width, new_height)
self.setPos(
self.x() + (width - new_width) / 2,
self.y() + (height - new_height) / 2
) # yapf: disable
self._updateTextPosition()
self._updateEdgePositions()
def _updateEdgePositions(self):
"""
Update positions for all view edges connected to this node.
"""
for edge in self.network.getEdges([self.model]):
if edge is not None:
edge.recalculateShape()
def _setText(self, text):
"""
Sets the text displayed in the node
:param text: the text to show
:type text: str
"""
if not self.text_item:
self._addLabel()
text = str(text)
self.text = text
self.text_item.setPlainText(text)
self._updateNodeSize()
self._updateTextPosition()
def _updateTextPosition(self):
"""
Reposition the text inside the node to be Horizontally and vertically
centered.
"""
nrect, trect = self.rect(), self.text_item.boundingRect()
text_width, text_height = trect.width(), trect.height()
node_width, node_height = nrect.width(), nrect.height()
# Ensure text does not fall out the bottom of the node
text_y_pos = min(self.text_y_pos, node_height - text_height - 2 * PSF)
# Horizontally center the text. There is a small offset for margin
self.text_item.setPos(
0.5 * (node_width - text_width) - self.text_offset, text_y_pos)
[docs]class PosTextNetworkNode(TextNetworkNode):
text_size = 5 * PSF
[docs] def syncModel(self):
NetworkNode.syncModel(self)
pos = self.model.pos()
if pos is not None:
x, y = pos
self._setText("(%0.1f, %0.1f)" % (x, y))
[docs] def mouseMoveEvent(self, e):
"""Move the item and tell the network to redraw connection lines"""
super().mouseMoveEvent(e)
self.syncModel()
[docs]class SmilesNetworkNode(TextNetworkNode):
"""
This renders the input variable 'val' as a 2D-Structure, assuming it's
a smiles that can be converted to a 2d structure.
"""
text_y_pos = 1000 * PSF # Display at very bottom of node
text_size = 6 * PSF
[docs] def __init__(self, x, y, val, allow_movement, text=''):
super().__init__(x, y, text, allow_movement)
self.val = val
r = QRectF(0, 0, 50 * PSF, 50 * PSF)
self.struct_item = structure2d.structure_item(rect=r)
self.struct_item.setParentItem(self)
self.struct_item.setPos(0, 0)
try:
self.struct_item.chmmol = canvas2d.ChmMol.fromSMILES(val)
except Exception:
raise ValueError("Could not convert SMILES to 2D structure")
canvas2d.Chm2DCoordGen.generateAndApply(
self.struct_item.chmmol, canvas2d.ChmAtomOption.H_ExplicitOnly,
canvas2d.ChmCoordGenOption.Rendering)
self.struct_item.generate_picture()
self.setRect(self.struct_item.boundingRect().adjusted(-5, -5, 5, 5))
[docs]class NetworkConnectionLabel(QGraphicsTextItem):
[docs] def __init__(self, *args, **kwargs):
QGraphicsSimpleTextItem.__init__(self, *args, **kwargs)
font = self.font()
font.setPointSize(6 * PSF)
self.setFont(font)
[docs]class NetworkEdge(QGraphicsPathItem):
"""
Right now this is just like a regular line, but will allow for
extension to avoid connection/node intersections if so desired in the
future.
"""
_triangle = QPolygonF([
QPointF(-100, -50),
QPointF(-100, 50),
QPointF(0, 0),
QPointF(-100, -50)
])
[docs] def __init__(self,
node1,
node2,
network,
dotted=False,
opacity=1.0,
label=None):
QGraphicsPathItem.__init__(self)
self.node1 = node1
self.node2 = node2
self.network = network
self.arrowhead_node = None
self.setFlags(QGraphicsItem.ItemIsSelectable)
self.setAcceptHoverEvents(True)
self.show_as_bad = False
pen_type = Qt.DotLine if dotted else Qt.SolidLine
color = QColor(DEFAULT_EDGE_COLOR_HEX)
self.qp = _get_formatted_pen(pen_type, color)
self.qb = QBrush(color)
color = QColor(SELECTED_EDGE_COLOR_HEX)
self.qp_selected = _get_formatted_pen(pen_type, color)
self.qb_selected = QBrush(color)
color = QColor(BAD_EDGE_COLOR_HEX)
self.qp_bad = _get_formatted_pen(pen_type, color)
self.qb_bad = QBrush(color)
# Set default pen and brush
self.setPen(self.qp)
self.setBrush(self.qb)
self.label = None
if label: # Also catches empty strings
self.setLabel(label)
self.recalculateShape()
self.original_opacity = opacity
self.setOpacity(opacity)
@property
def model_edge(self):
"""
:return: the model edge corresponding with this view edge
:rtype: Edge
"""
return self.network.model.getEdge(self.node1.model, self.node2.model)
def _getModelOpacity(self):
"""
Get opacity value from the model edge.
:return: the opacity, if available
:rtype: `float` or `None`
"""
return self.model_edge.getData('opacity')
[docs] def getOtherNode(self, node):
if node == self.node1:
return self.node2
elif node == self.node2:
return self.node1
return None
[docs] def syncModel(self):
self.recalculateShape()
opacity = self._getModelOpacity()
if opacity:
self.setOpacity(opacity)
self.original_opacity = opacity
[docs] def setLabel(self, text, html=False):
"""
Show text along edge
:param text: label text
:type text: str
"""
if not self.label:
self.label = NetworkConnectionLabel(self)
if html:
self.label.setHtml(text)
else:
self.label.setPlainText(text)
self.recalculateShape()
[docs] def setArrowhead(self, enable, head_node=None):
if not enable:
self.arrowhead_node = None
self.recalculateShape()
return
if head_node is None:
head_node = self.node1
self.arrowhead_node = head_node
self.arrowhead = QGraphicsPolygonItem(self._triangle, parent=self)
self.recalculateShape()
[docs] def showBadEdge(self, is_bad):
"""
Mark this edge as 'bad'. It will be shown with red color.
:param is_bad: True or False to indicate whether this edge is 'bad'.
:type is_bad: bool
"""
self.show_as_bad = is_bad
[docs] def mouseDoubleClickEvent(self, e):
self.network.zoomItem(self)
[docs] def hoverEnterEvent(self, e):
self.hovering = True
self.network.edgeHovered(self)
self.update()
[docs] def hoverLeaveEvent(self, e):
self.hovering = False
self.network.edgeUnhovered(self)
self.update()
[docs] def boundingRect(self): # TRY DELETING
return QGraphicsPathItem.boundingRect(self)
[docs] def paint(self, painter, option, widget):
myoption = QStyleOptionGraphicsItem(option)
myoption.state &= not QStyle.State_Selected
if self.isSelected():
self.setPen(self.qp_selected)
self.setBrush(self.qb_selected)
else:
if self.show_as_bad:
self.setPen(self.qp_bad)
self.setBrush(self.qb_bad)
else:
self.setPen(self.qp)
self.setBrush(self.qb)
if self.arrowhead_node is not None:
self.arrowhead.setBrush(self.brush())
self.arrowhead.setPen(NO_PEN)
QGraphicsPathItem.paint(self, painter, myoption, widget)
[docs] def recalculateShape(self):
"""
This function recalculates the shape (line, perhaps with arrow)
to draw for the connection between two nodes.
"""
pos1 = self.node1.centerPos()
pos2 = self.node2.centerPos()
path = QPainterPath(pos1)
path.lineTo(pos2)
self.setPath(path)
dx = pos2.x() - pos1.x()
dy = pos2.y() - pos1.y()
if abs(dx) < .001:
dx = .001
slope = dy / dx
self.angle = math.atan(slope) * 180 / 3.14159265
if self.label:
length = math.sqrt(dx * dx + dy * dy)
label_offset = 0.5 * (length - self.label.boundingRect().width())
fraction = label_offset / length
x_offset = fraction * dx
y_offset = fraction * dy
if dx > 0:
xpos = pos1.x() + x_offset
ypos = pos1.y() + y_offset
else:
xpos = pos2.x() - x_offset
ypos = pos2.y() - y_offset
self.label.setPos(xpos, ypos)
self.label.setRotation(self.angle)
if self.arrowhead_node:
head_node = self.arrowhead_node
ipos = self.nodeIntersectionPoint(head_node)
if not ipos:
return
self.arrowhead.setPos(ipos)
angle = self.angle
flip = dx > 0
if self.arrowhead_node == self.node1:
flip = not flip
if not flip:
angle += 180
self.arrowhead.setRotation(angle)
[docs] def nodeIntersectionPoint(self, node):
line = QLineF(self.node1.centerPos(), self.node2.centerPos())
rect = node.rect().adjusted(node.x(), node.y(), node.x(), node.y())
return intersect_line_and_rect(line, rect)
def __str__(self):
name1, name2 = self.node1.model.name, self.node2.model.name
return f'<{self.__class__.__name__}("{name1}" - "{name2}")>'
def __repr__(self):
return self.__str__()
[docs]def intersect_line_and_rect(line, rect):
"""
Returns the intersection point of a line and 4 line segments,
if it intersects multiple segments, only one intersection will be
returned.
"""
# Lines for each side of the incoming rect
side1 = QLineF(rect.topLeft(), rect.topRight())
side2 = QLineF(rect.topRight(), rect.bottomRight())
side3 = QLineF(rect.bottomRight(), rect.bottomLeft())
side4 = QLineF(rect.bottomLeft(), rect.topLeft())
for side in [side1, side2, side3, side4]:
p = QPointF(0, 0)
if line.intersect(side, p) == QLineF.BoundedIntersection:
return p
class _PickingArrow(QGraphicsPathItem):
def __init__(self, node1):
QGraphicsPathItem.__init__(self)
qb = QBrush(QColor(0, 0, 0))
qp = QPen()
qp.setWidth(2 * PSF)
qp.setJoinStyle(Qt.MiterJoin)
self.setPen(qp)
self.setBrush(qb)
self.setZValue(0)
self.node1 = node1
def setEndPos(self, pos):
pos1 = self.node1.centerPos()
pos2 = pos
p = QPainterPath(pos1)
p.lineTo(pos2)
p.addPolygon(calculateArrow(pos1, pos2))
self.setPath(p)
self.update()
class _localGraphicsScene(QGraphicsScene):
mouseMoved = pyqtSignal(QEvent)
def mouseMoveEvent(self, e):
self.mouseMoved.emit(e)
super().mouseMoveEvent(e)
class _localGraphicsView(QGraphicsView):
"""
:ivar geometryChanged: A signal emitted whenever network view
geometry has changed.
:vartype geometryChanged: `PyQt5.QtCore.pyqtSignal`
"""
keyPressed = pyqtSignal(QEvent)
scaleChanged = pyqtSignal()
geometryChanged = pyqtSignal()
def __init__(self, scene, parent):
QGraphicsView.__init__(self, scene, parent)
self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
def keyPressEvent(self, e):
self.keyPressed.emit(e)
def wheelEvent(self, e):
# angleDelta() returns a QPoint obj.
delta = e.angleDelta().y()
factor = 1.41**(delta / 240.0)
self.setTransformationAnchor(self.AnchorUnderMouse)
self.scale(factor, factor)
self.setTransformationAnchor(self.AnchorViewCenter)
self.scaleChanged.emit()
def mousePressEvent(self, event):
pos = event.pos()
button, item = event.button(), self.itemAt(pos)
if button == Qt.RightButton and item is None:
# This modification allows the user to move the network view
# scene by right click-dragging on an empty part of the view.
self.setDragMode(QGraphicsView.ScrollHandDrag)
event = QMouseEvent(QEvent.MouseButtonPress, pos, Qt.LeftButton,
Qt.LeftButton, Qt.NoModifier)
elif button == Qt.RightButton and isinstance(item, NetworkEdge):
# Add the control modifier when the user right-clicks an edge to
# prevent right-clicking from deselecting everything.
buttons, mods = event.buttons(), event.modifiers()
mods = mods | Qt.ControlModifier
event = QMouseEvent(event.type(), pos, button, buttons, mods)
super().mousePressEvent(event)
def mouseReleaseEvent(self, e):
if e.button() == Qt.RightButton:
e = QMouseEvent(QEvent.MouseButtonPress, e.pos(), Qt.LeftButton,
Qt.LeftButton, Qt.NoModifier)
QGraphicsView.mouseReleaseEvent(self, e)
self.setDragMode(QGraphicsView.RubberBandDrag)
def zoomIn(self):
"""
Changes scale by a fixed factor to zoom in.
"""
factor = 1.1
self.scale(factor, factor)
self.scaleChanged.emit()
def zoomOut(self):
"""
Changes scale by a fixed factor to zoom out.
"""
factor = 1.0 / 1.1
self.scale(factor, factor)
self.scaleChanged.emit()
def resizeEvent(self, e):
QGraphicsView.resizeEvent(self, e)
self.geometryChanged.emit()
#===============================================================================
# Network View Classes
#===============================================================================
[docs]class NetworkViewer(network_visualizer.AbstractNetworkView, QWidget):
"""
A network diagram representation of the model.
"""
pickingModeChanged = pyqtSignal(bool)
stopPickingConnection = pyqtSignal()
[docs] def __init__(self, node_class=NetworkNode, parent=None, model=None):
QWidget.__init__(self, parent=parent)
network_visualizer.AbstractNetworkView.__init__(self)
self.timer = None
self.last_zoom_item = None
self._setOptions()
self._setupDraw()
self._setDefaults()
self.node_class = node_class
self.setModel(model)
self.scene.selectionChanged.connect(self.onSelectionChanged)
self.view.keyPressed.connect(self.onKeyPressed)
self.scene.mouseMoved.connect(self.mouseMoveEvent)
def _setOptions(self):
self.allow_movement = True
self.allow_node_deletion = True
def _setDefaults(self):
self.picking_new_connection = False
self.picking_arrow = None
def _setupDraw(self):
self.scene = _localGraphicsScene()
self.view = _localGraphicsView(self.scene, self.parent())
layout = QVBoxLayout()
layout.addWidget(self.view)
self.setLayout(layout)
self.view.setRenderHints(QPainter.Antialiasing)
self.view.setDragMode(QGraphicsView.RubberBandDrag)
self.picking_arrow = None #Will hold a QGraphicsPathItem
self.min_scale = 0.3 * PSF
self.max_scale = 10.0 * PSF
self.scale_timeline = QTimeLine(300)
#===========================================================================
# Model-View Synchronization
#===========================================================================
[docs] def getSignalsAndSlots(self, model):
# Add the positionChanged signal for this view. See parent class for
# more information.
ss_list = super().getSignalsAndSlots(model)
ss_list.append((model.signals.positionChanged, self.syncNodesChanged))
return ss_list
[docs] def setModel(self, model, fit=True):
"""
:param fit: whether to fit the zoom to the new model
:type fit: bool
"""
super().setModel(model)
self.update()
if fit:
self.setDefaultScale()
#===========================================================================
# Required view implementation
#===========================================================================
[docs] def makeNodes(self, nodes):
return {node: self.node_class(node, self) for node in nodes}
[docs] def addNodes(self, viewnodes):
for viewnode in viewnodes:
self.scene.addItem(viewnode)
[docs] def removeNodes(self, viewnodes):
for viewnode in viewnodes:
self.nodeUnhovered(viewnode)
self.scene.removeItem(viewnode)
[docs] def updateNodes(self, nodes):
for node in nodes:
viewnode = self.nodes.get(node)
if viewnode is None:
return
viewnode.syncModel()
[docs] def makeEdges(self, edges):
edge_map = {}
for edge in edges:
viewnode1 = self.nodes[edge[0]]
viewnode2 = self.nodes[edge[1]]
edge_map[edge] = NetworkEdge(viewnode1, viewnode2, self)
return edge_map
[docs] def addEdges(self, viewedges):
for viewedge in viewedges:
self.scene.addItem(viewedge)
[docs] def removeEdges(self, viewedges):
for viewedge in viewedges:
self.edgeUnhovered(viewedge)
self.scene.removeItem(viewedge)
[docs] def updateEdges(self, edges):
for edge in edges:
view_edge = self.getEdge(edge)
if view_edge is not None:
view_edge.syncModel()
[docs] def selectItems(self, selected_view_objects):
for item in self.scene.items():
item.setSelected(item in selected_view_objects)
#===========================================================================
# Interactive UI
#===========================================================================
[docs] def edgeHovered(self, edge):
"""
Override this for connection hover behavior
"""
[docs] def edgeUnhovered(self, edge):
"""
Override this for connection hover behavior
"""
[docs] def exploreEdge(self, edge):
"""
Override this for edge context menu behavior.
"""
[docs] def nodeHovered(self, hovered_node):
"""
When a node is hovered, the all nodes with degree of separation greater
than 1 are dimmed in order to highlight the immediate neighbors.
"""
dim_nodes = self._getNonNeighborNodes(hovered_node)
self.timer = QTimer()
self.timer.timeout.connect(self._getHoverFunc(dim_nodes))
self.timer.setSingleShot(True)
self.timer.start(350)
def _getHoverFunc(self, dim_nodes):
"""
Returns a slot function to dim the requested nodes when invoked
:param dim_nodes: the nodes to dim
:type dim_nodes: set of Nodes
"""
def hoverfunc():
self.setNodesOpacity(dim_nodes, 0.1)
return hoverfunc
[docs] def nodeUnhovered(self, hovered_node):
"""
When a node is unhovered, return the dimmed nodes to the normal state.
"""
if self.timer:
self.timer.stop(
) # Node was unhovered before timer was fired; abort
self.setNodesOpacity(self.model.getNodes(), 1.0)
def _getNonNeighborNodes(self, in_node):
"""
This will return a set of all nodes with degree of separation greater
than 1.
"""
allnodes = self.model.getNodes()
neighbors = self.model.getNeighbors(in_node.model)
node_set = allnodes.difference(neighbors)
node_set.remove(in_node.model)
return node_set
[docs] def setNodesOpacity(self, nodes, opacity):
"""
Change the opacity of a set of nodes and their edges.
:param nodes: the nodes to change opacity
:type nodes: iterable of NetworkNode
:param opacity: opacity value from 0 to 1.
:type opacity: float
"""
for node in nodes:
try:
viewnode = self.nodes[node]
except KeyError: # The node may no longer exist
continue
viewnode.setOpacity(opacity)
edges = self.model.getEdges(node)
for edge in edges:
view_edge = self.getEdge(edge)
if view_edge is None: # The edge may no longer exist
continue
view_edge.setOpacity(opacity * view_edge.original_opacity)
[docs] def calcItemZoomFactor(self, item):
"""
Calculates the raw zoom factor necessary for the specified item to fill
most of the view
:param item: the item to fit in the view
:type item: QtWidgets.QGraphicsItem
"""
vp = self.view.viewport()
vp_size = min(vp.width(), vp.height())
item_rect = item.sceneBoundingRect()
item_size = max(item_rect.width(), item_rect.height())
ratio = vp_size / item_size
return 0.7 * ratio # Leave some room around the item
[docs] def calcFitZoomFactor(self):
"""
Calculates a raw zoom factor for the view to fit all the items in the
scene
"""
vp = self.view.viewport()
vp_size = min(vp.width(), vp.height())
item_rect = self.scene.itemsBoundingRect()
item_size = max(item_rect.width(), item_rect.height())
ratio = vp_size / item_size
return ratio
[docs] def zoomItem(self, item):
"""
Zoom into a specific item with animation.
:param item: item to zoom
:type item: QGraphicsItem
"""
movement_backup = self.allow_movement
self.setAllowMovement(False)
def restoreMovement():
self.setAllowMovement(movement_backup)
self.view.scaleChanged.emit()
def finishZoom():
self.view.centerOn(item)
zoom_factor = self.calcItemZoomFactor(item)
self.setRawScale(zoom_factor)
self.scale_timeline.finished.connect(restoreMovement)
def finishUnzoom():
self.view.centerOn(self.scene.itemsBoundingRect().center())
restoreMovement()
#If we're already zoomed in, then we need to zoom out before zooming in
zoom_factor = self.calcItemZoomFactor(item)
if self.view.transform().m11() > 0.7 * zoom_factor:
self.setRawScale(self.calcFitZoomFactor())
if not self.last_zoom_item == item:
self.scale_timeline.finished.connect(finishZoom)
else:
self.scale_timeline.finished.connect(finishUnzoom)
else:
finishZoom()
self.last_zoom_item = item
[docs] def mouseMoveEvent(self, e):
"""
If we're picking and the mouse is moving then update the path
of the arrow which is used for picking.
"""
if self.picking_arrow:
self.picking_arrow.setEndPos(e.scenePos())
#===========================================================================
# Adding edge from UI
#===========================================================================
[docs] def createEdge(self, vnode1, vnode2):
"""
MVC "controller" function to add a new edge to the graph between node1
and node2. This can be overridden to implement other edge creation
functionality.
:param vnode1: First node
:type vnode1: NetworkNode
:param vnode2: Second node
:type vnode2: NetworkNode
"""
self.model.setUndoPoint()
self.model.addEdge(vnode1.model, vnode2.model)
def _pickFirstConnectionNode(self, node):
self.picking_arrow = _PickingArrow(node)
self.scene.addItem(self.picking_arrow)
def _pickSecondConnectionNode(self, node2):
self.scene.removeItem(self.picking_arrow)
node1 = self.picking_arrow.node1
self.picking_arrow = None
self.picking_new_connection = False
valid, reason = self.model.getEdgeApproval(node1.model, node2.model)
if valid:
self.createEdge(node1, node2)
self.setNodesOpacity(self.model.getNodes(), 1.0)
else:
QMessageBox.warning(self, "Connection Not Allowed", reason)
def _startPickingConnection(self):
self.setCursor(QCursor(Qt.CrossCursor))
self.view.setDragMode(QGraphicsView.NoDrag)
self.picking_new_connection = True
def _stopPickingConnection(self):
self.setCursor(QCursor(Qt.ArrowCursor))
self.view.setDragMode(QGraphicsView.RubberBandDrag)
self.picking_new_connection = False
if self.picking_arrow:
self.scene.removeItem(self.picking_arrow)
self.picking_arrow = None
self.stopPickingConnection.emit()
#===========================================================================
# Selection
#===========================================================================
[docs] def selectedNodes(self):
"""
Get a list of currently selected view nodes
"""
nodes = []
for item in self.scene.selectedItems():
if isinstance(item, NetworkNode):
nodes.append(item)
return nodes
[docs] def selectedEdges(self):
"""
Get a list of currently selected view edges
"""
edges = set()
for item in self.scene.selectedItems():
if isinstance(item, NetworkEdge):
edges.add(item)
return edges
[docs] def onSelectionChanged(self):
"""
This is used to process UI changes in selection and updating the model
from the view selection. This is not for synchronizing the view to the
model selection. For that, use self.syncSelection().
"""
if self.skip_selectionChanged:
return
viewnodes = self.selectedNodes()
nodes = {n.model for n in viewnodes}
edges = {view_edge.model_edge for view_edge in self.selectedEdges()}
self.model.setSelectedObjs(nodes.union(edges), self)
if not self.picking_new_connection:
return
if len(viewnodes) != 1: # This would be a weird state
self.setPickingMode(False)
return
node = viewnodes.pop()
if not self.picking_arrow: # Pick first node
self._pickFirstConnectionNode(node)
else:
self._pickSecondConnectionNode(node)
self.setPickingMode(False)
[docs] def setPickingMode(self, state):
"""
Begins picking nodes to add connections, can be connected to a checkbox.
If a python checkbox is connected with "toggled(bool)" to this, the
checkbox also needs to have a slot to turn the checkbox off which will
be connected to the signal:
SIGNAL("stopPickingConnection()")
This is emitted once a connection has been created and the checkbox
should be turned off.
"""
self.model.setSelectedObjs([])
if state:
self._startPickingConnection()
else:
self._stopPickingConnection()
self.pickingModeChanged.emit(state)
#===========================================================================
# View to model operations
#===========================================================================
[docs] def deleteSelectedItems(self):
"""
Process item deletion request. Deletions are done in the model only;
the view is not modified here. Changing the model will automatically
result in a corresponding update to the view.
"""
if not self.scene.selectedItems():
return
self.model.deleteSelectedItems(include_nodes=self.allow_node_deletion)
[docs] def onKeyPressed(self, e):
if e.key() == Qt.Key_Delete or e.key() == Qt.Key_Backspace:
self.deleteSelectedItems()
elif e.key() == Qt.Key_Escape:
self.model.setSelectedObjs(set())
self.setPickingMode(False)
#===========================================================================
# Miscellaneous
#===========================================================================
[docs] def setAllowMovement(self, val):
"""
Allow nodes to be moved with click+drag
"""
self.allow_movement = val
flags = QGraphicsItem.ItemIsSelectable
if self.allow_movement:
flags |= QGraphicsItem.ItemIsMovable
for node in self.nodes:
viewnode = self.nodes[node]
viewnode.setFlags(flags)
[docs] def setNodeDeletion(self, val):
"""
Enable deletion of nodes
"""
self.allow_node_deletion = val
[docs] def setNetworkNodeDefaultClass(self, cl):
"""
Set the default view node class for
"""
assert issubclass(cl, NetworkNode) #Only allow NetworkNode subclasses
self.node_class = cl
[docs] def getFitRect(self, node_set=None):
"""
Calculate the rectangle that contains the nodes. If a node_set is
specified, only those nodes will be considered. Otherwise, all nodes
will be used. This function takes into account the size of the nodes.
:param node_set: a subset of nodes to fit
:type node_set: iterable
"""
if node_set is None:
node_set = list(self.nodes.values())
group = self.scene.createItemGroup(list(node_set))
rect = group.boundingRect()
self.scene.destroyItemGroup(group)
return rect
def _padRect(self, frect):
"""
Add margins to frect to allow zooming on nodes near edge of scene
:param frect: Current fit rectangle
:type frect: QtCore.QRectF
:return: New rect with margins added
:rtype: QtCore.QRectF
"""
dx = frect.width() * 0.1
dy = frect.height() * 0.1
return frect.adjusted(-dx, -dy, dx, dy)
[docs] def expandToNodes(self):
"""
Expand the scene to fit all nodes
"""
frect = self._padRect(self.getFitRect())
curr_rect = self.view.sceneRect()
old_top, old_left, old_bot, old_right = curr_rect.getCoords()
fit_top, fit_left, fit_bot, fit_right = frect.getCoords()
# Top left is toward the origin so use min
new_top = min(old_top, fit_top)
new_left = min(old_left, fit_left)
# Bottom right is away from the origin so use max
new_bot = max(old_bot, fit_bot)
new_right = max(old_right, fit_right)
frect.setCoords(new_top, new_left, new_bot, new_right)
self.scene.setSceneRect(frect)
[docs] def setDefaultScale(self):
""" This zooms to fits all nodes and sets the default scale """
frect = self.getFitRect()
self.scene.setSceneRect(self._padRect(frect))
self.view.fitInView(frect, mode=Qt.KeepAspectRatio)
self.min_scale = self.view.transform().m11() * 0.8
self.max_scale = 10.0
self.view.scaleChanged.emit()
[docs] def scaleSmoother(self, val):
"""
Sets the scale between self.start_scale and self.finish_scale,
based on 'val', which goes from 0->1 as the animation proceeds.
"""
current_scale = self.view.transform().m11()
self.view.scale(1 / current_scale, 1 / current_scale)
new_scale = val * (self.finish_scale -
self.start_scale) + self.start_scale
self.view.scale(new_scale, new_scale)
[docs] def setRawScale(self, scale):
"""
Sets the scale using the actual multipliers QT does.
We use a QTimeLine to do this as an animation.
The "finish_scale" is always set here, whether the timeline is
already zooming or not, so that it'll always zoom as far as the
users last scroll desired.
"""
self.finish_scale = scale
val = (scale - self.min_scale) * 100. / (self.max_scale -
self.min_scale)
self.view.scaleChanged.emit()
if self.scale_timeline.state() == QTimeLine.Running:
return
self.start_scale = self.view.transform().m11()
self.scale_timeline = QTimeLine(300)
self.scale_timeline.setFrameRange(0, 20)
self.scale_timeline.valueChanged.connect(self.scaleSmoother)
self.scale_timeline.start()
[docs] def setScale(self, val):
""" Accepts any scale value in [0, 100] """
scale = self.min_scale + val / 100. * (self.max_scale - self.min_scale)
self.setRawScale(scale)
def _get_formatted_pen(pen_style, color):
"""
Return a new pen instance with a standard formatting applied to it.
:param pen_style: the desired pen style
:type pen_style: Qt.PenStyle
:param color: the desired pen color
:type color: QtGui.QColor
:return: a formatted pen object
:rtype: QtGui.QPen
"""
pen = QPen(pen_style)
pen.setJoinStyle(Qt.MiterJoin)
pen.setWidth(2 * PSF)
pen.setColor(color)
return pen