elasticnodes.py¶
import sys
import weakref
import math
from PySide6 import QtCore, QtGui, QtWidgets
def random(boundary):
return QtCore.QRandomGenerator.global_().bounded(boundary)
class Edge(QtWidgets.QGraphicsItem):
Pi = math.pi
TwoPi = 2.0 * Pi
Type = QtWidgets.QGraphicsItem.UserType + 2
def __init__(self, sourceNode, destNode):
QtWidgets.QGraphicsItem.__init__(self)
self.arrowSize = 10.0
self.sourcePoint = QtCore.QPointF()
self.destPoint = QtCore.QPointF()
self.setAcceptedMouseButtons(QtCore.Qt.NoButton)
self.source = weakref.ref(sourceNode)
self.dest = weakref.ref(destNode)
self.source().addEdge(self)
self.dest().addEdge(self)
self.adjust()
def type(self):
return Edge.Type
def sourceNode(self):
return self.source()
def setSourceNode(self, node):
self.source = weakref.ref(node)
self.adjust()
def destNode(self):
return self.dest()
def setDestNode(self, node):
self.dest = weakref.ref(node)
self.adjust()
def adjust(self):
if not self.source() or not self.dest():
return
line = QtCore.QLineF(self.mapFromItem(self.source(), 0, 0), self.mapFromItem(self.dest(), 0, 0))
length = line.length()
if length == 0.0:
return
edgeOffset = QtCore.QPointF((line.dx() * 10) / length, (line.dy() * 10) / length)
self.prepareGeometryChange()
self.sourcePoint = line.p1() + edgeOffset
self.destPoint = line.p2() - edgeOffset
def boundingRect(self):
if not self.source() or not self.dest():
return QtCore.QRectF()
penWidth = 1
extra = (penWidth + self.arrowSize) / 2.0
return QtCore.QRectF(self.sourcePoint,
QtCore.QSizeF(self.destPoint.x() - self.sourcePoint.x(),
self.destPoint.y() - self.sourcePoint.y())).normalized().adjusted(-extra, -extra, extra, extra)
def paint(self, painter, option, widget):
if not self.source() or not self.dest():
return
# Draw the line itself.
line = QtCore.QLineF(self.sourcePoint, self.destPoint)
if line.length() == 0.0:
return
painter.setPen(QtGui.QPen(QtCore.Qt.black, 1, QtCore.Qt.SolidLine, QtCore.Qt.RoundCap, QtCore.Qt.RoundJoin))
painter.drawLine(line)
# Draw the arrows if there's enough room.
angle = math.acos(line.dx() / line.length())
if line.dy() >= 0:
angle = Edge.TwoPi - angle
sourceArrowP1 = self.sourcePoint + QtCore.QPointF(math.sin(angle + Edge.Pi / 3) * self.arrowSize,
math.cos(angle + Edge.Pi / 3) * self.arrowSize)
sourceArrowP2 = self.sourcePoint + QtCore.QPointF(math.sin(angle + Edge.Pi - Edge.Pi / 3) * self.arrowSize,
math.cos(angle + Edge.Pi - Edge.Pi / 3) * self.arrowSize)
destArrowP1 = self.destPoint + QtCore.QPointF(math.sin(angle - Edge.Pi / 3) * self.arrowSize,
math.cos(angle - Edge.Pi / 3) * self.arrowSize)
destArrowP2 = self.destPoint + QtCore.QPointF(math.sin(angle - Edge.Pi + Edge.Pi / 3) * self.arrowSize,
math.cos(angle - Edge.Pi + Edge.Pi / 3) * self.arrowSize)
painter.setBrush(QtCore.Qt.black)
painter.drawPolygon(QtGui.QPolygonF([line.p1(), sourceArrowP1, sourceArrowP2]))
painter.drawPolygon(QtGui.QPolygonF([line.p2(), destArrowP1, destArrowP2]))
class Node(QtWidgets.QGraphicsItem):
Type = QtWidgets.QGraphicsItem.UserType + 1
def __init__(self, graphWidget):
QtWidgets.QGraphicsItem.__init__(self)
self.graph = weakref.ref(graphWidget)
self.edgeList = []
self.newPos = QtCore.QPointF()
self.setFlag(QtWidgets.QGraphicsItem.ItemIsMovable)
self.setFlag(QtWidgets.QGraphicsItem.ItemSendsGeometryChanges)
self.setCacheMode(self.DeviceCoordinateCache)
self.setZValue(-1)
def type(self):
return Node.Type
def addEdge(self, edge):
self.edgeList.append(weakref.ref(edge))
edge.adjust()
def edges(self):
return self.edgeList
def calculateForces(self):
if not self.scene() or self.scene().mouseGrabberItem() is self:
self.newPos = self.pos()
return
# Sum up all forces pushing this item away.
xvel = 0.0
yvel = 0.0
for item in self.scene().items():
if not isinstance(item, Node):
continue
line = QtCore.QLineF(self.mapFromItem(item, 0, 0), QtCore.QPointF(0, 0))
dx = line.dx()
dy = line.dy()
l = 2.0 * (dx * dx + dy * dy)
if l > 0:
xvel += (dx * 150.0) / l
yvel += (dy * 150.0) / l
# Now subtract all forces pulling items together.
weight = (len(self.edgeList) + 1) * 10.0
for edge in self.edgeList:
if edge().sourceNode() is self:
pos = self.mapFromItem(edge().destNode(), 0, 0)
else:
pos = self.mapFromItem(edge().sourceNode(), 0, 0)
xvel += pos.x() / weight
yvel += pos.y() / weight
if QtCore.qAbs(xvel) < 0.1 and QtCore.qAbs(yvel) < 0.1:
xvel = yvel = 0.0
sceneRect = self.scene().sceneRect()
self.newPos = self.pos() + QtCore.QPointF(xvel, yvel)
self.newPos.setX(min(max(self.newPos.x(), sceneRect.left() + 10), sceneRect.right() - 10))
self.newPos.setY(min(max(self.newPos.y(), sceneRect.top() + 10), sceneRect.bottom() - 10))
def advance(self):
if self.newPos == self.pos():
return False
self.setPos(self.newPos)
return True
def boundingRect(self):
adjust = 2.0
return QtCore.QRectF(-10 - adjust, -10 - adjust,
23 + adjust, 23 + adjust)
def shape(self):
path = QtGui.QPainterPath()
path.addEllipse(-10, -10, 20, 20)
return path
def paint(self, painter, option, widget):
painter.setPen(QtCore.Qt.NoPen)
painter.setBrush(QtCore.Qt.darkGray)
painter.drawEllipse(-7, -7, 20, 20)
gradient = QtGui.QRadialGradient(-3, -3, 10)
if option.state & QtWidgets.QStyle.State_Sunken:
gradient.setCenter(3, 3)
gradient.setFocalPoint(3, 3)
gradient.setColorAt(1, QtGui.QColor(QtCore.Qt.yellow).lighter(120))
gradient.setColorAt(0, QtGui.QColor(QtCore.Qt.darkYellow).lighter(120))
else:
gradient.setColorAt(0, QtCore.Qt.yellow)
gradient.setColorAt(1, QtCore.Qt.darkYellow)
painter.setBrush(QtGui.QBrush(gradient))
painter.setPen(QtGui.QPen(QtCore.Qt.black, 0))
painter.drawEllipse(-10, -10, 20, 20)
def itemChange(self, change, value):
if change == QtWidgets.QGraphicsItem.ItemPositionChange:
for edge in self.edgeList:
edge().adjust()
self.graph().itemMoved()
return QtWidgets.QGraphicsItem.itemChange(self, change, value)
def mousePressEvent(self, event):
self.update()
QtWidgets.QGraphicsItem.mousePressEvent(self, event)
def mouseReleaseEvent(self, event):
self.update()
QtWidgets.QGraphicsItem.mouseReleaseEvent(self, event)
class GraphWidget(QtWidgets.QGraphicsView):
def __init__(self):
QtWidgets.QGraphicsView.__init__(self)
self.timerId = 0
scene = QtWidgets.QGraphicsScene(self)
scene.setItemIndexMethod(QtWidgets.QGraphicsScene.NoIndex)
scene.setSceneRect(-200, -200, 400, 400)
self.setScene(scene)
self.setCacheMode(QtWidgets.QGraphicsView.CacheBackground)
self.setRenderHint(QtGui.QPainter.Antialiasing)
self.setTransformationAnchor(QtWidgets.QGraphicsView.AnchorUnderMouse)
self.setResizeAnchor(QtWidgets.QGraphicsView.AnchorViewCenter)
node1 = Node(self)
node2 = Node(self)
node3 = Node(self)
node4 = Node(self)
self.centerNode = Node(self)
node6 = Node(self)
node7 = Node(self)
node8 = Node(self)
node9 = Node(self)
scene.addItem(node1)
scene.addItem(node2)
scene.addItem(node3)
scene.addItem(node4)
scene.addItem(self.centerNode)
scene.addItem(node6)
scene.addItem(node7)
scene.addItem(node8)
scene.addItem(node9)
scene.addItem(Edge(node1, node2))
scene.addItem(Edge(node2, node3))
scene.addItem(Edge(node2, self.centerNode))
scene.addItem(Edge(node3, node6))
scene.addItem(Edge(node4, node1))
scene.addItem(Edge(node4, self.centerNode))
scene.addItem(Edge(self.centerNode, node6))
scene.addItem(Edge(self.centerNode, node8))
scene.addItem(Edge(node6, node9))
scene.addItem(Edge(node7, node4))
scene.addItem(Edge(node8, node7))
scene.addItem(Edge(node9, node8))
node1.setPos(-50, -50)
node2.setPos(0, -50)
node3.setPos(50, -50)
node4.setPos(-50, 0)
self.centerNode.setPos(0, 0)
node6.setPos(50, 0)
node7.setPos(-50, 50)
node8.setPos(0, 50)
node9.setPos(50, 50)
self.scale(0.8, 0.8)
self.setMinimumSize(400, 400)
self.setWindowTitle(self.tr("Elastic Nodes"))
def itemMoved(self):
if not self.timerId:
self.timerId = self.startTimer(1000 / 25)
def keyPressEvent(self, event):
key = event.key()
if key == QtCore.Qt.Key_Up:
self.centerNode.moveBy(0, -20)
elif key == QtCore.Qt.Key_Down:
self.centerNode.moveBy(0, 20)
elif key == QtCore.Qt.Key_Left:
self.centerNode.moveBy(-20, 0)
elif key == QtCore.Qt.Key_Right:
self.centerNode.moveBy(20, 0)
elif key == QtCore.Qt.Key_Plus:
self.scaleView(1.2)
elif key == QtCore.Qt.Key_Minus:
self.scaleView(1 / 1.2)
elif key == QtCore.Qt.Key_Space or key == QtCore.Qt.Key_Enter:
for item in self.scene().items():
if isinstance(item, Node):
item.setPos(-150 + random(300), -150 + random(300))
else:
QtWidgets.QGraphicsView.keyPressEvent(self, event)
def timerEvent(self, event):
nodes = [item for item in self.scene().items() if isinstance(item, Node)]
for node in nodes:
node.calculateForces()
itemsMoved = False
for node in nodes:
if node.advance():
itemsMoved = True
if not itemsMoved:
self.killTimer(self.timerId)
self.timerId = 0
def wheelEvent(self, event):
self.scaleView(math.pow(2.0, -event.delta() / 240.0))
def drawBackground(self, painter, rect):
# Shadow.
sceneRect = self.sceneRect()
rightShadow = QtCore.QRectF(sceneRect.right(), sceneRect.top() + 5, 5, sceneRect.height())
bottomShadow = QtCore.QRectF(sceneRect.left() + 5, sceneRect.bottom(), sceneRect.width(), 5)
if rightShadow.intersects(rect) or rightShadow.contains(rect):
painter.fillRect(rightShadow, QtCore.Qt.darkGray)
if bottomShadow.intersects(rect) or bottomShadow.contains(rect):
painter.fillRect(bottomShadow, QtCore.Qt.darkGray)
# Fill.
gradient = QtGui.QLinearGradient(sceneRect.topLeft(), sceneRect.bottomRight())
gradient.setColorAt(0, QtCore.Qt.white)
gradient.setColorAt(1, QtCore.Qt.lightGray)
painter.fillRect(rect.intersected(sceneRect), QtGui.QBrush(gradient))
painter.setBrush(QtCore.Qt.NoBrush)
painter.drawRect(sceneRect)
# Text.
textRect = QtCore.QRectF(sceneRect.left() + 4, sceneRect.top() + 4,
sceneRect.width() - 4, sceneRect.height() - 4)
message = self.tr("Click and drag the nodes around, and zoom with the "
"mouse wheel or the '+' and '-' keys")
font = painter.font()
font.setBold(True)
font.setPointSize(14)
painter.setFont(font)
painter.setPen(QtCore.Qt.lightGray)
painter.drawText(textRect.translated(2, 2), message)
painter.setPen(QtCore.Qt.black)
painter.drawText(textRect, message)
def scaleView(self, scaleFactor):
factor = self.matrix().scale(scaleFactor, scaleFactor).mapRect(QtCore.QRectF(0, 0, 1, 1)).width()
if factor < 0.07 or factor > 100:
return
self.scale(scaleFactor, scaleFactor)
if __name__ == "__main__":
app = QtWidgets.QApplication(sys.argv)
widget = GraphWidget()
widget.show()
sys.exit(app.exec_())
© 2021 The Qt Company Ltd. Documentation contributions included herein are the copyrights of their respective owners. The documentation provided herein is licensed under the terms of the GNU Free Documentation License version 1.3 as published by the Free Software Foundation. Qt and respective logos are trademarks of The Qt Company Ltd. in Finland and/or other countries worldwide. All other trademarks are property of their respective owners.