Рисуем графы в PyQT

Время прочтения: 17 мин.

Некоторое время назад я решал задачу в области графовой аналитики, искал связи между клиентами. В качестве входных данных у меня были переводы между клиентами и основной задачей было: выяснить связи между клиентами и организациями.

Алгоритм работы такой: на входе я брал *.csv файлы из базы данных, которые сводил в pandas, и затем посредством группировок и фильтров выводил связи по каждому клиенту и организации, а после закидывал это все в networkx и получал на выходе граф, с которым уже работал совместно с аналитиком.

Аналитики, посмотрев на это все, сказали, а может сделаем что-нибудь быстрое для нас, чтобы мы сами там крутили все и смотрели. Конечно же, мой ответ был – а может что-нибудь из open-source? Возьмите Gephi, например. Прошла неделя, приходят обратно – вот это не так, это не то и вообще сложно. Я, честно, «отбивался», но как видите, не выдержал.

Итак, формируем задачу: нужно сделать простое приложение для построения графов, не потратив на него много времени, а также, чтобы запускалось и работало без бубнов. Язык – Python. Чтобы не «запариваться» с оберткой, а сосредоточиться на внутренних алгоритмах, в качестве интерфейса взял PYQT.

Саму разработку разделил на два этапа, а именно — сначала накидал макет интерфейса (макеты в Paint наше все), а потом уже добавил встроенные методы для расчета и отрисовки графа.

Интерфейс

Давайте остановимся на организации пространства в основном окне PYQT. Вообще, для создания интерфейсов PyQT имеет отдельное GUI, где можно просто нарисовать то, что вы хотите, и на выходе программа даст вам код. Мне было интересно разобраться во внутренней организации и попробовать написать с нуля. И я не придумал как туда засунуть сразу граф.

Сама программа состоит из основного класса (назовем его, например, GraphWindow), который наследует класс из библиотеки QtWidgets. QMainWindow и который мы в итоге будем отображать пользователю при запуске.

Желаемая схема будущего интерфейса выглядит вот так:

Вернемся к организации пространства, чтобы получить необходимое расположение блоков. Фактически PYQT не имеет такого инструмента как верстка, но тем не менее использует схожую схему расположения блоков как в HTML

Создадим файл main.py, ставим все библиотеки, пишем импорты.

from PyQt5.QtWidgets import QApplication, QWidget, QFileDialog, QListWidget, QMessageBox, QColorDialog
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtGui import QPixmap

import sys
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas

Напишем основной класс интерфейса, и впишем в метод конструктора базовые параметры окна:

class GraphWindow(QtWidgets.QMainWindow): - наследуем класс от виджета основного окна

    def __init__(self):
        super().__init__() – инициируем конструктор наследуемого класса
        self.resize(1280, 1080) – сразу выставим размеры окна
        self.move(250, 0) – перенесем немного в центр
        self.title = 'Graph' – название окна
	# Контейнер обертка для блоков
	self._main = QtWidgets.QWidget(self)

Разберемся из чего состоит окно. Документация нам говорит о том, что для организации блоков можно использовать два типа обертки, наследуемые от родительского класса QtWidgets QVBoxLayout и QHBoxLayout – V(vertical) и H(horizontal) соответственно, говорят нам о расположении блоков внутри обертки.

Для того, чтобы было понятно, давайте нарисуем блоки прямо на макете:

Последовательно собираем наш интерфейс из вертикальных и горизонтальных блоков.

# Основной контейнер для обоих частей
self.horizontal_container = QtWidgets.QHBoxLayout(self._main)

# Делим окно на два - левая и правая часть
self.vertical_container_left = QtWidgets.QVBoxLayout(self)
self.vertical_container_right = QtWidgets.QVBoxLayout(self)
self.horizontal_container.addLayout(self.vertical_container_left)
self.horizontal_container.addLayout(self.vertical_container_right)

Для того, чтобы у нас была возможность выбирать с какими полями работать, создадим контейнер для списков полей из исходного файла.

# Контейнер для списков выбора полей из датафрейма
self.horizontal_container_lists = QtWidgets.QHBoxLayout(self)

Контейнер для списков полей состоит из 3 оберток (исходные узлы, входящие узлы и связь между ними). Так как мы еще добавляем лейблы (подписи) для них, оборачиваем каждый список в контейнер. Поля оформляем в QlistWidget.

