graph_neuron

main.py

import sys
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *


class Node(QGraphicsEllipseItem):
    def __init__(self, parent=None):
        QGraphicsEllipseItem.__init__(self, QRectF(-20, -20, 40, 40), parent)
        self.edges = []
        self.setBrush(Qt.white)

    def addEdge(self, edge):
        self.edges.append(edge)

    def itemChange(self, change, value):
        if change == QGraphicsItem.ItemSelectedChange:
            self.setBrush(Qt.green if value else Qt.darkGray)

        if change == QGraphicsItem.ItemPositionHasChanged:
            for edge in self.edges:
                edge.adjust()

        return QGraphicsItem.itemChange(self, change, value)


class Edge(QGraphicsLineItem):
    def __init__(self, source, dest, parent=None):
        QGraphicsLineItem.__init__(self, parent)
        self.source = source
        self.dest = dest
        self.source.addEdge(self)
        self.dest.addEdge(self)
        self.setPen(QPen(Qt.black, 1.75))
        self.adjust()

    def adjust(self):
        self.prepareGeometryChange()
        self.setLine(QLineF(self.dest.pos(), self.source.pos()))


class ViewClass(QGraphicsView):
    def __init__(self, layers, parent=None):
        QGraphicsView.__init__(self, parent)
        self.setScene(QGraphicsScene(QRectF(-100, -100, 600, 600)))
        self.setRenderHint(QPainter.Antialiasing)
        dst_layer = 100 
        dst_neuron = 50

        groups = []
        ycenters = []

        for layer in range(len(layers)-1, -1, -1):
        	group = QGraphicsItemGroup()
        	group.setZValue(10)
        	self.scene().addItem(group)
        	for neuron in range(0, layers[layer]):
        		node = Node(group)
        		self.scene().addItem(node)
        		node.setPos(dst_layer * layer, dst_neuron * neuron)
        		node.setPen(Qt.green if layer==0 else Qt.red if layer==(len(layers)-1) else Qt.blue)
        		group.addToGroup(node)
        	groups.append(group)
        	ycenters.append(group.boundingRect().center().y())

        ys = [max(ycenters) - center for center in ycenters]
        for group, y in zip(groups, ys):
        	[item.moveBy(0, y) for item in group.childItems()]

        for i in range(len(groups)-1):
        	last_group = groups[i]
        	next_group = groups[i+1]
        	for last_item in last_group.childItems():
        		for next_item in next_group.childItems():
        			edge = Edge(last_item, next_item)
        			self.scene().addItem(edge)




if __name__ == '__main__':
    app = QApplication(sys.argv)
    layers = [8, 3, 4, 2, 1]
    wd = ViewClass(layers)
    wd.show()
    sys.exit(app.exec_())