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.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 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(),:].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 ```