CODEWITHSUNDEEP
python-3.xpandasnetworkxpyvis
Evan Camila
Evan Camila

Reputation: 23

Pyvis and Networkx: How to make nodes different color based on source or target

I've strung together a few answers from SO and got this far but can't get nodes to show different colors. My code is attempting to get Team nodes to be Red while State nodes are Black. But eventually I'd like to have node colors defined in the source file via the Type attribute (nfl=red, nba=black, mlb=green). Appreciate any help.

Source file csv:

team,state,Weight,type
Knicks,New York,9,nba
Warriors,California,9,nba
Giants,California,3,mlb
49ers,California,6,nfl
Cowboys,Texas,1,nfl
Giants,New York,2,nfl

Code:

import networkx as nx
import pandas as pd
from pyvis.network import Network

df = pd.read_csv('./data_files/sample.csv')

G = nx.from_pandas_edgelist(df,
                            source='team',
                            target='state',
                            edge_attr='Weight')

colors = []

for node in G:
    if node in df["team"].values:
        colors.append("red")
    else: 
        colors.append("black")

print(colors)

nx.draw(G, node_color=colors, with_labels=True)

net = Network(notebook=True, filter_menu=True, cdn_resources='remote')

net.from_nx(G)

net.show("./output/teams.html")

Upvotes: 2

Views: 2206

Answers (2)

Robert Haas
Robert Haas

Reputation: 900

Here's a short solution using gravis to create the graph visualization you requested. Source node colors are determined by the type column, and edge widths come from the Weight column in your DataFrame.

Disclaimer: I'm the author of gravis and provide this answer because use cases like yours were part of the motivation to build this package and therefore it is straightforward to generate a fitting solution with it.

import gravis as gv
import networkx as nx
import pandas as pd

df = pd.read_csv('./data_files/sample.csv')
graph = nx.from_pandas_edgelist(df, source='team', target='state', edge_attr=('Weight', 'type'))
mapping = dict(nfl='red', nba='black', mlb='green')
for (source, target), attr in graph.edges.items():
    graph.nodes[source]['color'] = mapping[attr['type']]
    graph.edges[(source, target)]['size'] = attr['Weight'] / 2

fig = gv.d3(graph)
fig.display()  # opens a web browser with the visualization

Instead of fig.display() you can also use fig.export_html(filename) to create a standalone HTML file, or fig.to_html() to get HTML text you can serve in a web app. If you like Jupyter notebooks, you can also embed the visualization in the output cell: enter image description here

Upvotes: 0

Timeless
Timeless

Reputation: 37857

If you need to conditionally color one of the nodes based on type, you can make a map :

default = "blue"

colors = dict(nfl="red", nba="black", mlb="green")

dmap = df.set_index("team")["type"].map(colors).to_dict()

nx.draw(G, node_color=[dmap.get(n, default) for n in G.nodes], with_labels=True)

NB : If the nodes are unique, you can map both columns together, this way :

dmap = (pd.lreshape(df, {"nodes": ["team", "state"]})
     .set_index("nodes")["type"].map(colors).to_dict())

If you need to set the colors as attributes before saving an html, you can use :

nx.set_node_attributes(G, dmap, "color") # << add this line

net.from_nx(G)

net.write_html("./output/teams.html")

Preview (of teams.html):

enter image description here

Upvotes: 1

Related Questions