Karate club embedding visualization

tsne-karate.py

import numpy as np
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

# retrieve the labels for each node
labels = np.asarray([G.nodes[i]['club'] != 'Mr. Hi' for i in G.nodes]).astype(np.int64)

# assigning colours to node labels
color_map = []
for i in labels:
    if i == 0:
        color_map.append('blue')
    else: 
        color_map.append('red')  

# transform the embeddings from 128 dimensions to 2D space
m = TSNE(learning_rate=20, random_state=42)
tsne_features = m.fit_transform(list(embeddings.values()))

# plot the transformed embeddings
plt.figure(figsize=(9,6)) 
plt.scatter(x = tsne_features[:,0], 
            y = tsne_features[:,1],
            c = color_map,
            s =600,
            alpha=0.6)

# adds annotations
for i, label in enumerate(np.arange(0,34)):
    plt.annotate(label, (tsne_features[:,0][i], tsne_features[:,1][i]))

# save the visualization
plt.savefig('tsne.png', bbox_inches='tight',dpi = 1000)