Modifying edges dynamically in DataBatch

[NajwaLaabid]

I have a DataBatch object made up of several graphs. In one step of my algorithm, I change the edges of the graphs (in their batch form). I would like to get a new object with the modified edge_index/edge_attr, and the full functionality of DataBatch. Here is an example:

graph1 = <DataBatch object of N graphs>
new_edge_index, new_edge_attr = get_new_edges(graph1)
graph2 = <DataBatch object with the same nodes as graph1 and new_edge_index/attr>

Here are a few things I tried:

Solution 1:

graph2 = Data(x=graph1.x, edge_index=new_edge_index, edge_attr=new_edge_attr, batch=graph1.batch)

With this solution, I cannot get a list of the graphs making up the batch via indexing (e.g. graph2[graph2.batch[:3]])

Solution 2:

graph2 = graph1.clone()
graph2.edge_index = new_edge_index
graph2.edge_attr = new_edge_attr

In this case, indexing graph2 by batch index returns the same graphs as graph1 (i.e. the new edges do not reflect on the graphs making up the batch).

I would like to avoid using python lists to keep all my computations on the GPU.

Thanks in advance!

[rusty1s]

There is currently no straightforward way to maintain indexing logic in modified Batch objects, but it is possible by modifying some internal state of Batch, i.e., batch._slice_dict and batch._inc_dict, see here. I suggest you do

batch2 = copy.copy(batch1)
batch2.edge_index = new_edge_index
batch2.edge_attr = new_edge_attr

and modify batch._slice_dict["edge_index"], batch._slice_dict["edge_attr"] and batch._inc_dict["edge_index"] afterwards.

[NajwaLaabid]

Thanks for the prompt reply! I get the following error using your implementation when I try to access graphs from the batch via indexing :
batch2[:10]

/projappl/project_2006950/newretro/lib/python3.10/site-packages/torch_geometric/data/separate.py in _separate(key, value, idx, slices, incs, batch, store, decrement)
     62         key = str(key)
     63         cat_dim = batch.__cat_dim__(key, value, store)
---> 64         start, end = int(slices[idx]), int(slices[idx + 1])
     65         value = value.narrow(cat_dim or 0, start, end - start)
     66         value = value.squeeze(0) if cat_dim is None else value

ValueError: only one element tensors can be converted to Python scalars

[rusty1s]

How did you modify _slice_dict and _inc_dict?

[NajwaLaabid]

Yup that was the mistake. I had to compute the compressed index of the edge_index attribute for _slice_dict["edge_index"] and _slice_dict["edge_attr"].

I am not sure what _inc_dict["edge_index"] contains. The documentation says stores how individual elements need to be incremented, e.g.edge_index is incremented by the cumulated sum of previous elements., which sounds to me like the cumulated sum of edge_index (so same as _slice_dict["edge_index"] ). However, the array seems to contain the cumulated sum of the indices of the nodes.

Here is the implementation that worked for me as of now:

batch2 = copy.deepcopy(batch1) # deep copy to get a new object
batch2.edge_index = new_edge_index
batch2.edge_attr = new_edge_attr

cum_edge_index = get_cum_index(new_edge_index, batch1.ptr) # see implementation below
batch2._inc_dict["edge_index"] = batch1._inc_dict["edge_index"] # because batch1 and batch2 have the same nodes
batch2._slices_dict["edge_index"] = cum_edge_index
batch2._slices_dict["edge_index"] = cum_edge_index 

from itertools import pairwise

def get_cum_index(edge_index, cum_node_cnt):
    total = 0
    s = []
    for i, j in pairwise(cum_node_cnt):
        n =  edge_index[(i <= edge_index) & (edge_index < j)].view(-1, 2)
        total += n.shape[0]
        s.append(n.shape[0])
        
    assert total==edge_index.shape[0], 'the nb of edges per graph does not match the total number of edges in the batch.'
    
    cum_idx = torch.tensor([0]+s).cumsum(dim=0)

    return cum_idx

Does this seem right? Can you clarify the use of _inc_dict?

[rusty1s]

Yes, this looks good. _inc_dict indeed refers to the accumulated amount to which you want to increase edge_index per example. I don't think this needs changes in your case since the number of nodes stays the same.

[NajwaLaabid]

