python算法学习双曲嵌入论文代码实现数据集介绍

import nltk
# nltk.download('wordnet') # 第一次运行需运行此命令，安装wordnet数据集
from nltk.corpus import wordnet as wn
from math import *
import random
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import networkx as nx

network = {} # 构建层级网络
last_level = 8 # 最深的层设置为8层
levelOfNode = {} # 数据的层级信息，0为哺乳动物（根节点），1为哺乳动物下一结构
# 递归构建network
def get_hyponyms(synset, level):
    if (level == last_level):
        levelOfNode[str(synset)] = level
        return
    if not str(synset) in network:
        network[str(synset)] = [str(s) for s in synset.hyponyms()]
        levelOfNode[str(synset)] = level
    for hyponym in synset.hyponyms():
        get_hyponyms(hyponym, level + 1)
# 构建以哺乳动物为根节点的层次结构数据集
mammal = wn.synset('mammal.n.01')
get_hyponyms(mammal, 0)
levelOfNode[str(mammal)] = 0
# 将终端叶子节点补到network字典中
for a in levelOfNode:
    if not a in network:
        network[a] = []

def norm(x):
    return np.dot(x, x)

def traverse(graph, start, node):
    node_name = node.name().split(".")[0]
    graph.depth[node_name] = node.shortest_path_distance(start)
    for child in node.hyponyms():
        child_name = child.name().split(".")[0]
        graph.add_edge(node_name, child_name) # 添加边
        traverse(graph, start, child) # 递归构建

def hyponym_graph(start):
    G = nx.Graph() # 定义一个图
    G.depth = {}
    traverse(G, start, start)
    return G

def graph_draw(graph):
    plt.figure(figsize=(10, 10)) # 展示整体的网络
    # plt.figure(figsize=(3, 3)) # 展示大象网络
    nx.draw(graph,
         node_size = [10 * graph.degree(n) for n in graph],
         node_color = [graph.depth[n] for n in graph],
         alpha = 0.8,
         font_size = 4,
         width = 0.5,
         with_labels = True)
    
def get_keys(d, value):
        return [k for k,v in d.items() if v == value]
    root_name = get_keys(graph.depth, 0)[0]
    plt.savefig("~/hyperE/fig/" + root_name + ".png", dpi = 300)
    
graph = hyponym_graph(mammal)
graph_draw(graph)

elephant = wn.synset('elephant.n.01')
graph = hyponym_graph(elephant)
graph_draw(graph)