Stacking KNN graph

[fmocking]

Hi, I was trying to figure out how to get NxKxC from edge_index produced by knn_graph. Probably torch.arange with kn_nodes should work to group K neighbors. However, when there are fewer than k+1 nodes in the graph, this method doesn't work. Basically, my main goal is to concatenate the C dim of K neighbors to get Nx(KC) rather than (N*K)xC. Maybe apply a linear operation to K dim rather than C dim. I can not think of a way to do that right now. Please advise.

In other words, is there a way to have a 'dense' representation of the K neighbors?

[rusty1s]

Yes, reshaping only works in case of more than k nodes in the graph.

One way to fix this is to utilize to_dense_batch (untested):

row, col = knn(...)
row = torch.arange(num_nodes).view(-1, 1).repeat(1, k).view(-1)

col, mask = to_dense_batch(col, batch=row, fill_value=-1)

Let me know if this works for you!

[fmocking]

Hi, thank you for the suggestion! I've tested the code. However, I'm getting the following error when k is larger or equal to the size of the graph. Here is the complete example:

from torch_geometric.utils import to_dense_batch
k=4
pos = torch.tensor([[-1., -1.], [-1., 1.], [1., -1.], [1., 1.]])
graph1 = Data(pos=pos)
pos2 = torch.tensor([[0., 0.], [1., 2.], [2., 3.], [1., 1.], [-1., -1.]])
graph2 = Data(pos=pos2)

b = Batch.from_data_list([graph1,graph2])
row, col  = knn_graph(b.pos,k=k, batch=b.batch )

num_nodes = q.pos.shape[0]
row = torch.arange(num_nodes).view(-1, 1).repeat(1, k).view(-1)

col, mask = to_dense_batch(col, batch=row, fill_value=-1)

Here is the error:

torch_scatter/utils.py:12, in broadcast(src, other, dim)
     10 for _ in range(src.dim(), other.dim()):
     11     src = src.unsqueeze(-1)
---> 12 src = src.expand(other.size())
     13 return src

RuntimeError: The expanded size of the tensor (32) must match the existing size (36) at non-singleton dimension 0.  Target sizes: [32].  Tensor sizes: [36]

It works when k <= 3 but it fails when it exceeds the size of the graph with the minimum number of nodes in batch.

[rusty1s]

I tried to run it on my own. Here is a working example:

import torch

from torch_geometric.data import Batch, Data
from torch_geometric.nn import knn_graph
from torch_geometric.utils import to_dense_batch

k = 4
x1 = torch.randn(4, 16)
pos1 = torch.tensor([[-1., -1.], [-1., 1.], [1., -1.], [1., 1.]])
graph1 = Data(x=x1, pos=pos1)
x2 = torch.randn(5, 16)
pos2 = torch.tensor([[0., 0.], [1., 2.], [2., 3.], [1., 1.], [-1., -1.]])
graph2 = Data(x=x2, pos=pos2)

batch = Batch.from_data_list([graph1, graph2])
row, col = knn_graph(batch.pos, k=k, batch=batch.batch)

x_neigh = batch.x[row]
x_neigh, mask = to_dense_batch(x_neigh, batch=col, fill_value=0.)
print(x_neigh.shape)

[fmocking]

I tested it, perfectly working. Thanks!