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")
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):