# Контейнер для выбора поля 1 узла
self.vertical_container_nodes_first = QtWidgets.QVBoxLayout(self)
# Лейбл для листа с узлами 1
self.nodes_list_first_label = QtWidgets.QLabel(self)
self.nodes_list_first_label.setFixedSize(100, 40)
# Лист для выбора поля для узла 1
self.nodes_list_first_widget = QtWidgets.QListWidget(self)
self.nodes_list_first_widget.setFixedSize(100, 150)
self.vertical_container_nodes_first.addWidget(self.nodes_list_first_label)
self.vertical_container_nodes_first.addWidget(self.nodes_list_first_widget)

# Контейнер для выбора поля 2 узла
self.vertical_container_nodes_second = QtWidgets.QVBoxLayout(self)
# Лейбл для листа с узлами 2
self.nodes_list_second_label = QtWidgets.QLabel(self)
self.nodes_list_second_label.setFixedSize(100, 40)
# Лист для выбора поля с узлами 2
self.nodes_list_second_widget = QListWidget(self)
self.nodes_list_second_widget.setFixedSize(100, 150)

self.vertical_container_nodes_second.addWidget(self.nodes_list_second_label)
self.vertical_container_nodes_second.addWidget(self.nodes_list_second_widget)

# Контейнер для выбора поля входящих ребер
self.vertical_container_edges_in = QtWidgets.QVBoxLayout(self)
# Лейбл для листа с ребрами
self.edges_in_list_label = QtWidgets.QLabel(self)
self.edges_in_list_label.setFixedSize(100, 35)

# Лист для выбора поля с ребрами
self.edges_in_list_widget = QListWidget(self)
self.edges_in_list_widget.setFixedSize(100, 150)

self.vertical_container_edges_in.addWidget(self.edges_in_list_label)
self.vertical_container_edges_in.addWidget(self.edges_in_list_widget)

В процессе работы поступила просьба — а можно ли узлы как-то дополнительно помечать, допустим, если клиент еще очень важный, соотвественно, добавить им спецпризнак. Не вопрос, добавляем в контейнер еще один список.

# Контейнер для выбора спец признак
self.vertical_container_special_features = QtWidgets.QVBoxLayout(self)
# Лейбл для листа со сцец празнаками
self.special_features_list_label = QtWidgets.QLabel(self)
self.special_features_list_label.setFixedSize(100, 35)

# Лист для выбора поля с ребрами
self.special_features_list_widget = QListWidget(self)
self.special_features_list_widget.setFixedSize(100, 150)

self.vertical_container_special_features.addWidget(self.special_features_list_label)
self.vertical_container_special_features.addWidget(self.special_features_list_widget)

После того как мы ввели все необходимые контейнеры для выбора полей, переносим их в горизонтальную обертку:

# Добавляем контейнеры для списков
self.horizontal_container_lists.addLayout(self.vertical_container_nodes_first)
self.horizontal_container_lists.addLayout(self.vertical_container_nodes_second)
self.horizontal_container_lists.addLayout(self.vertical_container_edges_in)
self.horizontal_container_lists.addLayout(self.vertical_container_special_features)

В процессе опять поступает просьба – хотим, чтобы узлы можно было раскрашивать разными цветами. Хорошо, что PyQT имеет встроенный класс для Color Picker, добавляем его наверх нашего макета.

self.horizontal_container_df_options = QtWidgets.QHBoxLayout(self)

# Выбор цвета для узла 1
self.horizontal_container_color_nodes_first = QtWidgets.QHBoxLayout(self)
self.color_picker_nodes_first_label = QtWidgets.QLabel(self)
self.color_picker_nodes_first_label.setFixedSize(65, 35)
self.color_picker_nodes_first = QtWidgets.QPushButton(self.nodes_first_color, self)
self.color_picker_nodes_first.setFixedSize(100, 30)

self.horizontal_container_color_nodes_first.addWidget(self.color_picker_nodes_first_label)
self.horizontal_container_color_nodes_first.addWidget(self.color_picker_nodes_first)

# Выбор цвета для узла 2
self.horizontal_container_color_nodes_second = QtWidgets.QHBoxLayout(self)
self.color_picker_nodes_second_label = QtWidgets.QLabel(self)
self.color_picker_nodes_second_label.setFixedSize(65, 35)
self.color_picker_nodes_second = QtWidgets.QPushButton(self.nodes_second_color, self)
self.color_picker_nodes_second.setFixedSize(100, 30)

self.horizontal_container_color_nodes_second.addWidget(self.color_picker_nodes_second_label)
self.horizontal_container_color_nodes_second.addWidget(self.color_picker_nodes_second)

