Adding negative edges after link neighbor loader

[jaxels20]

I have a hetero graph in PyTorch geometric and want to mini-batch it using LinkNeighborLoader. I know that it can sample and add negative edges for me, but I want to do it manually. How can I, after a batch is created, add negative edges to the batch? I know the global indices of the nodes I want to add a negative edge between. My batches look like this when created:

HeteroData(                                                                                                                                                                                                
  drug={                                                                                                                                                                                                   
    node_id=[1406],                                                                                                                                                                                        
    node_type=[1406],                                                                                                                                                                                      
    x=[1406, 10],                                                                                                                                                                                          
    n_id=[1406]                                                                                                                                                                                            
  },                                                                                                                                                                                                       
  disease={                                                                                                                                                                                                
    node_id=[644],                                                                                                                                                                                         
    node_type=[644],                                                                                                                                                                                       
    x=[644, 10],                                                                                                                                                                                           
    n_id=[644]                                                                                                                                                                                             
  },                                                                                                                                                                                                       
  (drug, may_treat, disease)={                                                                                                                                                                             
    edge_index=[2, 2281],                                                                                                                                                                                  
    edge_label=[128],                                                                                                                                                                                      
    edge_label_index=[2, 128],                                                                                                                                                                             
    e_id=[2281],                                                                                                                                                                                           
    input_id=[64]                                                                                                                                                                                          
  },                                                                                                                                                                                                       
  (disease, rev_may_treat, drug)={                                                                                                                                                                         
    edge_index=[2, 2959],                                                                                                                                                                                  
    e_id=[2959]                                                                                                                                                                                            
  }                                                                                                                                                                                                        
)

Thanks for the help in advance :)
Please just write if you need more information about the problem?

[rusty1s]

You can take a look at the utils.negative_sampling function in PyG. Then, you can do

batch[edge_type_to_add_negative_samples].neg_edge_index = negative_sampling(...)