Spaces:
Runtime error
Runtime error
| import torch | |
| ##import metis | |
| import numpy as np | |
| import networkx as nx | |
| from sklearn.model_selection import train_test_split | |
| class ClusteringMachine(object): | |
| """ | |
| Clustering the graph, feature set, and target. If the graph is not huge enough, we suggest using 'none' here. | |
| """ | |
| def __init__(self, args, graph, field_index, target): | |
| """ | |
| :param args: Arguments object with parameters. | |
| :param graph: Networkx Graph. | |
| :param field_index: field_index matrix (ndarray). | |
| :param target: Target vector (ndarray). | |
| """ | |
| self.args = args | |
| self.graph = graph | |
| self.field_index = field_index | |
| self.target = target | |
| self._set_sizes() | |
| self._set_loss_weight() | |
| def _set_sizes(self): | |
| """ | |
| Setting the field and class count. | |
| """ | |
| self.user_count = self.field_index.shape[0] | |
| self.field_count = np.max(self.field_index)+1 | |
| self.field_size = self.field_index.shape[1] | |
| self.class_count = np.max(self.target)+1 | |
| print("####\tData Info\t####") | |
| print("user count:\t", self.user_count) | |
| print("field count:\t", self.field_count) | |
| print("field size:\t", self.field_size) | |
| print("class count:\t", self.class_count) | |
| def _set_loss_weight(self): | |
| class_weight = self.target.shape[0] / (self.class_count * np.bincount(self.target.squeeze())) | |
| if self.args.weight_balanced == 'True': | |
| self.class_weight = torch.FloatTensor(class_weight) | |
| else: | |
| self.class_weight = torch.ones(self.class_count) | |
| def decompose(self): | |
| """ | |
| Decomposing the graph, partitioning the features and target, creating Torch arrays. | |
| """ | |
| if self.args.clustering_method == "none": | |
| print("\nWithout graph clustering.\n") | |
| self.clusters = [0] | |
| self.cluster_membership = {node: 0 for node in self.graph.nodes()} | |
| print('cluster memebership', self.cluster_membership) | |
| #elif self.args.clustering_method == "metis": | |
| # print("\nMetis graph clustering started.\n") | |
| # self.metis_clustering() | |
| else: | |
| print("\nRandom graph clustering started.\n") | |
| self.random_clustering() | |
| self.generate_data_partitioning() | |
| self.transfer_edges_and_nodes() | |
| def random_clustering(self): | |
| """ | |
| Random clustering the nodes. | |
| """ | |
| self.clusters = [cluster for cluster in range(self.args.cluster_number)] | |
| self.cluster_membership = {node: np.random.choice(self.clusters) for node in self.graph.nodes()} | |
| #def metis_clustering(self): | |
| """ | |
| Clustering the graph with Metis. | |
| """ | |
| # (st, parts) = metis.part_graph(self.graph, self.args.cluster_number, seed=self.args.seed) | |
| # self.clusters = list(set(parts)) | |
| # self.cluster_membership = {node: membership for node, membership in enumerate(parts)} | |
| def generate_data_partitioning(self): | |
| """ | |
| Creating data partitions and train-val-test splits. | |
| """ | |
| if self.args.clustering_method != "metis": | |
| self.sg_nodes = {} | |
| self.sg_targets = {} | |
| self.sg_edges = {} | |
| self.sg_train_nodes = {} | |
| self.sg_val_nodes = {} | |
| self.sg_test_nodes = {} | |
| self.sg_field_index = {} | |
| for cluster in self.clusters: | |
| print('Cluster', cluster) | |
| subgraph = self.graph.subgraph([node for node in sorted(self.graph.nodes()) if self.cluster_membership[node] == cluster]) | |
| self.sg_nodes[cluster] = [node for node in sorted(subgraph.nodes())] | |
| self.sg_targets[cluster] = self.target[self.sg_nodes[cluster],:] | |
| mapper = {node: i for i, node in enumerate(sorted(self.sg_nodes[cluster]))} | |
| self.sg_edges[cluster] = [[mapper[edge[0]], mapper[edge[1]]] for edge in subgraph.edges()] + [[mapper[edge[1]], mapper[edge[0]]] for edge in subgraph.edges()] | |
| self.sg_train_nodes[cluster], sg_val_test_nodes = \ | |
| train_test_split(list(mapper.values()), test_size = 1-self.args.train_ratio, random_state=self.args.seed, shuffle=True) | |
| self.sg_val_nodes[cluster], self.sg_test_nodes[cluster] = \ | |
| train_test_split(sg_val_test_nodes, test_size = 0.5, random_state=self.args.seed, shuffle=True) | |
| self.sg_train_nodes[cluster] = sorted(self.sg_train_nodes[cluster]) | |
| self.sg_val_nodes[cluster] = sorted(self.sg_val_nodes[cluster]) | |
| self.sg_test_nodes[cluster] = sorted(self.sg_test_nodes[cluster]) | |
| self.sg_field_index[cluster] = self.field_index[self.sg_nodes[cluster],:] | |
| def transfer_edges_and_nodes(self): | |
| """ | |
| Transfering the data to PyTorch format. | |
| """ | |
| for cluster in self.clusters: | |
| self.sg_nodes[cluster] = torch.LongTensor(self.sg_nodes[cluster]) | |
| self.sg_targets[cluster] = torch.LongTensor(self.sg_targets[cluster]) | |
| self.sg_edges[cluster] = torch.LongTensor(self.sg_edges[cluster]).t() | |
| self.sg_train_nodes[cluster] = torch.LongTensor(self.sg_train_nodes[cluster]) | |
| self.sg_val_nodes[cluster] = torch.LongTensor(self.sg_val_nodes[cluster]) | |
| self.sg_test_nodes[cluster] = torch.LongTensor(self.sg_test_nodes[cluster]) | |
| self.sg_field_index[cluster] = torch.LongTensor(self.sg_field_index[cluster]) |