self.horizontal_container_df_options.addLayout(self.horizontal_container_separator)
self.horizontal_container_df_options.addLayout(self.horizontal_container_color_nodes_first)

# Добавляем в контейнер сепаратора
self.vertical_container_left.addLayout(self.horizontal_container_df_options)

Для удобства работы сделаем кнопку загрузку файла.

# Кнопка загрузки файла
self.pushButton = QtWidgets.QPushButton(self)
self.pushButton.setObjectName("pushButton")
# self.pushButton.setFixedSize(150, 50)
self.pushButton.move(150, 10)
self.vertical_container_left.addWidget(self.pushButton)

После обработки нам еще понадобится кнопка, чтобы построить граф по имеющимся данным.

# Кнопка для построения графа
self.pushButton_3 = QtWidgets.QPushButton(self)
self.pushButton_3.setObjectName("pushButton_3")

self.vertical_container_left.addLayout(self.horizontal_container_lists)
self.vertical_container_left.addWidget(self.pushButton_3)

Так как попросили еще добавить фильтры, добавим контейнеры для них. После внесем их в следующий ряд.

# Контейнер для выбора узла
self.horizontal_container_lists_second = QtWidgets.QHBoxLayout(self)

# Контейнер для выбора узла для отсмотра
self.vertical_container_get_node = QtWidgets.QVBoxLayout(self)
# Контейнер для выбора типа узлов
self.vertical_container_edge_type = QtWidgets.QVBoxLayout(self)
# Контейнер для ввода суммы
self.vertical_container_sum_input = QtWidgets.QVBoxLayout(self)
# Контейнер для индикации найденных узлов
self.vertical_container_edges_found = QtWidgets.QVBoxLayout(self)

# Третий ряд контейнеров
self.horizontal_container_lists_second.addLayout(self.vertical_container_get_node)
self.horizontal_container_lists_second.addLayout(self.vertical_container_edge_type)
self.horizontal_container_lists_second.addLayout(self.vertical_container_sum_input)
self.horizontal_container_lists_second.addLayout(self.vertical_container_edges_found)

Теперь нам необходимо создать объекты под сами фильтры. Фактически это будут не только выбор узла, но и фильтры по суммам узлов, количеству и типу.

# Лист для выбора узла
self.get_node_list_widget = QListWidget(self)
self.get_node_list_widget.setFixedSize(150, 150)
# Лейбл для листа с сотрудниками
self.get_node_list_label = QtWidgets.QLabel(self)
self.get_node_list_label.setFixedSize(190, 50)
self.vertical_container_get_node.addWidget(self.get_node_list_label)
self.vertical_container_get_node.addWidget(self.get_node_list_widget)

# Лейбл для выбора типа ребра
self.edge_type_label = QtWidgets.QLabel(self)
self.edge_type_label.setFixedSize(150, 50)
# Лист для выбора типа ребра
self.edge_type_list = QtWidgets.QListWidget(self)
self.edge_type_list.setFixedSize(150, 150)
self.vertical_container_edge_type.addWidget(self.edge_type_label)
self.vertical_container_edge_type.addWidget(self.edge_type_list)

# Фильтр по количеству узлов
self.sum_limit_value_label = QtWidgets.QLabel(self)
self.sum_limit_value_label.setFixedSize(100, 50)
self.sum_limit_value = QtWidgets.QLineEdit(self)
self.sum_limit_value.setFixedSize(100, 20)
self.sum_limit_value.setText("0")
sum_validator = QtGui.QIntValidator()
sum_validator.setRange(0, 1000000)
self.sum_limit_value.setValidator(sum_validator)
self.vertical_container_sum_input.addWidget(self.sum_limit_value_label)
self.vertical_container_sum_input.addWidget(self.sum_limit_value)

# Количество уникальных связей
self.uniq_edges_found_label = QtWidgets.QLabel(self)
self.uniq_edges_found_label.setFixedSize(100, 50)
self.uniq_edges_found_text = QtWidgets.QLineEdit(self)
self.uniq_edges_found_text.setFixedSize(100, 20)
self.uniq_edges_found_text.setReadOnly(True)
self.vertical_container_sum_input.addWidget(self.uniq_edges_found_label)
self.vertical_container_sum_input.addWidget(self.uniq_edges_found_text)

# self.vertical_container_sum_input.addStretch()

