infomap代码实现

from infomap import Infomap
import tqdm

min_sim = 0.5
word_vecs = np.array(temp_v)
#word_vecs = model.wv.vectors
word_vecs /= (word_vecs ** 2).sum(axis=1, keepdims=True) ** 0.5


word2id = {j: i for i, j in enumerate(sub_n_words)}
new_words = sub_n_words
new_vecs = word_vecs
links = {}

# 每个词找与它相似度不小于0.6的词（不超过50个），来作为图上的边
for i in range(len(new_words)):
    sims = np.dot(new_vecs, new_vecs[i])
    idxs = sims.argsort()[::-1][1:]
    for j in idxs:
        if sims[j] >= min_sim:
            links[(i, j)] = float(sims[j])
        else:
            break

infomapWrapper = Infomap("--two-level --directed")

for (i, j), sim in links.items():
    _ = infomapWrapper.addLink(i, j, sim)

infomapWrapper.run()
tree = infomapWrapper.tree

word2class = {}
class2word = {}
for node in tree:
    if node.is_leaf:
        if new_words[node.node_id] not in word2class:
            word2class[new_words[node.node_id]] = []
        word2class[new_words[node.node_id]].append(node.module_id)
        if node.module_id not in class2word:
            class2word[node.module_id] = []
        class2word[node.module_id].append(new_words[node.node_id])

for i in class2word.keys():
    print(class2word[i])