def __call__(self, edge_index, x=None, edge_weight=None):
if self.method_type == "emb":
if isinstance(edge_index, np.ndarray):
edge_index = torch.from_numpy(edge_index)
edge_index = (edge_index[:, 0], edge_index[:, 1])
data = Graph(edge_index=edge_index, edge_weight=edge_weight)
self.model = build_model(self.args)
embeddings = self.model(data)
elif self.method_type == "gnn":
num_nodes = edge_index.max().item() + 1
if x is None:
print("No input node features, use random features instead.")
x = np.random.randn(num_nodes, self.num_features)
if isinstance(x, np.ndarray):
x = torch.from_numpy(x).float()
if isinstance(edge_index, np.ndarray):
edge_index = torch.from_numpy(edge_index)
edge_index = (edge_index[:, 0], edge_index[:, 1])
data = Graph(x=x, edge_index=edge_index, edge_weight=edge_weight)
torch.save(data, self.data_path)
dataset = NodeDataset(path=self.data_path, scale_feat=False, metric="accuracy")
self.args.dataset = dataset
model = train(self.args)
embeddings = model.embed(data.to(model.device))
embeddings = embeddings.detach().cpu().numpy()
return embeddings
def build_g_feat(self, A):
edge2type = {}
edges = []
weights = []
for k, mat in enumerate(A):
edges.append(mat[0].cpu().numpy())
weights.append(mat[1].cpu().numpy())
for u, v in zip(*edges[-1]):
edge2type[(u, v)] = k
edges = np.concatenate(edges, axis=1)
weights = np.concatenate(weights)
edges = torch.tensor(edges).to(self.device)
weights = torch.tensor(weights).to(self.device)
g = Graph(edge_index=edges, edge_weight=weights)
g = g.to(self.device)
e_feat = []
for u, v in zip(*g.edge_index):
u = u.cpu().item()
v = v.cpu().item()
e_feat.append(edge2type[(u, v)])
e_feat = torch.tensor(e_feat, dtype=torch.long).to(self.device)
g.edge_type = e_feat
self.g = g
