Affinity Propagation Clustering and Edge Filtering

[cameronosmith]

Hi, thank you for the amazing library :).

I am trying to implement this paper Learning Physical Graph Representations from Visual Scenes, and they use the label propagation clustering algorithm, where each node is assigned a unique index and at each label propagation iteration its index is replaced with the most common index of its neighbors. This hard clustering algorithm is very simple, but I'm struggling to figure out how to implement it in a MessagePassing format. The node list is large (~4000), so I can't afford to keep a one-hot num-nodes-wide vector at each node and I'm not sure how to keep a list of the neighbors labels at each node in a batch-wise fashion. Is this possible to implement efficiently and if so do you have any ideas on how to do so?

Also, they learn an affinity matrix to filter a grid's edges, which takes the form of concatenating every node pair on the grid, passing it through a function which yields a scalar and thresholding on that scalar to determine if the edge should exist. The way I'm doing this now is converting the graph to dense-node format, and getting every node pair as something like the following:
graph.x[:,:,graph.edge_index] to format as (num_batch x node_dim x 2 x num_node_pairs), passing it into the projection and thresholding functions. Is there a way I can perform this while in batched graph format? Effectively I need to just filter the existing edges by a function of the edge's node features. I'm thinking this is most effective as an edge_weight attribute (binary) instead of actually removing the edges.

[rusty1s]

What you may wanna try is to represent node indices as sparse matrices, and utilize sparse-sparse matrix multiplication, e.g.:

from torch_sparse import SparseTensor
from torch_scatter import scatter_max

adj_t = SparseTensor(row=edge_index[0], col=edge_index[1], sparse_sizes=(N, N)).t()
x = SparseTensor.eye(N)

out = adj_t @ x 

# Get the argmax of each row:
row, col, value = out.coo()

max, argmax = scatter_max(out[col], row, dim=0, dim_size=N)

# Assign most common label:
x = SparseTensor(row=torch.arange(N), col=argmax, sparse_sizes=(N, N))

Regarding your second question, this feels really similar to a classic link prediction problem, e.g.:

row, col = edge_index
out = torch.cat([x[row], x[col]], dim=-1)
# Get a score for every edge, e.g.:
score = f(x)
mask = score > threshold
edge_index = edge_index[:, mask]

[cameronosmith]

Thank you very much for this super useful code, really.
I'm not familiar working with sparse tensors so I'm not sure how to proceed, but when running the code above in calling scatter_max yields the error:

`
----> 1 max, argmax = scatter_max(out[col], row, dim=0, dim_size=N)

RuntimeError: scatter_max() Expected a value of type 'Tensor' for argument 'src' but instead found type 'SparseTensor'.
Position: 0
Value: SparseTensor(row=tensor([ 0, 0, 0, ..., 36099, 36099, 36099], device='cuda:0'),
col=tensor([ 0, 1, 64, ..., 4031, 4094, 4095], device='cuda:0'),
val=tensor([1., 1., 1., ..., 4., 4., 4.], device='cuda:0'),
size=(36100, 4096), nnz=320356, density=0.22%)
Declaration: scatter_max(Tensor src, Tensor index, int dim=-1, Tensor? out=None, int? dim_size=None) -> ((Tensor, Tensor))
Cast error details: Unable to cast Python instance to C++ type (compile in debug mode for details)
`
Do I need to convert the sparse tensors to non-sparse tensors, or is there a scatter_max like function for sparse tensors?

[rusty1s]

Oh, sry. You need to pass in value instead of out[col].