UPDATE: I needed to sort the final edge_idx list before constructing the new batch. The fixed code is marked with #UPDATE.

Unfortunately, my code is still not working. I suspect the edge indexing to still be off somehow. I am posting a full working example in case I missed something else.

For context, I am trying to generate new graphs from an initial graph by 'flipping' some edges (i.e. add new edges and remove old ones) as part of an MCMC sampling algorithm. Since I have more than 10000 graphs in my training data, I would like to perform this sampling on a batch of graphs at once. So the batch version of my algorithm does the following:

• For a given batch of graphs, generate at random an edge for each graph in a vector form (so if I have 64 graphs, I get a vector of 64 edges).
• Add the edges non-existing in the original batch edge_index object, and remove the ones already existing.
• Create a graph batch object with the new graphs.

I am trying to do everything using PyTorch tensors to maximize GPU utilization. Here is a working minimal example:

from itertools import pairwise
import copy
import torch
from torch_geometric.data import Data
from torch_geometric.data import Batch 
from torch_geometric.loader import DataLoader 

# create dummy data
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
edge_attr = torch.tensor([[0, 1, 1, 2]], dtype=torch.long).T
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

graph0 = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
train_loader = DataLoader([graph0]*64, batch_size=64)
init_graph = [b for b in train_loader][0]

# generate a batch example
new_graph = get_sample(copy.deepcopy(init_graph))

# print any graph with negative edges
graphs_with_neg_edges = [g for g in new_graph.to_data_list() \
                  if (g.edge_index<0).any()]

print(f'len(graphs_with_neg_edges) {len(graphs_with_neg_edges)}')
if len(graphs_with_neg_edges)>0: print(f'graphs_with_neg_edges[0].edge_index {graphs_with_neg_edges[0].edge_index}\n')

The above print statement always returns some graphs with random negative edges. Something like:

len(graphs_with_neg_edges) 17
graphs_with_neg_edges[0].edge_index tensor([[  49,   50, -133, -132, -127, -126],
        [  50,   49, -132, -133, -126, -127]])

Here is the implementation:

