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

Issue

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

Solution

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

Answered By - Timeless