# Фильтр суммы связей
self.edges_sum_filter_label = QtWidgets.QLabel(self)
self.edges_sum_filter_label.setFixedSize(100, 50)
self.edges_sum_filter_text = QtWidgets.QLineEdit(self)
self.edges_sum_filter_text.setFixedSize(100, 20)
self.edges_sum_filter_text.setText("0")
sum_validator = QtGui.QIntValidator()
sum_validator.setRange(0, 1000000)
self.sum_limit_value.setValidator(sum_validator)

self.vertical_container_edges_found.addWidget(self.edges_sum_filter_label)
self.vertical_container_edges_found.addWidget(self.edges_sum_filter_text)

# Сумма связей
self.edges_sum_label = QtWidgets.QLabel(self)
self.edges_sum_label.setFixedSize(100, 50)
self.edges_sum_text = QtWidgets.QLineEdit(self)
self.edges_sum_text.setFixedSize(100, 20)
self.edges_sum_text.setReadOnly(True)
self.vertical_container_edges_found.addWidget(self.edges_sum_label)
self.vertical_container_edges_found.addWidget(self.edges_sum_text)
# self.vertical_container_edges_found.addStretch()

И в очередной раз заносим все в контейнер

# ДОБАВЛЯЕМ ВСЕ В КОНТЕЙНЕР
self.vertical_container_left.addLayout(self.horizontal_container_lists_second)

Так как один из фильтров предполагает выбор одного из узлов (для более детального результата) добавим кнопку отрисовки графа по одному из узлов.

# Кнопка для построения графа по узлу
self.pushButton_2 = QtWidgets.QPushButton(self)
self.pushButton_2.setObjectName("pushButton_2")

self.vertical_container_left.addWidget(self.pushButton_2)
self.vertical_container_left.addStretch()

После добавления всех элементов, мы наконец-то можем инициировать основной виджет с всеми компонентами.

# Инициируем основной виджет
self._main.setFocus()
self.setCentralWidget(self._main)

И здесь важное отступление, я кое-что забыл. Вспоминаем, изначально мы хотели отрисовать граф с использованием networkx. Посмотрим в документации библиотеки.

Метод отрисовки выдает нам объект библиотеки matplotlib. Для того, чтобы связать PyQT с данным объектом, импортируем из бэкэнда matplot класс для работы с canvas.

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas

И добавляем его в правую часть приложения (то место, где будет отрисован граф).

# Добавляем блок для отрисовки графа
self.fig = plt.figure(figsize=(5, 4))
self.canvas = FigureCanvas(self.fig)
self.canvas.draw()
self.vertical_container_right.addWidget(self.canvas)

После всех действий запускаем метод, который связывает лейблы и навешивает обработчики на кнопки.

self.retranslateUi()

Посмотрим на сам метод.

def retranslateUi(self):
    """
    Функция навешивает обработчики на кнопки и списки
    Function sets texts and clicks on lists and buttons
    :return:
    """
    _translate = QtCore.QCoreApplication.translate
    self.nodes_list_first_label.setText(_translate("MainWindow", "Выберите поле,\nсодержащее \n узлы(1)"))
    self.nodes_list_second_label.setText(
        _translate("MainWindow", "Выберите поле,\nсодержащее \n узлы(2)"))
    self.edges_in_list_label.setText(
        _translate("MainWindow", "Выберите поле, \n для ребра"))
    self.special_features_list_label.setText(_translate("MainWindow", "Выберите столбец \n с спец признаками \n ("
                                                                      "необязательно)"))

    self.separator_label.setText(_translate("MainWindow", "Разделитель текста \n (по умолчанию ,)"))

    self.pushButton.setText(_translate("MainWindow", "Загрузить \n файл"))
    self.pushButton_2.setText(_translate("MainWindow", "Построить \n граф по узлу"))
    self.pushButton_3.setText(_translate("MainWindows", "Построить общий граф"))
    self.get_node_list_label.setText(
        _translate("MainWindow",
                   "Выберите узел, по которому \n нужно построить граф "
                   "\n (можно выбрать после выбора \n поля с узлами)"))
    self.edge_type_label.setText(
        _translate("MainWindow",
                   "Выберите тип \n связи"))
    self.uniq_edges_found_label.setText(
        _translate("MainWindow", "Количество \n уникальных \n связей"))
    self.sum_limit_value_label.setText(
        _translate("MainWindow", "Фильтр по \n количеству связей"))
    self.edges_sum_label.setText(
        _translate("MainWindow", "Cумма по связям"))

    self.edges_sum_filter_label.setText(
        _translate("MainWindow", "Фильтр по сумме \n связей")
    )

    self.edge_type_list.addItems(['Все', 'Исходящие', 'Входящие'])
    self.color_picker_nodes_first_label.setText(_translate("MainWindow", "Цвет узла 1"))
    self.color_picker_nodes_second_label.setText(_translate("MainWindow", "Цвет узла 2"))
    self.pushButton.clicked.connect(self.pushButton_handler)
    
    self.pushButton_2.clicked.connect(self.draw_single_graph)
    self.pushButton_3.clicked.connect(self.draw_summary_graph)
    self.pushButton_2.setEnabled(False)
    self.pushButton_3.setEnabled(False)
    self.color_picker_nodes_first.clicked.connect(self.open_color_dialog_first)
    self.color_picker_nodes_second.clicked.connect(self.open_color_dialog_second)
    self.nodes_list_first_widget.itemClicked.connect(self.get_uniq_first_node)
    self.nodes_list_second_widget.itemClicked.connect(self.get_second_node_column)
    self.edges_in_list_widget.itemClicked.connect(self.get_edge_in_column)
    self.special_features_list_widget.itemClicked.connect(self.get_special_node_column)
    self.get_node_list_widget.itemClicked.connect(self.get_special_node)
    self.edge_type_list.itemClicked.connect(self.set_edge_type)
    self.setWindowTitle(_translate("Graph-program", "Graph-program"))

