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

Порой Data Scientist’ам приходится иметь дело с графами. Чаще всего, это дело не такое уж и сложное, но бывают разногласия, которые начинаются при представлении результатов заказчику данного графа, ведь у каждого своё представлении о прекрасном. Особенно, когда дело касается расположения узлов. Сегодня мы снова сколлаборируем библиотеки NetworkX и Plotly (как уже было ранее ), а также научимся отрисовывать 3D графы для более комфортного взаимодействия с заказчиком, который сможет сам покрутить полученные результаты. Пожалуй, начнем.

Импортируем нужные нам библиотеки:

import networkx as nx 
import plotly.graph_objects as go

Не будем мудрить и за датасетом обратимся к официальной документации NetworkX, там находим данные, связанные с американским футболом, под названием «Football», благо документации у данной библиотеки на 700+ страниц. Там находим небольшое описание данных и код, который позволяет получить и отстроить граф в 2D. Для вашего удобства размещаю эту часть ниже:

import urllib.request
import io
import zipfile
import matplotlib.pyplot as plt
url = "http://www-personal.umich.edu/~mejn/netdata/football.zip"
sock = urllib.request.urlopen(url)  
s = io.BytesIO(sock.read()) 
sock.close()
zf = zipfile.ZipFile(s)  
txt = zf.read("football.txt").decode() 
gml = zf.read("football.gml").decode()  
gml = gml.split("\n")[1:]

Создадим граф с помощью полученных данных:

G = nx.parse_gml(gml)  

Давайте посмотрим на данный граф в 2D, выставим цвет узла, его размер и толщину линии связи:

pos = nx.spring_layout(G, seed=1) 
options = {
    "node_color": "k",
    "node_size": 10,
    "width": .2,}
nx.draw(G, pos, **options)
plt.show();

Для создания графа в с тремя координатами установим размерность равную трем:

positionsNew = nx.spring_layout(G, dim = 3, seed = 1)

Посмотрим координаты для первой команды:

positionsNew['BrighamYoung']
array([ 0.14127402, -0.13225209, -0.0650812 ])

Сохраним координаты узлов x, y, z из набора координат positionsNew в отдельные списки, а в positionsEdges будем хранить все связи нашего графа, аналогично сделаем и для связей:

xNodes = [positionsNew[i][0] for i in list(positionsNew.keys())]
yNodes = [positionsNew[i][1] for i in list(positionsNew.keys())]
zNodes = [positionsNew[i][2] for i in list(positionsNew.keys())]
positionsEdges = G.edges()
xPosEdges = []
yPosEdges = []
zPosEdges = []
for posEdge in positionsEdges:
    xCoords = [positionsNew[posEdge[0]][0], positionsNew[posEdge[1]][0], None]
    xPosEdges.extend(xCoords)
    yCoords = [positionsNew[posEdge[0]][1], positionsNew[posEdge[1]][1], None]
    yPosEdges.extend(yCoords)
    zCoords = [positionsNew[posEdge[0]][2], positionsNew[posEdge[1]][2], None]
    zPosEdges.extend(zCoords)

Далее настроим сам график в plotly, начнем со связей в x, y, z передаем координаты, режим — линии, цвет черный:

plotly_edges = go.Scatter3d(x = xPosEdges,
                            y = yPosEdges,
                            z = zPosEdges,
                            mode = 'lines',
                            line = dict(color = 'black', 
                                        width = 1)
                            ,)

Аналогично с узлами, маркеры выберем неожиданно квадратные. Маркеров в документации оказалось слишком много на разный вкус:

plotly_nodes = go.Scatter3d(x = xNodes,
                            y = yNodes,
                            z = zNodes,
                            mode = 'markers',
                            marker = dict(symbol = 'square',
                                          size = 3,
                                          color = 'black',
                                          line = dict(color = 'black', 
                                                      width = .1))
                                                    ,)

Создадим небольшие настройки осей, отключим подписи и фон, оставим лишь ориентацию сторон — x, y, z. Подпишем наш граф и зададим ширину и высоту на выходе:

settings = {'showbackground': False, 'showticklabels': False}
layout = go.Layout(title = "Test for NTA",
                   width = 1600,
                   height = 900,
                   scene = {'xaxis': settings,
                             'yaxis': settings,
                             'zaxis': settings},
                  )

Ну и финальная часть, все объединяем и отрисовываем:

data = [plotly_edges, plotly_nodes]
fig = go.Figure(data = data, layout = layout)
fig.show()

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

fig.write_html('figure.html')