(prefix, normalize=True, load_walks=False)
| 31 | |
| 32 | |
| 33 | def load_data(prefix, normalize=True, load_walks=False): |
| 34 | def obj2int(n): |
| 35 | return n |
| 36 | |
| 37 | def just_return(n): |
| 38 | return n |
| 39 | G_data = json.load(open(prefix + "-G.json")) |
| 40 | G = json_graph.node_link_graph(G_data) |
| 41 | if isinstance(G.nodes()[0], int): |
| 42 | conversion = obj2int |
| 43 | else: |
| 44 | conversion = just_return |
| 45 | |
| 46 | if os.path.exists(prefix + "-feats.npy"): |
| 47 | feats = np.load(prefix + "-feats.npy") |
| 48 | else: |
| 49 | print("No features present.. Only identity features will be used.") |
| 50 | feats = None |
| 51 | id_map = json.load(open(prefix + "-id_map.json")) |
| 52 | id_map = {conversion(k): int(v) for k, v in id_map.items()} |
| 53 | walks = [] |
| 54 | class_map = json.load(open(prefix + "-class_map.json")) |
| 55 | if isinstance(list(class_map.values())[0], list): |
| 56 | lab_conversion = just_return |
| 57 | else: |
| 58 | lab_conversion = obj2int |
| 59 | |
| 60 | class_map = {conversion(k): lab_conversion(v) |
| 61 | for k, v in class_map.items()} |
| 62 | |
| 63 | # Remove all nodes that do not have val/test annotations |
| 64 | # (necessary because of networkx weirdness with the Reddit data) |
| 65 | broken_count = 0 |
| 66 | for node in G.nodes(): |
| 67 | if 'val' not in G.node[node] or 'test' not in G.node[node]: |
| 68 | G.remove_node(node) |
| 69 | broken_count += 1 |
| 70 | print("Removed {:d} nodes that lacked proper annotations due to networkx \ |
| 71 | versioning issues".format(broken_count)) |
| 72 | |
| 73 | # Make sure the graph has edge train_removed annotations |
| 74 | # (some datasets might already have this..) |
| 75 | print("Loaded data.. now preprocessing..") |
| 76 | for edge in G.edges(): |
| 77 | if (G.node[edge[0]]['val'] or G.node[edge[1]]['val'] or |
| 78 | G.node[edge[0]]['test'] or G.node[edge[1]]['test']): |
| 79 | G[edge[0]][edge[1]]['train_removed'] = True |
| 80 | else: |
| 81 | G[edge[0]][edge[1]]['train_removed'] = False |
| 82 | |
| 83 | if normalize and feats is not None: |
| 84 | from sklearn.preprocessing import StandardScaler |
| 85 | train_ids = np.array([id_map[n] for n in G.nodes() |
| 86 | if not G.node[n]['val'] |
| 87 | and not G.node[n]['test']]) |
| 88 | train_feats = feats[train_ids] |
| 89 | scaler = StandardScaler() |
| 90 | scaler.fit(train_feats) |
no test coverage detected