graph pyside6

graph_pyside6.py

from PySide6.QtWidgets import *
from PySide6.QtCore import *
from PySide6.QtGui import *

import math
import networkx as nx
import igraph as ig


class Node(QGraphicsItem):
    def __init__(self, name, parent=None):
        super().__init__(parent)
        self._name = name
        self.edges = []

        self.color = "red" if name.startswith("x") else "#2FB4AE"
        self.radius = 30 if name.startswith("x") else 30
        self._rect = QRect(0, 0, self.radius * 2, self.radius * 2)

        self.setFlag(self.ItemIsMovable)
        self.setFlag(self.ItemSendsGeometryChanges)
        self.setCacheMode(self.DeviceCoordinateCache)

    def boundingRect(self):
        return self._rect

    def paint(self, painter, option, widget=None):
        painter.setRenderHints(QPainter.Antialiasing)
        painter.setPen(
            QPen(
                QColor("#2FB4AE").darker(),
                2,
                Qt.SolidLine,
                Qt.RoundCap,
                Qt.RoundJoin,
            )
        )
        painter.setBrush(QBrush(QColor(self.color)))
        painter.drawEllipse(self.boundingRect())
        painter.setPen(QPen(QColor("white")))
        painter.drawText(self.boundingRect(), Qt.AlignCenter, self._name)

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

    def itemChange(self, change, value):

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

        return super().itemChange(change, value)


class Edge(QGraphicsItem):
    def __init__(self, source, dest, parent=None):
        super().__init__(parent)
        self.source = source
        self.dest = dest

        self.height = 2
        self.color = "#b7b7b8"
        self.arrowSize = 20

        self.source.add_edge(self)
        self.dest.add_edge(self)

        self.line = QLineF()

        self.setZValue(-1)

        self.adjust()

    def boundingRect(self):

        return (
            QRectF(self.line.p1(), self.line.p2())
            .normalized()
            .adjusted(
                -self.height - self.arrowSize,
                -self.height - self.arrowSize,
                self.height + self.arrowSize,
                self.height + self.arrowSize,
            )
        )

    def adjust(self):
        self.prepareGeometryChange()
        self.line = QLineF(
            self.source.pos() + self.source.boundingRect().center(),
            self.dest.pos() + self.dest.boundingRect().center(),
        )
        # self.update()

    def draw_arraow(self, painter, start, end):

        painter.setBrush(QBrush(self.color))

        line = QLineF(end, start)

        angle = math.atan2(-line.dy(), line.dx())
        arrowP1 = line.p1() + QPointF(
            math.sin(angle + math.pi / 3) * self.arrowSize,
            math.cos(angle + math.pi / 3) * self.arrowSize,
        )
        arrowP2 = line.p1() + QPointF(
            math.sin(angle + math.pi - math.pi / 3) * self.arrowSize,
            math.cos(angle + math.pi - math.pi / 3) * self.arrowSize,
        )

        center = (arrowP1 - arrowP2) / 2

        arrowHead = QPolygonF()
        arrowHead.clear()
        arrowHead.append(line.p1())
        arrowHead.append(arrowP1)
        arrowHead.append(arrowP2)
        painter.drawLine(line)
        painter.drawPolygon(arrowHead)

        painter.setPen(
            QPen(
                QColor("red"),
                10,
                Qt.SolidLine,
                Qt.RoundCap,
                Qt.RoundJoin,
            )
        )

    def paint(self, painter, option, widget=None):

        if self.source and self.dest:
            painter.setRenderHints(QPainter.Antialiasing)

            painter.setPen(
                QPen(
                    QColor(self.color),
                    self.height,
                    Qt.SolidLine,
                    Qt.RoundCap,
                    Qt.RoundJoin,
                )
            )
            painter.drawLine(self.line)

            self.draw_arraow(painter, self.line.p1(), self.get_target())

    def get_target(self):
        target = self.line.p1()
        center = self.line.p2()
        radius = self.dest.radius
        vector = target - center
        length = math.sqrt(vector.x() ** 2 + vector.y() ** 2)
        normal = vector / length
        target = QPointF(
            center.x() + (normal.x() * radius), center.y() + (normal.y() * radius)
        )
        return target


class PedView(QGraphicsView):
    def __init__(self, parent=None):
        super().__init__()

        scene = QGraphicsScene()

        self.setScene(scene)

        self.resize(600, 600)

    def read_pedfile(self, pedfile):
        self.graph = nx.DiGraph()
        edges = []
        self.index = []
        i = 0
        with open(pedfile, "r") as file:
            for line in file:
                rows = line.strip().split("\t")
                index = rows[1]
                father = rows[2]
                mother = rows[3]

                self.index.append(0)

                i += 1

                if father and father != "NA":
                    edges.append((father, index))

                if mother and mother != "NA":
                    edges.append((mother, index))

        self.graph.add_edges_from(edges)
        self.load()

    def load(self):

        # graph = nx.DiGraph()
        # graph.add_edges_from(
        #     [
        #         ("Grand Pere", "union1"),
        #         ("Grand Mere", "union1"),
        #         ("union1", "union1b"),
        #         ("union1b", "unionf1"),
        #         ("union1b", "unionf2"),
        #         ("unionf1", "Boby"),
        #         ("unionf2", "Elise"),
        #         ("Elise", "union2"),
        #         ("Charles", "union2"),
        #         ("union2", "union2b"),
        #         ("union2b", "unionf3"),
        #         ("unionf3", "Stephane"),
        #     ]
        # )

        # pos = nx.spring_layout(graph, iterations=5000, scale=100, k=50)

        self.graph = nx.DiGraph()
        self.graph.add_edges_from(
            [
                ("Pere", "xmate"),
                ("Mere", "xmate"),
                ("xmate", "xsib"),
                ("xsib", "Lucas"),
                ("xsib", "Nina"),
                ("Lucas", "xmate2"),
                ("Monica belluci", "xmate2"),
                ("xmate2", "xsib2"),
                ("xsib2", "Lenine"),
            ]
        )
        print(self.graph.nodes())

        h = ig.Graph.from_networkx(self.graph)

        index = [0, 0, 0, 1, 2, 2, 2, 2, 3, 4]

        layout = h.layout_sugiyama(layers=index, maxiter=1000)
        layout.scale(100)
        pos = layout.coords

        nodes = {}
        for i in h.vs():
            node = i["_nx_name"]
            item = Node(node)
            x, y = pos[i.index]
            item.setPos(float(x), float(y))
            self.scene().addItem(item)
            nodes[i.index] = item

        for a, b in h.get_edgelist():
            source = nodes[a]
            dest = nodes[b]

            self.scene().addItem(Edge(source, dest))

        # nodes = {}
        # for node in graph:
        #     item = Node(node)
        #     x, y = pos[node]
        #     item.setPos(float(x), float(y))
        #     self.scene().addItem(item)
        #     nodes[node] = item

        # for a, b in graph.edges:

        #     source = nodes[a]
        #     dest = nodes[b]

        #     self.scene().addItem(Edge(source, dest))


if __name__ == "__main__":
    import sys

    app = QApplication(sys.argv)

    view = PedView()
    view.read_pedfile("pedigree.tfam")
    view.show()

    app.exec()