В нижней части мы навешиваем обработчики на кнопки/листы/фильтры. Каждую из функций рассмотрим ниже.

Запустим наш получившийся интерфейс и посмотрим, что у нас вышло.

Здесь нужно учесть, что, если функций в контексте не существует, приложение не запустится.

if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    window = GraphWindow()
    window.show()
    app.exec_()

Зовем заказчика, показываем. В целом интерфейс устраивает, приступаем к следующей части – написанию всех необходимых методов.

Сформируем список для разработки:

1. Загрузка из файла
2. Выбор узлов
3. Выбор ребер
4. Расчет узлов
5. Расчет ребер и их весов
6. Отрисовка графа
7. Фильтры

Начнем по порядку:

1.Загрузка из файла

В качестве решения я выбрал библиотеку pandas – во многом в ней уже есть функционал для работы с данными, поэтому остается только подвязать все зависимости.

Для начала напишем класс для работы с входным файлом. На вход мы получаем ссылку на файл и преобразуем его в pandas dataframe. В дальнейшем мы можем легко в рамках приложения обращаться к нему, изменять его. Помимо этого в метод init pyqt введем параметр self.pandas_df

class DataFile:
    def __init__(self,
                 pandas_df: (pd.DataFrame, str),
                 nrows=None,
                 encoding=None,
                 sep=','
                 ):

        self.pandas_df = pandas_df

        if self.pandas_df is not None:
            if type(pandas_df) == str:
                if pandas_df.split('.')[-1] == 'csv':
                    self.pandas_df = pd.read_csv(pandas_df, sep=sep, encoding=encoding, nrows=nrows)
                elif pandas_df.split('.')[-1] in ['xlsx', 'xls']:
                    self.pandas_df = pd.read_excel(pandas_df, nrows=nrows)
                elif pandas_df.split('.')[-1] == 'txt':
                    self.pandas_df = pd.read_table(pandas_df, sep=sep, encoding=encoding, nrows=nrows)
                else:
                    msg = QMessageBox()
                    msg.setText(f'Only "csv", "xls(x)" and "txt" file formats are supported, '
                                f'but given file path ends with "{pandas_df.split(".")[-1]}"')
                    msg.exec_()

                    raise TypeError(f'Only "csv", "xls(x)" and "txt" file formats are supported, '
                                    f'but given file path ends with "{pandas_df.split(".")[-1]}"')
        else:
            msg = QMessageBox()
            msg.setText(f'pd.DataFrame or str types are expected, but got: {type(pandas_df)}')
            msg.exec_()
            raise TypeError(f'pd.DataFrame or str types are expected, but got: {type(pandas_df)}')

2. Для выбора столбов из датафрейма, напишем еще один метод. Метод будет сохранять имена столбцов в выделенные параметры при инициации окна. Делаем отдельные методы для исходящего и входящего узла.

Важное уточнение. В методе get_uniq_first_node мы помимо выбора столбца, сделаем выбор уникальных значений, для того, чтобы у нас была возможность отрисовать тот или иной узел на отдельном графе. Если что-то не получится, выбрасываем исключение и уведомление пользователю.