# helper functions
def get_sample(graph):  
    '''
      Generate new samples from given graph batch.
    '''
    all_edge_idx = graph.edge_index.T.detach().clone() # no need for gradients?
    all_edge_attr = graph.edge_attr.detach().clone()

    # define masks to remove duplicate directional edges
    mask_attr_ = torch.tensor([1, torch.nan]).view(-1, 1) # works because even edge always duplicate
    mask_idx_ = mask_attr_.repeat_interleave(2).view(-1, 2)
    nb_dir_edges = int(all_edge_idx.shape[0]/mask_idx_.shape[0])

    # repeat the mask to be of the size of the tensors
    mask_idx = mask_idx_.repeat(repeats=(nb_dir_edges,1))
    mask_attr = mask_attr_.repeat(repeats=(nb_dir_edges,1))

    # apply mask
    undir_edge_idx = all_edge_idx*mask_idx
    undir_edge_idx = undir_edge_idx[~undir_edge_idx.isnan()].int().reshape(-1, 2)
    undir_edge_attr = all_edge_attr*mask_attr
    undir_edge_attr = undir_edge_attr[~undir_edge_attr.isnan()].int().reshape(-1, 1)
    
    # randomly select edges to transform

    # group the nodes per batch index (i.e. get graph indices from the batch)
    graph_idx, cnt = graph.batch.unique(return_counts=True)
    
    # select edges at random from nodes
    edge_i = torch.randint(low=0, high=cnt.max(), size=cnt.shape)%cnt
    edge_j = torch.randint(low=0, high=cnt.max(), size=cnt.shape)%cnt
    # if a selected edge is a loop, shift its value by -1
    # edge_i[edge_i==edge_j] -= 1 #BUG: fails for edge_i==edge_j==0
    # UPDATE: resample to avoid self-loops
    while (edge_i==edge_j).any():
        edge_j = torch.randint(low=0, high=cnt.max(), size=cnt.shape)%cnt
    # get the cumulated (collated) indices to use on the batch graph
    cs = cnt.cumsum(dim=0).roll(shifts=1)
    cs[0] = 0
    edge_i += cs
    edge_j += cs

    # sort edges to have the smallest first along axis 0 (for easier debugging)
    new_edge_idx = torch.stack((edge_i, edge_j), dim=0).sort(dim=0)[0].T
    
    # randomly select attributes for given edges 
    # Equivalent to flipping: remove edges intersecting with original graph
    # add non intersecting edges with chosen attributes
    id_ = torch.randint(low=0, high=5, size=(new_edge_idx.shape[0],1))
    new_edge_attr = torch.tensor([-2, -1, 1.5, 1, 2.0])[id_]
    
    intersect_idx, toadd_diff_idx = diff_2d_tensor(new_edge_idx, undir_edge_idx)
    intersect_idx2, tokeep_diff_idx = diff_2d_tensor(undir_edge_idx, new_edge_idx)

    tokeep_edges = undir_edge_idx[tokeep_diff_idx]
    tokeep_attr = undir_edge_attr[tokeep_diff_idx]

    toadd_edges = new_edge_idx[toadd_diff_idx]
    toadd_attr = new_edge_attr[toadd_diff_idx]

    edge_idx = torch.cat([tokeep_edges, toadd_edges], dim=0)
    edge_attr = torch.cat([tokeep_attr, toadd_attr], dim=0)
    
    # UPDATE: graph batch expects a sorted edge_list
    edge_idx = sort_row(edge_idx)

    edge_idx_flipped = edge_idx.flip(dims=(1,))
    edge_idx_all = torch.stack((edge_idx, edge_idx_flipped), dim=2).view(-1, 2)
    edge_attr_all = edge_attr.repeat_interleave(repeats=2).view(-1, 1)
    
    total = 0
    s = []
    for i, j in pairwise(graph.ptr):
        n =  edge_idx_all[(i <= edge_idx_all) & (edge_idx_all < j)].view(-1, 2) # explore why count is off
        total += n.shape[0]
        s.append(n.shape[0])

    assert total==edge_idx_all.shape[0], 'The sum of graph edges does not match the total number of edges.'

    dict_idx = torch.tensor([0]+s).cumsum(dim=0)
    
    new_graph = copy.deepcopy(graph)
    new_graph.edge_index = edge_idx_all.T
    new_graph.edge_attr = edge_attr_all
    
    new_graph._slice_dict = copy.deepcopy(new_graph._slice_dict)
    new_graph._inc_dict = copy.deepcopy(new_graph._inc_dict)
    
    new_graph._slice_dict["edge_index"] = dict_idx.detach().clone()
    new_graph._slice_dict["edge_attr"] = dict_idx.detach().clone()
    new_graph._inc_dict["edge_index"] = graph._inc_dict["edge_index"].detach().clone()

    return new_graph


def diff_2d_tensor(tsr1, tsr2):
   '''
      Return the intersection and difference between 2 2d tensors.
   '''
    values, intersect_idx = torch.topk(((tsr1.t()==tsr2.unsqueeze(-1))\
                                      .all(dim=1)).int(), 1, 1)
    intersect_idx = intersect_idx[values!=0]
    
    diff_idx_mask = torch.ones(tsr1.shape[0], dtype=torch.bool)
    diff_idx_mask[intersect_idx] = False

    intersect_edges_ = tsr1[intersect_idx]
    diff_idx_edges_ = tsr1[diff_idx_mask] # edges to keep
    
    return intersect_idx.long(), diff_idx_mask

def sort_row(tsr):
    '''
        Sort a 2d tensor A s.t. for i<j, A[i,0]<A[j,0] or (A[i,0]==A[j,0] and A[i,1]<A[j,1]).
        Assumes all the values in the tensor are positive, with a known maximum value.
    '''
    M = tsr.max()
    augmented_ = tsr.select(1, 0) * M + tsr.select(1, 1)
    ind = augmented_.sort().indices
    res = tsr.index_select(0, ind)
    
    return res

[rusty1s]

The issue seems to be that the final edge_index in new_graph is not sorted, while Batch expects edge_index to be sorted such that edge indices for a single example are contiguous in memory.

[NajwaLaabid]

Yes exactly, I figured it out yesterday :) There is also another bug in this line:
'''
edge_i[edge_i==edge_j] -= 1
'''
It leads to a negative edge value if edge_i==edge_j==0. I will update my answer with the working code now. Thanks a lot for your time!