utils_graphsaint.py

import scipy.sparse as sp
import numpy as np
import sys
import json
import os
from sklearn.preprocessing import StandardScaler
from torch_geometric.data import InMemoryDataset, Data
import torch
from itertools import repeat
from torch_geometric.data import NeighborSampler

class DataGraphSAINT:

    def __init__(self, dataset, **kwargs):
        dataset_str='data/'+dataset+'/'
        adj_full = sp.load_npz(dataset_str+'adj_full.npz') 
        self.nnodes = adj_full.shape[0] 
        self.pagerank_scores = None
        if dataset in ['ogbn-arxiv','ogbn-arxiv-xrt']: 
            adj_full = adj_full + adj_full.T
            adj_full[adj_full > 1] = 1

        role = json.load(open(dataset_str+'role.json','r')) 
        idx_train = role['tr'] 
        idx_test = role['te'] 
        idx_val = role['va'] 

        if 'label_rate' in kwargs: 
            label_rate = kwargs['label_rate']
            if label_rate < 1:
                idx_train = idx_train[:int(label_rate*len(idx_train))]

        self.adj_train = adj_full[np.ix_(idx_train, idx_train)]
        self.adj_val = adj_full[np.ix_(idx_val, idx_val)] 
        self.adj_test = adj_full[np.ix_(idx_test, idx_test)] 

        feat = np.load(dataset_str+'feats.npy') 

        feat_train = feat[idx_train] 
        scaler = StandardScaler()
        scaler.fit(feat_train)
        feat = scaler.transform(feat)

        self.feat_train = feat[idx_train] 
        self.feat_val = feat[idx_val] 
        self.feat_test = feat[idx_test] 

        class_map = json.load(open(dataset_str + 'class_map.json','r')) 
        labels = self.process_labels(class_map) 

        self.labels_train = labels[idx_train]
        self.labels_val = labels[idx_val] 
        self.labels_test = labels[idx_test] 

        self.data_full = GraphData(adj_full, feat, labels, idx_train, idx_val, idx_test) 
        self.class_dict = None
        self.class_dict2 = None

        self.adj_full = adj_full 
        self.feat_full = feat 
        self.labels_full = labels 
        self.idx_train = np.array(idx_train) 
        self.idx_val = np.array(idx_val) 
        self.idx_test = np.array(idx_test)
        self.samplers = None

    def process_labels(self, class_map):

        num_vertices = self.nnodes  
        if isinstance(list(class_map.values())[0], list):  
            num_classes = len(list(class_map.values())[0])  
            self.nclass = num_classes 
            class_arr = np.zeros((num_vertices, num_classes))  
            for k,v in class_map.items():
                class_arr[int(k)] = v  
        else:
            class_arr = np.zeros(num_vertices, dtype=np.int_)  
            for k, v in class_map.items():
                class_arr[int(k)] = v  
            class_arr = class_arr - class_arr.min()  
            self.nclass = max(class_arr) + 1  
        return class_arr  

    def retrieve_class(self, c, num=256):
        if self.class_dict is None:
            self.class_dict = {}
            for i in range(self.nclass):
                self.class_dict['class_%s'%i] = (self.labels_train == i)  
        idx = np.arange(len(self.labels_train)) 
        idx = idx[self.class_dict['class_%s'%c]] 
        return np.random.permutation(idx)[:num]  

    
    def compute_pagerank_scores(self,adj):
        # Convert the sparse adjacency matrix to a NetworkX graph
        import networkx as nx
        import numpy as np
        from scipy.sparse import coo_matrix
        # G = nx.from_scipy_sparse_matrix(adj)
        adj_matrix = adj.to_scipy(layout='csr')  

        G = nx.DiGraph(adj_matrix) 

        # pagerank_scores = nx.pagerank(G)
        # Calculate the PageRank scores using NetworkX
        pagerank_scores = nx.pagerank(G)

        # Convert the dictionary of PageRank scores to a NumPy array
        num_nodes = adj.size(0)
        pagerank_array = np.zeros(num_nodes)
        for node, score in pagerank_scores.items():
            pagerank_array[node] = score

        return pagerank_array

    
    def retrieve_class_sampler_pagerank(self, c, adj, transductive, num=256, args=None):
        import torch
        import torch_geometric.utils as pyg_utils
        # Initialize PageRank centrality during the first execution
        if self.pagerank_scores is None:
            self.pagerank_scores = self.compute_pagerank_scores(adj)  # Custom function to compute PageRank centrality

     
        sizes = []
        if args.nlayers == 1:
            sizes = [30]
        if args.nlayers == 2:
            if args.dataset in ['reddit', 'flickr']:
                if args.option == 0:
                    sizes = [15, 8]
                if args.option == 1:
                    sizes = [20, 10]
                if args.option == 2:
                    sizes = [25, 10]
            else:
                sizes = [10, 5]

     
        if self.class_dict2 is None:
            print(sizes)
            self.class_dict2 = {}
            for i in range(self.nclass):
                if transductive:
                    idx_train = np.array(self.idx_train)
                    idx = idx_train[self.labels_train == i]
                else:
                    idx = np.arange(len(self.labels_train))[self.labels_train == i]
                self.class_dict2[i] = idx

    
        if self.samplers is None:
            self.samplers = []
            for i in range(self.nclass):
                node_idx = torch.LongTensor(self.class_dict2[i])
                if len(node_idx) == 0:
                    continue

                self.samplers.append(NeighborSampler(adj,
                                                    node_idx=node_idx,
                                                    sizes=sizes, batch_size=num,
                                                    num_workers=8, return_e_id=False,
                                                    num_nodes=adj.size(0),
                                                    shuffle=True))
        
        pagerank_scores_c = self.pagerank_scores[self.class_dict2[c]]
        pagerank_scores_c = torch.from_numpy(pagerank_scores_c)  
        top_nodes_idx = torch.argsort(pagerank_scores_c, descending=True)[:num]
        batch = self.class_dict2[c][top_nodes_idx]
        

        out = self.samplers[c].sample(batch)
        return out

    def retrieve_class_sampler_mutil(self, c_s, adj, transductive, num=256, args=None):
        sizes = []
        if args.nlayers == 1:
            sizes = [30]
        if args.nlayers == 2:
            if args.dataset in ['reddit', 'flickr']:
                if args.option == 0:
                    sizes = [15, 8]
                if args.option == 1:
                    sizes = [20, 10]
                if args.option == 2:
                    sizes = [25, 10]
            else:
                sizes = [10, 5]

        if self.class_dict2 is None:
            print(sizes)
            self.class_dict2 = {}
            for i in range(self.nclass):
                if transductive:
                    idx_train = np.array(self.idx_train)
                    idx = idx_train[self.labels_train == i]
                else:
                    idx = np.arange(len(self.labels_train))[self.labels_train == i]
                self.class_dict2[i] = idx

        if self.samplers is None:
            self.samplers = {}
            for i in range(0, self.nclass, len(c_s)):
                # print(i)
                # classes = min(self.nclass, i + len(c_s))
                for j in range(len(c_s)):
                    if j == 0:
                        node_idx = torch.LongTensor(self.class_dict2[i + j])
                    else:
                        node_idx = torch.cat([node_idx, torch.LongTensor(self.class_dict2[i + j])], dim=0)

                if len(node_idx) == 0:
                    continue

                self.samplers[i] = (NeighborSampler(adj,
                                                     node_idx=node_idx,
                                                     sizes=sizes, batch_size=num*len(c_s),
                                                     num_workers=8, return_e_id=False,
                                                     num_nodes=adj.size(0),
                                                     shuffle=True))
        number = c_s[0]
        batches = []

        for i in range(number, number + len(c_s)):
            samples = np.random.permutation(self.class_dict2[i])[:num]
            batches.append(samples)

        batch = np.concatenate(batches, axis=0)

        out = self.samplers[number].sample(batch)
        return out

    def retrieve_class_sampler(self, c, adj, transductive, num=256, args=None):
        sizes = []
        if args.nlayers == 1:
            sizes = [30]
        if args.nlayers == 2:
            if args.dataset in ['reddit', 'flickr']:
                if args.option == 0:
                    sizes = [15, 8]
                if args.option == 1:
                    sizes = [20, 10]
                if args.option == 2:
                    sizes = [25, 10]
            else:
                sizes = [10, 5]

        if self.class_dict2 is None:
            print(sizes)
            self.class_dict2 = {}
            for i in range(self.nclass):
                if transductive:
                    idx_train = np.array(self.idx_train)
                    idx = idx_train[self.labels_train == i]
                else:
                    idx = np.arange(len(self.labels_train))[self.labels_train == i]
                self.class_dict2[i] = idx

        if self.samplers is None:
            self.samplers = []
            for i in range(self.nclass):
                node_idx = torch.LongTensor(self.class_dict2[i])
                if len(node_idx) == 0:
                    continue

                self.samplers.append(NeighborSampler(adj,
                                                     node_idx=node_idx,
                                                     sizes=sizes, batch_size=num,
                                                     num_workers=8, return_e_id=False,
                                                     num_nodes=adj.size(0),
                                                     shuffle=True))
        batch = np.random.permutation(self.class_dict2[c])[:num]
        out = self.samplers[c].sample(batch)
        return out