def get_uniq_first_node(self, item):
    """
    Функция берет уникальные значения из выбранного столбца с первыми узлами и записывает в новый список
    Function gets uniq values from column name with nodes and writes it to new list (QLISTWIDGET)
    :param item: column name from qlistwidget
    :return:
    """
    self.get_node_list_widget.clear()
    try:
        self.get_node_list_widget.addItems(self.df[item.text()].unique().tolist())
        self.nodes_list_first_column = item.text()
    except TypeError as ex:
        # Если в выбранном столбце дерьмовые данные, бросаем exception
        msg = QMessageBox()
        msg.setText("Неверный тип данных в столбце, необходима строка или число. Выберите другой столбец")
        msg.exec_()
def get_second_node_column(self, item):
    """
    Set second node column from df
    :param item: chosen item from qlistwidget
    :return:
    """

    self.nodes_list_second_column = item.text()

3. Метод выбора столбца с ребром (связью) сделаем по аналогии с предыдущим методом выбора узла.

Так как у нас была привязка к денежному эквиваленту, сразу делаем приведение к формату целого числа, и, если это нам не удается – выбрасываем пользователю сообщение.

def get_edge_in_column(self, item):
    """
    Set edge column from df
    :param item: chosen item from qlistwidget
    :return:
    """

    self.edges_in_list_column = item.text()
    try:
        self.df[self.edges_in_list_column] = self.df[self.edges_in_list_column].astype(int)
    except Exception as ex:
        # Если в выбранном столбце некачественные данные, бросаем exception
        msg = QMessageBox()
        msg.setText("Неверный формат данных узлов" + str(ex))
        msg.exec_()
    self.pushButton_3.setEnabled(True)

4. Для выбора спецпризнака добавляем точно такой же метод.

def get_special_node_column(self, item):
    """
    Get special node to visualize
    :param item: clicked item from qlistwidget
    :return:
    """
    self.special_list_column = item.text()

Теперь переходим к расчету самого графа и выделения связей. Подготовим исходные данные для анализа:

Сначала очищаем входные данные от пустых значений, если он есть.

Затем проверяем через фильтр (который стоит по умолчанию) на тип связи, который нужен пользователю (входящая, исходящая, все). Для этого оставляем только те связи, которые присутствуют только в одном столбце. Для примера, если есть связь x -> y и y -> x, убирая все х из второго столбца, у нас останутся только исходящие связи. Используем методы pandas для фильтра.

def prepare_graph_df(self, data_frame, is_edge_filtered):
    """
    # Формируем датафрейм с необходимыми стобцами и дф, в которых ребра и связи не пустые
    Параметр filtered nodes идет из узлов, если есть ограничение по количеству связей, убираем из дфа ненужные нам.
    """
    try:
        graph_df_full = data_frame.loc[
            (data_frame[self.nodes_list_first_column].notnull()) &
            (data_frame[self.nodes_list_second_column].notnull())]

        # Выбираем связи:
        if is_edge_filtered:
            if self.edge_type == 'Исходящие':
                # Убираем из вторых узлов первые - значит только исходящие связи
                graph_df_full = graph_df_full.loc[
                    graph_df_full[self.nodes_list_first_column] == self.nodes_first_clicked
                    ]
            elif self.edge_type == 'Входящие':
                graph_df_full = graph_df_full.loc[
                    graph_df_full[self.nodes_list_second_column] == self.nodes_first_clicked
                    ]
        return graph_df_full
    except Exception as ex:
        # Если в выбранном столбце дерьмовые данные, бросаем exception
        msg = QMessageBox()
        msg.setText("Вы не выбрали один из столбцов, " + str(ex))
        msg.exec_()

После того, как мы очистили данные и отфильтровали их по типу – нам нужно выбрать все узлы, которые будут отрисованы.

Сначала создаем пустые массивы для узлов и их цветов. Networkx принимает на вход два таких массива и сопоставляет их для отрисовки. После чего из входного датафрейма формируем два новых – один с первыми узлами, а второй со вторыми и присваиваем им в новый столбец «COLOR» цвета, которые выбрал пользователь через colorpicker.

def open_color_dialog_first(self):
    """
    Function shows open-color dialog and sets value for first nodes
    :return:
    """
    color = QColorDialog.getColor()
    self.nodes_first_color = color.name()
    self.color_picker_nodes_first.setText(self.nodes_first_color)

