How to add row stochastic adjacency matrix for clustering and sampling (using NeighborLoader and Metis)

[TOP-RX]

Hello,

I want to add D^(-1)A to data for sampling, where D is the degree matrix and A is the adjacency matrix, while keeping the original adjacency matrix, but I found only original adjacency can be sampled, this is what I did for a simple test:

from torch_geometric_autoscale import get_data

root='/tmp/datasets'
data, in_channels, out_channels = get_data(root, 'products')
deg = data.adj_t.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow_(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
DA = deg_inv_sqrt.view(-1, 1) * deg_inv_sqrt.view(-1, 1) * data.adj_t
data.DA = DA
train_loader = NeighborLoader(data, input_nodes=None,
                              num_neighbors=[-1, -1], shuffle=True)`

For each batch, the shape of batch.DA is always num_nodes * num_nodes, which seems is not sampled by Neighborloader, but batch.adj_t seems has the correct shape.

Besides, I also try to do clustering using METIS, like following:

data.DA = DA
perm, ptr = metis(data.adj_t, num_parts=100, log=True)
data = permute(data, perm, log=True)

However, the output before permute and after permute are same.

I am wondering how to add this D^(-1)A to the data correctly to be sampled and clustered. Thanks.

[rusty1s]

That is currently working as intended. Sampling is only performed on a single graph. If you want to keep the original adjacency matrix, why not do something like

data.adj_t, data.orig_adj_t = DA, data.adj_t

before passing to the NeighborLoader? You can also use e_id of the returned mini-batch to filter the original adjacency matrix.