class GraphData:

    def __init__(self, adj, features, labels, idx_train, idx_val, idx_test):

        self.adj = adj
        self.features = features
        self.labels = labels
        self.idx_train = idx_train
        self.idx_val = idx_val
        self.idx_test = idx_test


class Data2Pyg:
    def __init__(self, data, device='cuda', transform=None, **kwargs):
        self.data_train = Dpr2Pyg(data.data_train, transform=transform)[0].to(device) 
        self.data_val = Dpr2Pyg(data.data_val, transform=transform)[0].to(device) 
        self.data_test = Dpr2Pyg(data.data_test, transform=transform)[0].to(device) 
        self.nclass = data.nclass 
        self.nfeat = data.nfeat 
        self.class_dict = None 

    def retrieve_class(self, c, num=256):
        if self.class_dict is None:
            self.class_dict = {}
            for i in range(self.nclass):
                self.class_dict['class_%s'%i] = (self.data_train.y == i).cpu().numpy() 
        idx = np.arange(len(self.data_train.y)) 
        idx = idx[self.class_dict['class_%s'%c]] 
        return np.random.permutation(idx)[:num] 



class Dpr2Pyg(InMemoryDataset):

    def __init__(self, dpr_data, transform=None, **kwargs):
        root = 'data/' 
        self.dpr_data = dpr_data
        super(Dpr2Pyg, self).__init__(root, transform)
        pyg_data = self.process()
        self.data, self.slices = self.collate([pyg_data])
        self.transform = transform

    def process(self):
        dpr_data = self.dpr_data
        edge_index = torch.LongTensor(dpr_data.adj.nonzero())  
        if sp.issparse(dpr_data.features):  
            x = torch.FloatTensor(dpr_data.features.todense()).float()  
        else:
            x = torch.FloatTensor(dpr_data.features).float()  
        y = torch.LongTensor(dpr_data.labels)  
        data = Data(x=x, edge_index=edge_index, y=y)  
        data.train_mask = None  
        data.val_mask = None 
        data.test_mask = None 
        return data


    def get(self, idx):

        data = self.data.__class__()

        if hasattr(self.data, '__num_nodes__'):
            data.num_nodes = self.data.__num_nodes__[idx]

        for key in self.data.keys:
            item, slices = self.data[key], self.slices[key]
            s = list(repeat(slice(None), item.dim()))
            s[self.data.__cat_dim__(key, item)] = slice(slices[idx], slices[idx + 1])
            data[key] = item[s]

        return data

    @property
    def raw_file_names(self):

        return ['some_file_1', 'some_file_2', ...]  

    @property
    def processed_file_names(self):
        return ['data.pt']  

    def _download(self):

        pass