После того как мы присвоили цвета, соединяем датафреймы и убираем все дубликаты. Таким образом у нас получатся все возможные узлы графа, и каждое со своим цветом. Затем мы еще раз фильтруем узлы по входному массиву из узлов, которые не прошли по количеству связей. Сам расчет количества делается дальше.

def get_all_nodes(self, graph_df_full, filtered_nodes):
    """
    Method takes 3 df with nodes and returns arrays with values and colors for it
    :param filtered_nodes: If any nodes are filtered due to edges count we get array of nodes
    :param graph_df_full: graph, that has both in and out operations
    :return: Array
    """
    nodes_array = []
    nodes_color = []
    first_nodes_df = pd.DataFrame(columns=['NODES', 'COLOR'])
    second_nodes_df = pd.DataFrame(columns=['NODES', 'COLOR'])
    first_nodes_df['NODES'] = graph_df_full[self.nodes_list_first_column]
    first_nodes_df = first_nodes_df.assign(COLOR=self.nodes_first_color)
    first_nodes_df = first_nodes_df.drop_duplicates(subset=['NODES'])

    # Если выбрали колонку
    # где есть спецпризнак у первого спецзла, присваиваем новые цвета у него.
    if len(self.special_list_column) > 0:
        first_nodes_df = self.set_special_color(first_nodes_df)
    second_nodes_df['NODES'] = graph_df_full[self.nodes_list_second_column]
    second_nodes_df = second_nodes_df.assign(COLOR=self.nodes_second_color)
    second_nodes_df = second_nodes_df.drop_duplicates(subset=['NODES'])

    summary_nodes = pd.concat([first_nodes_df, second_nodes_df])

    summary_nodes = summary_nodes.drop_duplicates(subset=['NODES'])
    if len(filtered_nodes) > 0:
        summary_nodes = summary_nodes.loc[summary_nodes['NODES'].isin(filtered_nodes)]

    # Добавляем узлы
    nodes_array.extend(summary_nodes['NODES'].tolist())
    nodes_color.extend(summary_nodes['COLOR'].tolist())
    return nodes_array, nodes_color

Рассчитаем ребра графа. Сам метод получился довольно громоздкий, но я перефразирую одну цитату сами знаете кого: «Если костыль неизбежен, костылить надо первым».

Первым этапом мы группируем узлы (каждому по отдельности), чтобы получить все возможные вариации ребер в рамках датафрейма, делаем агрегацию по сумме и количеству, а после сбрасываем индекс для вывода таблицы.

Затем мы проверяем, вводил ли пользователь что-либо в поле фильтра по сумме, и, если да, то пробуем преобразовать в целое число.

def get_all_edges(self, graph_df_full):
    """
    Function counts edges and weights for graph
    :return: prepared edges array
    """
    filtered_nodes = []
    sum_limit = 0
    edges_sum_filter = 0
    # Убираем дубликаты и считаем веса / суммы / количество
    edges_full = graph_df_full.groupby(by=[self.nodes_list_first_column,
                                           self.nodes_list_second_column]).agg(
        VALUE=(self.edges_in_list_column, 'sum'),
        MAX_COUNT=(self.edges_in_list_column, 'count'))
    edges_full = edges_full.reset_index()
    # Фильтр по сумме
    try:
        edges_sum_filter = int(self.edges_sum_filter_text.text())
    except Exception as ex:
        edges_sum_filter = 0

Если все-таки в фильтре что-то есть – фильтруем датафрейм по сумме связи, и после добавляем все полученные узлы в массив с отфильтрованными узлами.

У нас еще есть фильтр по количеству связей узлов. В целом принцип фильтрации ничем не отличается от предыдущего.

if edges_sum_filter > 0:
    edges_full = edges_full.loc[edges_full['VALUE'] > int(self.edges_sum_filter_text.text())]
    filtered_nodes = edges_full[self.nodes_list_first_column].append(
        edges_full[self.nodes_list_second_column]).unique()

# Фильтруем по количеству связей
try:
    sum_limit = int(self.sum_limit_value.text())
except Exception as ex:
    sum_limit = 0

