I'm new to Graph Convolutional Networks (and pytorch in general) so I'm trying to verify that the message passing layer is working as expected before I go on to adding layers to the network.
But when I look at the output of the propagation, it doesn't line up with what I expect given a mean aggregation of the nodes that I've done manually.
My data are a collection of single cells (n) with transcript counts for each (p). I've created an adjacency matrix for these cells based on their similarity in the dimensionally-reduced space (using UMAP and k-nearest neighbors).
What I'm trying to do next is essentially smooth the values of the transcript counts in each node (cell) based on aggregating the transcript counts of neighboring cells.
I've constructed the message passing layer and deactivated many of the functions that would ordinarily introduce some transformation in the node values. But when I take the first cell in the matrix and manually average its features based on its own local neighbors, I get vastly different feature values.
What am I doing wrong? Have a made a mistake in the code? Or is there still some part of the message passing that is transforming the data?
For most of my code I borrowed heavily from this example on PyTorch geometric's docs https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html
Reproducible code below:
#pytorch imports
import torch
import torch.nn as nn
from torch.nn import functional as F
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_remaining_self_loops, degree
#other imports
import numpy as np
import os
import sys
import pandas as pd
#Create GCN Class
class TranscriptConv(MessagePassing):
def __init__(self, in_channels, out_channels): #Just message passing - so in_channels and out_channels are equal (ie. number of features in = number of features out)
super().__init__(aggr = 'mean')
#self.lin = nn.Linear(in_channels, out_channels, bias=False) #commented out to avoid transformation
#self.bias = nn.Parameter(torch.Tensor(out_channels))
#self.reset_parameters()
#def reset_parameters(self): #Called during __init__ to make sure all parameters are zeroed.
#self.lin.reset_parameters()
#self.bias.data.zero_()
def message(self, x_j, norm):
#flow = "source_to_target" #this doesn't seem to be declared in the example (because 'source to target is the default')
return norm.view(-1,1) * x_j # .view() creates a new tensor, with -1 indicating that this dimension is inferred from the input
def forward(self, data): #input to this is the transcript x cell data itself and the edge_index
x = data.X
edge_index = data.edges
edge_index, _ = add_remaining_self_loops(edge_index, num_nodes= data.len) #this is needed so that it's own value is included in the mean
#x = self.lin(x) #linear summation of input x. commented out to avoid transformation
#degree calculation, normalization
#N = D^(-.5)AD^(-.5) #This isn't code - just the equation written down.
#For directed graph the above has become N = D^(-1)A
row, col = edge_index
#row is list of edges coming in with each value an index to a node. col is list of edges going out with each value an index to a node.
deg = degree(row, x.size(0), dtype=x.dtype) #This is the diagonal matrix D in the above equation. degree only calculates on 1-dimensional index
#deg_inv_sqrt = deg.pow(-0.5) #element wise power. changed to .pow(1) below for directed edges
deg_inv_sqrt = deg.pow(-1)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0 #values which were 0 before being inversed put back to 0
#when the edge is bi-direcitonal, it has to be normalised by in-degrees at target and source. for directed edge, just normalised by indegree at target.
norm = deg_inv_sqrt[row] #put in same order as edge index
norm = torch.ones(row.size()) # no normalisation here just to examine aggregation
#And propogate. This is where the magic happens: function calls message(), aggregate(), and update() internally
out = self.propagate(edge_index, x=x, norm=norm) #norm is required argument.
#Leave intercept in, but not sure if necessary
#out += self.bias #commented out here to avoid transformation
return out
#Create toy data
num_cells =100
num_features = 10
feature_mat = np.matrix(np.random.normal(size = (num_cells, num_features)))
adj_mat = np.matrix(np.random.binomial(p = 0.1, n = 1, size = (num_cells, num_cells)))
adj_mat.diagonal = 1 #set self-loops to 1 for the manual aggregating. Convolution adds remaining self-loops when this isn't used.
#Create edge index
edge_index = torch.tensor(data = adj_mat).nonzero().t().contiguous()
np.unique(edge_index[0].numpy(), return_counts=True) #these are the in-degrees for each cell
#DataLoader
class Data(torch.utils.data.DataLoader):
def __init__(self, X, edges):
self.X = torch.from_numpy(X)
self.edges = edges
self.len = self.X.shape[0]
def __getitem__(self, index):
return self.X[index]
def __len__(self):
return self.len
cell_graph = Data(X = feature_mat, edges = edge_index)
#instantiate and propagate
conv = TranscriptConv(num_features, num_features)
smoothed_features = conv(cell_graph)
averaged_firstCell = feature_mat[adj_mat[0,:].nonzero()[1],:].mean(axis = 0) #average features along neighborhood of first cell.
#The output of the below should be equal?
print(averaged_firstCell) #first cell mean-aggregated manually
print(smoothed_features[0,:]) #first cell mean-aggregated through GCN's message passing layer
```
