def test(model, data, evaluator): print('Evaluating full-batch GNN on CPU...') weights = [(conv.lin_rel.weight.t().cpu().detach().numpy(), conv.lin_rel.bias.cpu().detach().numpy(), conv.lin_root.weight.t().cpu().detach().numpy()) for conv in model.convs] model = SAGEInference(weights) x = data.x.numpy() adj = SparseTensor(row=data.edge_index[0], col=data.edge_index[1]) adj = adj.sum(dim=1).pow(-1).view(-1, 1) * adj adj = adj.to_scipy(layout='csr') out = model(x, adj) y_true = data.y y_pred = torch.from_numpy(out).argmax(dim=-1, keepdim=True) train_acc = evaluator.eval({ 'y_true': y_true[data.train_mask], 'y_pred': y_pred[data.train_mask] })['acc'] valid_acc = evaluator.eval({ 'y_true': y_true[data.valid_mask], 'y_pred': y_pred[data.valid_mask] })['acc'] test_acc = evaluator.eval({ 'y_true': y_true[data.test_mask], 'y_pred': y_pred[data.test_mask] })['acc'] return train_acc, valid_acc, test_acc
def get_adj(row, col, N, asymm_norm=False, set_diag=True, remove_diag=False): adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N)) if set_diag: print('... setting diagonal entries') adj = adj.set_diag() elif remove_diag: print('... removing diagonal entries') adj = adj.remove_diag() else: print('... keeping diag elements as they are') if not asymm_norm: print('... performing symmetric normalization') deg = adj.sum(dim=1).to(torch.float) deg_inv_sqrt = deg.pow(-0.5) deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0 adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1) else: print('... performing asymmetric normalization') deg = adj.sum(dim=1).to(torch.float) deg_inv = deg.pow(-1.0) deg_inv[deg_inv == float('inf')] = 0 adj = deg_inv.view(-1, 1) * adj adj = adj.to_scipy(layout='csr') return adj
def spectral(data, name, embedding_dim=128): try: result = load_embedding(name, 'spectral') return result except FileNotFoundError: print( f'cache/feature/spectral_{name}.pt not found! Regenerating it now') from julia.api import Julia jl = Julia(compiled_modules=False) from julia import Main Main.include(f'{CWD}/norm_spec.jl') print('Setting up spectral embedding') if data.setting == 'inductive': N = data.num_train_nodes edge_index = to_undirected(data.train_edge_index, num_nodes=data.num_train_nodes) else: N = data.num_nodes edge_index = to_undirected(data.edge_index, num_nodes=data.num_nodes) np_edge_index = np.array(edge_index.T) row, col = edge_index adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N)) adj = adj.to_scipy(layout='csr') result = torch.tensor(Main.main(adj, embedding_dim)).float() save_embedding(result, name, 'spectral') return result
def init_adj(self, edge_index): """ cache normalized adjacency and normalized strict two-hop adjacency, neither has self loops """ n = self.num_nodes if isinstance(edge_index, SparseTensor): dev = adj_t.device adj_t = edge_index adj_t = scipy.sparse.csr_matrix(adj_t.to_scipy()) adj_t[adj_t > 0] = 1 adj_t[adj_t < 0] = 0 adj_t = SparseTensor.from_scipy(adj_t).to(dev) elif isinstance(edge_index, torch.Tensor): row, col = edge_index adj_t = SparseTensor(row=col, col=row, value=None, sparse_sizes=(n, n)) adj_t.remove_diag(0) adj_t2 = matmul(adj_t, adj_t) adj_t2.remove_diag(0) adj_t = scipy.sparse.csr_matrix(adj_t.to_scipy()) adj_t2 = scipy.sparse.csr_matrix(adj_t2.to_scipy()) adj_t2 = adj_t2 - adj_t adj_t2[adj_t2 > 0] = 1 adj_t2[adj_t2 < 0] = 0 adj_t = SparseTensor.from_scipy(adj_t) adj_t2 = SparseTensor.from_scipy(adj_t2) adj_t = gcn_norm(adj_t, None, n, add_self_loops=False) adj_t2 = gcn_norm(adj_t2, None, n, add_self_loops=False) self.adj_t = adj_t.to(edge_index.device) self.adj_t2 = adj_t2.to(edge_index.device)
def preprocess(data, preprocess="diffusion", num_propagations=10, p=None, alpha=None, use_cache=True, post_fix=""): if use_cache: try: x = torch.load(f'embeddings/{preprocess}{post_fix}.pt') print('Using cache') return x except: print( f'embeddings/{preprocess}{post_fix}.pt not found or not enough iterations! Regenerating it now' ) # Creates a new file with open(f'embeddings/{preprocess}{post_fix}.pt', 'w') as fp: pass if preprocess == "community": return community(data, post_fix) if preprocess == "spectral": return spectral(data, post_fix) print('Computing adj...') N = data.num_nodes data.edge_index = to_undirected(data.edge_index, data.num_nodes) row, col = data.edge_index adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N)) adj = adj.set_diag() deg = adj.sum(dim=1).to(torch.float) deg_inv_sqrt = deg.pow(-0.5) deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0 adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1) adj = adj.to_scipy(layout='csr') sgc_dict = {} print(f'Start {preprocess} processing') if preprocess == "sgc": result = sgc(data.x.numpy(), adj, num_propagations) # if preprocess == "lp": # result = lp(adj, data.y.data, num_propagations, p = p, alpha = alpha, preprocess = preprocess) if preprocess == "diffusion": result = diffusion(data.x.numpy(), adj, num_propagations, p=p, alpha=alpha) torch.save(result, f'embeddings/{preprocess}{post_fix}.pt') return result
def test(model, predictor, data, split_edge, evaluator, batch_size, device): predictor.eval() print('Evaluating full-batch GNN on CPU...') weights = [(conv.weight.cpu().detach().numpy(), conv.bias.cpu().detach().numpy()) for conv in model.convs] model = GCNInference(weights) x = data.x.numpy() adj = SparseTensor(row=data.edge_index[0], col=data.edge_index[1]) adj = adj.set_diag() deg = adj.sum(dim=1) deg_inv_sqrt = deg.pow(-0.5) deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0 adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1) adj = adj.to_scipy(layout='csr') h = torch.from_numpy(model(x, adj)).to(device) def test_split(split): source = split_edge[split]['source_node'].to(device) target = split_edge[split]['target_node'].to(device) target_neg = split_edge[split]['target_node_neg'].to(device) pos_preds = [] for perm in DataLoader(range(source.size(0)), batch_size): src, dst = source[perm], target[perm] pos_preds += [predictor(h[src], h[dst]).squeeze().cpu()] pos_pred = torch.cat(pos_preds, dim=0) neg_preds = [] source = source.view(-1, 1).repeat(1, 1000).view(-1) target_neg = target_neg.view(-1) for perm in DataLoader(range(source.size(0)), batch_size): src, dst_neg = source[perm], target_neg[perm] neg_preds += [predictor(h[src], h[dst_neg]).squeeze().cpu()] neg_pred = torch.cat(neg_preds, dim=0).view(-1, 1000) return evaluator.eval({ 'y_pred_pos': pos_pred, 'y_pred_neg': neg_pred, })['mrr_list'].mean().item() train_mrr = test_split('eval_train') valid_mrr = test_split('valid') test_mrr = test_split('test') return train_mrr, valid_mrr, test_mrr
def spectral(data, post_fix): from julia.api import Julia jl = Julia(compiled_modules=False) from julia import Main Main.include("./norm_spec.jl") print('Setting up spectral embedding') data.edge_index = to_undirected(data.edge_index) np_edge_index = np.array(data.edge_index.T) N = data.num_nodes row, col = data.edge_index adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N)) adj = adj.to_scipy(layout='csr') result = torch.tensor(Main.main(adj, 128)).float() torch.save(result, f'embeddings/spectral{post_fix}.pt') return result
def svd(data, name, embedding_dim=64): try: result = load_embedding(name, 'svd') return result except FileNotFoundError: print(f'cache/feature/svd_{name}.pt not found! Regenerating it now') if data.setting == 'inductive': N = data.num_train_nodes row, col = data.train_edge_index else: N = data.num_nodes row, col = data.edge_index adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N)) adj = adj.to_scipy(layout='csc') ut, s, vt = sparsesvd(adj, embedding_dim) result = torch.tensor(np.dot(ut.T, np.diag(s)), dtype=torch.float) save_embedding(result, name, 'svd') return result
def spectral(data, post_fix): # from julia.api import Julia # jl = Julia(compiled_modules=False) # from julia import Main # Main.include("./norm_spec.jl") print('Setting up spectral embedding') data.edge_index = to_undirected(data.edge_index) np_edge_index = np.array(data.edge_index.T) N = data.num_nodes row, col = data.edge_index adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N)) adj = adj.to_scipy(layout='csr') tsvd = TruncatedSVD(n_components=128) adj_tsvd = tsvd.fit(adj).transform(adj) result = torch.tensor(adj_tsvd).float() # result = torch.tensor(adj.todense()).float() torch.save(result, f'embeddings/spectral{post_fix}.pt') return result
def main(): parser = argparse.ArgumentParser(description='OGBN-papers100M (MLP)') parser.add_argument('--data_root_dir', type=str, default='../../dataset') parser.add_argument('--num_propagations', type=int, default=3) parser.add_argument('--dropedge_rate', type=float, default=0.4) parser.add_argument('--node_emb_path', type=str, default=None) parser.add_argument('--output_path', type=str, required=True) args = parser.parse_args() # SGC pre-processing ###################################################### dataset = PygNodePropPredDataset(name='ogbn-papers100M', root=args.data_root_dir) split_idx = dataset.get_idx_split() data = dataset[0] x = None if args.node_emb_path: x = np.load(args.node_emb_path) else: x = data.x.numpy() N = data.num_nodes print('Making the graph undirected.') ### Randomly drop some edges to save computation data.edge_index, _ = dropout_adj(data.edge_index, p=args.dropedge_rate, num_nodes=data.num_nodes) data.edge_index = to_undirected(data.edge_index, data.num_nodes) print(data) row, col = data.edge_index print('Computing adj...') adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N)) adj = adj.set_diag() deg = adj.sum(dim=1).to(torch.float) deg_inv_sqrt = deg.pow(-0.5) deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0 adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1) adj = adj.to_scipy(layout='csr') train_idx, valid_idx, test_idx = split_idx['train'], split_idx[ 'valid'], split_idx['test'] all_idx = torch.cat([train_idx, valid_idx, test_idx]) mapped_train_idx = torch.arange(len(train_idx)) mapped_valid_idx = torch.arange(len(train_idx), len(train_idx) + len(valid_idx)) mapped_test_idx = torch.arange( len(train_idx) + len(valid_idx), len(train_idx) + len(valid_idx) + len(test_idx)) sgc_dict = {} sgc_dict['label'] = data.y.data[all_idx].to(torch.long) sgc_dict['split_idx'] = { 'train': mapped_train_idx, 'valid': mapped_valid_idx, 'test': mapped_test_idx } print('Start SGC processing') for _ in tqdm(range(args.num_propagations)): x = adj @ x sgc_dict['sgc_embedding'] = torch.from_numpy(x[all_idx]).to(torch.float) torch.save(sgc_dict, args.output_path)
def amfs( A: SparseTensor, Sigma=None, level=None, delta=0.1, thresh_kld=1e-6, priority=True, verbose=False, ) -> Tuple[List[lil_matrix], np.ndarray]: r""" AMFS bipartite approximation for graph wavelet signal processing [3]_. Parameters ---------- A: SparseTensor The adjacency matrix. Sigma: scipy.spmatrix, optional The covariance matrix specified by the Laplacian matrix L. If None, :math:`\Sigma^{-1}=L+\delta I` level: int, optional The number of bipartite subgraphs, i.e., the decomposition level. If None, :math:`level=\lceil log_2( \mathcal{X}) \rceil`, where :math:`\mathcal{X}` is the chromatic number of :obj:`A`. delta: float, optional :math:`1/\delta` is interpreted as the variance of the DC compnent. Refer to [4]_ for more details. thresh_kld: float, optional Threshold of Kullback-Leibler divergence to perform `AMFS` decomposition. priority: bool,optional If True, KLD holds priority. verbose: bool,optional Returns ------- bptG: List[SparseTensor] The bipartite subgraphs. beta: Tensor(N, M) The indicator of bipartite sets References ---------- .. [3] Jing Zen, et al, "Bipartite Subgraph Decomposition for Critically Sampledwavelet Filterbanks on Arbitrary Graphs," IEEE trans on SP, 2016. .. [4] A. Gadde, et al, "A probablistic interpretation of sampling theory of graph signals". ICASSP, 2015. """ N = A.size(-1) # compute_sigma consists of laplace matrix which prefers "coo" A = A.to_scipy(layout="coo").astype("d") if Sigma is None: Sigma = compute_sigma(A, delta) else: assert Sigma.shape == (N, N) if level is None: chromatic = dsatur(A).n_color level = np.ceil(np.log2(chromatic)) A = A.tolil() beta = np.zeros((N, level), dtype=bool) bptG = [lil_matrix((N, N), dtype=A.dtype) for _ in range(level)] for i in range(level): if verbose: print(f"\n|----decomposition in level: {i:4d} ----|") s1, s2 = amfs1level(A, Sigma, delta, thresh_kld, priority, verbose) bt = beta[:, i] bt[s1] = 1 # set s1 True mask = bipartite_mask(bt) bptG[i][mask] = A[mask] A[mask] = 0 return bptG, beta
def osglm(A: SparseTensor, lc: Optional[int] = None, vtx_color: Optional[VertexColor] = None): r""" The oversampled bipartite graph approximation method proposed in [1]_ Parameters ---------- A: SparseTensor The adjacent matrix of graph. lc: int The ordinal of color marking the boundary such that all nodes with a smaller color ordinal are grouped into the low-pass channel while those with a larger color ordinal are in the high-pass channel. vtx_color:iter The graph coloring result Returns ------- bptG: lil_matrix The oversampled graph(with additional nodes) beta : np.ndarray append_nodes: np.ndarray The indices of those appended nodes vtx_color: np.ndarray The node colors References ---------- .. [1] Akie Sakiyama, et al, "Oversampled Graph Laplacian Matrix for Graph Filter Banks", IEEE trans on SP, 2016. """ if vtx_color is None: from thgsp.alg import dsatur vtx_color = dsatur(A) vtx_color = np.asarray(vtx_color) n_color = max(vtx_color) + 1 if lc is None: lc = n_color // 2 assert 1 <= lc < n_color A = A.to_scipy(layout="csr").tolil() # the foundation bipartite graph Gb Gb = lil_matrix(A.shape, dtype=A.dtype) N = A.shape[-1] bt = np.in1d(vtx_color, range(lc)) idx_s1 = np.nonzero(bt)[0] # L idx_s2 = np.nonzero(~bt)[0] # H mask = bipartite_mask(bt) # the desired edges Gb[mask] = A[mask] A[mask] = 0 eye_mask = eye(N, N, dtype=bool) A[eye_mask] = 1 # add vertical edges degree = A.sum(0).getA1() # 2D np.matrix -> 1D np.array append_nodes = (degree != 0).nonzero()[0] Nos = len(append_nodes) + N # oversampled size bptG = [lil_matrix((Nos, Nos), dtype=A.dtype)] # the expanded graph bptG[0][:N, N:] = A[:, append_nodes] bptG[0][:N, :N] = Gb bptG[0][N:, :N] = A[append_nodes, :] beta = np.zeros((Nos, 1), dtype=bool) beta[idx_s1, 0] = 1 # appended nodes corresponding to idx_s2 are assigned to # the L channel of oversampled graph with idx_s1 _, node_ordinal_append, _ = np.intersect1d(append_nodes, idx_s2, return_indices=True) beta[N + node_ordinal_append, 0] = 1 return bptG, beta, append_nodes, vtx_color
def tree_decomposition(mol, return_vocab=False): r"""The tree decomposition algorithm of molecules from the `"Junction Tree Variational Autoencoder for Molecular Graph Generation" <https://arxiv.org/abs/1802.04364>`_ paper. Returns the graph connectivity of the junction tree, the assignment mapping of each atom to the clique in the junction tree, and the number of cliques. Args: mol (rdkit.Chem.Mol): A :obj:`rdkit` molecule. return_vocab (bool, optional): If set to :obj:`True`, will return an identifier for each clique (ring, bond, bridged compounds, single). (default: :obj:`False`) :rtype: (LongTensor, LongTensor, int) """ import rdkit.Chem as Chem # Cliques = rings and bonds. cliques = [list(x) for x in Chem.GetSymmSSSR(mol)] xs = [0] * len(cliques) for bond in mol.GetBonds(): if not bond.IsInRing(): cliques.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()]) xs.append(1) # Generate `atom2clique` mappings. atom2clique = [[] for i in range(mol.GetNumAtoms())] for c in range(len(cliques)): for atom in cliques[c]: atom2clique[atom].append(c) # Merge rings that share more than 2 atoms as they form bridged compounds. for c1 in range(len(cliques)): for atom in cliques[c1]: for c2 in atom2clique[atom]: if c1 >= c2 or len(cliques[c1]) <= 2 or len(cliques[c2]) <= 2: continue if len(set(cliques[c1]) & set(cliques[c2])) > 2: cliques[c1] = set(cliques[c1]) | set(cliques[c2]) xs[c1] = 2 cliques[c2] = [] xs[c2] = -1 cliques = [c for c in cliques if len(c) > 0] xs = [x for x in xs if x >= 0] # Update `atom2clique` mappings. atom2clique = [[] for i in range(mol.GetNumAtoms())] for c in range(len(cliques)): for atom in cliques[c]: atom2clique[atom].append(c) # Add singleton cliques in case there are more than 2 intersecting # cliques. We further compute the "initial" clique graph. edges = {} for atom in range(mol.GetNumAtoms()): cs = atom2clique[atom] if len(cs) <= 1: continue # Number of bond clusters that the atom lies in. bonds = [c for c in cs if len(cliques[c]) == 2] # Number of ring clusters that the atom lies in. rings = [c for c in cs if len(cliques[c]) > 4] if len(bonds) > 2 or (len(bonds) == 2 and len(cs) > 2): cliques.append([atom]) xs.append(3) c2 = len(cliques) - 1 for c1 in cs: edges[(c1, c2)] = 1 elif len(rings) > 2: cliques.append([atom]) xs.append(3) c2 = len(cliques) - 1 for c1 in cs: edges[(c1, c2)] = 99 else: for i in range(len(cs)): for j in range(i + 1, len(cs)): c1, c2 = cs[i], cs[j] count = len(set(cliques[c1]) & set(cliques[c2])) edges[(c1, c2)] = min(count, edges.get((c1, c2), 99)) # Update `atom2clique` mappings. atom2clique = [[] for i in range(mol.GetNumAtoms())] for c in range(len(cliques)): for atom in cliques[c]: atom2clique[atom].append(c) if len(edges) > 0: edge_index_T, weight = zip(*edges.items()) row, col = torch.tensor(edge_index_T).t() inv_weight = 100 - torch.tensor(weight) clique_graph = SparseTensor(row=row, col=col, value=inv_weight, sparse_sizes=(len(cliques), len(cliques))) junc_tree = minimum_spanning_tree(clique_graph.to_scipy('csr')) row, col, _ = SparseTensor.from_scipy(junc_tree).coo() edge_index = torch.stack([row, col], dim=0) edge_index = to_undirected(edge_index, num_nodes=len(cliques)) else: edge_index = torch.empty((2, 0), dtype=torch.long) rows = [[i] * len(atom2clique[i]) for i in range(mol.GetNumAtoms())] row = torch.tensor(list(chain.from_iterable(rows))) col = torch.tensor(list(chain.from_iterable(atom2clique))) atom2clique = torch.stack([row, col], dim=0).to(torch.long) if return_vocab: vocab = torch.tensor(xs, dtype=torch.long) return edge_index, atom2clique, len(cliques), vocab else: return edge_index, atom2clique, len(cliques)
def harary( A: SparseTensor, vtx_color: Optional[VertexColor] = None, threshold: float = 0.97 ) -> Tuple[List[lil_matrix], np.ndarray, np.ndarray, np.ndarray, Dict[int, int]]: """ Harary bipartite decomposition Parameters ---------- A: :py:class:`SparseTensor` The adjacency matrix vtx_color: array_like, optional All valid type for :py:func:`np.asarray` is acceptable, including :py:class:`torch.Tensor` on cpu. If None, this function will invoke :py:func:`thgsp.alg.dsatur` silently. threshold: float, optional Returns ------- bptG: List[lil_matrix] An array consisting of :obj:`M` bipartite subgraphs formatted as :class:`scipy.sparse.lil_matrix` beta: array :obj:`beta[:,i]` is the bipartite set indicator of :obj:`i`-th subgraph beta_dist: array A table showing the relationship between :obj:`beta` and :obj:`channels` new_vtx_color: array The node colors mapper: dict Map **new_vtx_color** to the original ones. For example mapper={1:2, 2:3, 3:1} map 1,2 and 3-th color to 2,3 and 1, respectively. """ if vtx_color is None: vtx_color = dsatur(A) vtx_color = np.asarray(vtx_color) n_color = max(vtx_color) + 1 if n_color > 256: raise RuntimeError( "Too many colors will lead to a too complicated channel division") A = A.to_scipy(layout="csr").tolil() M = int(np.ceil(np.log2(n_color))) # the number of bipartite graphs N = A.shape[-1] # the number of nodes new_color_ordinal = new_order(n_color) mapper = {c: i for i, c in enumerate(new_color_ordinal)} new_vtx_color = [mapper[c] for c in vtx_color] beta_dist = distribute_color(n_color, M) bptG = [lil_matrix((N, N), dtype=A.dtype) for _ in range(M)] link_weights = -np.ones(M) beta = np.zeros((N, M), dtype=bool) for i in range(M): colors_L = (beta_dist[:, i] == 1).nonzero()[0] bt = np.in1d(new_vtx_color, colors_L) beta[:, i] = bt mask = bipartite_mask(bt) bpt_edges = A[mask] bptG[i][mask] = bpt_edges link_weights[i] = bpt_edges.sum() A[mask] = 0 ratio_link_weights = link_weights.cumsum(0) / link_weights.sum() bpt_idx = (ratio_link_weights >= threshold).nonzero()[0] M1 = bpt_idx[0] + 1 bptG = bptG[:M1] max_color = np.power(2, M1) beta_dist = distribute_color(max_color, M1) beta = beta[:, :M1] return bptG, beta, beta_dist, vtx_color, mapper