if sum_limit > 0:
    edges_count = pd.concat([edges_full.groupby(by=[self.nodes_list_first_column]).agg(
        MAX_COUNT_EDGES_IN=(self.nodes_list_first_column, 'count')).reset_index()
                            .rename(columns={self.nodes_list_first_column: 'node'}),

                             edges_full.groupby(by=[self.nodes_list_second_column]).agg(
                                 MAX_COUNT_EDGES_IN=(self.nodes_list_second_column, 'count')).reset_index()
                            .rename(columns={self.nodes_list_second_column: 'node'})])
    filtered_nodes = edges_count.groupby(by=['node']).sum().reset_index()
    filtered_nodes = filtered_nodes.loc[filtered_nodes['MAX_COUNT_EDGES_IN'] > int(self.sum_limit_value.text())]
    filtered_nodes = filtered_nodes['node'].unique()

    edges_full = edges_full.loc[edges_full[self.nodes_list_first_column].isin(filtered_nodes) |
                                edges_full[self.nodes_list_second_column].isin(filtered_nodes)]
    filtered_nodes = edges_full[self.nodes_list_first_column].append(
        edges_full[self.nodes_list_second_column]).unique()

После фильтров, все что нам остается – посчитать веса всех связей для визуального отображения на графе. Сам метод возвращает все ребра, а также массив отфильтрованных узлов, по которому мы фильтруем исходные узлы для отрисовки.

# Cчитаем веса
edges_full['WEIGHT'] = edges_full['VALUE'].apply(
    lambda x: (x / edges_full['VALUE'].sum()) *
              self.count_weight_coeff(len(edges_full[self.nodes_list_first_column])))
# Количество уникальных связей
self.uniq_edges_found_text.setText(str(len(edges_full[self.nodes_list_first_column])))
self.edges_sum_text.setText(str(edges_full['VALUE'].sum()))

return edges_full[[self.nodes_list_first_column,
                   self.nodes_list_second_column, 'WEIGHT']].to_numpy(), filtered_nodes

Здесь у меня появилась проблема, при отображении узлов с малым количеством весов, некоторые линии получались слишком тонкие, а некоторые слишком толстые. Для решения проблемы я ввел метод для коррекции весов на визуальной части графа.

def count_weight_coeff(self, nodes_list_length):
    """
    Count weights according to count, to get good visualization
    :return:
    """
    if nodes_list_length > 0 & nodes_list_length <= 50:
        return 5
    elif nodes_list_length > 50 & nodes_list_length <= 500:
        return 100
    elif nodes_list_length > 500:
        return 1000

Переходим к отрисовке графа. Фактически у нас уже создан Canvas контейнер, где мы будем отрисовать объект matplot. Все, что нам остается – создать объект графа встроенными методами networkx, наполнить узлами, ребрами и после отрисовать. Перед выводом не забываем очищать рисунок, чтобы все отображалось корректно.

Если у нас что-то пошло не так, просто выдаем пустой рисунок.

def draw_summary_graph(self):
    self.fig = plt.cla()
    graph = nx.DiGraph()
    # Готовим данные для графа (узлы)
    graph_df_full = self.prepare_graph_df(self.df, False)
    edges, filtered_nodes = self.get_all_edges(graph_df_full)
    for edge in edges:
        graph.add_edge(edge[0], edge[1], colour=20, weight=edge[2])

    edges = graph.edges()
    weights = [graph[u][v]['weight'] for u, v in edges]
    # Получаем узлы и цвета для них
    nodes_list, nodes_color = self.get_all_nodes(graph_df_full, filtered_nodes)
    # Добавляем ребра и веса
    try:
        nx.draw(graph, nodelist=nodes_list, node_color=nodes_color, width=weights, with_labels=True)
    except nx.NetworkXError:
        self.fig = plt.cla()
    self.fig = plt.figure(1, figsize=(3, 3))
    self.canvas.draw()

Запустим нашу программу, посмотрим результат.

Заказчик доволен. Все что осталось — упаковать его в exe файлик.

Для этого воспользуемся pyinstaller. Он упаковывает все в единый файл вместе с окружением.

C:\Users\Iov - RV\PycharmProjects\graph - pyqt - new > pyinstaller
main.py - -one - file

Запустим полученный exe для тестов.

Видим, что при запуске у нас разворачивается окружение и запускается проект. Здесь важное уточнение: у меня изначально не собирался проект, проблема была в версии matplotlib. В моем случае спасло положение — понижение версии библиотеки.

В целом результатом я доволен. Аналитики довольны, для их задач этого инструмента хватает.

При необходимости всегда можно доработать и внести правки в инструмент. Для примера могу сказать, что в дальнейшем мне удалось соединить внутренний механизм с нейронными сетями, чтобы обрабатывать и классифицировать данные.