Added batching in transductive setting

configs/dataset/graph/US-county-demos.yaml

@@ -30,6 +30,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 0
   pin_memory: False

configs/dataset/graph/amazon_ratings.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 0
   pin_memory: False

configs/dataset/graph/cocitation_citeseer.yaml

@@ -28,6 +28,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/graph/cocitation_cora.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/graph/cocitation_pubmed.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/graph/manual_dataset.yaml

@@ -28,6 +28,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 
+  batch_size: -1 
   num_workers: 1
   pin_memory: False

configs/dataset/graph/minesweeper.yaml

@@ -28,6 +28,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 0
   pin_memory: False

configs/dataset/graph/questions.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/graph/roman_empire.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 0
   pin_memory: False

configs/dataset/graph/tolokers.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/hypergraph/coauthorship_cora.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/hypergraph/coauthorship_dblp.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/hypergraph/cocitation_citeseer.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/hypergraph/cocitation_cora.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

configs/dataset/hypergraph/cocitation_pubmed.yaml

@@ -27,6 +27,6 @@ split_params:
 
 # Dataloader parameters
 dataloader_params:
-  batch_size: 1 # Fixed
+  batch_size: -1 # Fixed
   num_workers: 1
   pin_memory: False

test/data/batching/test_neighbor_cells_loader.py

@@ -0,0 +1,132 @@
+""" Test for the NeighborCellsLoader class."""
+import os
+import shutil
+import rootutils
+from hydra import compose
+import torch
+
+from topobenchmark.data.preprocessor import PreProcessor
+from topobenchmark.data.utils.utils import load_manual_graph
+from topobenchmark.data.batching import NeighborCellsLoader
+from topobenchmark.run import initialize_hydra
+
+initialize_hydra()
+
+path = "./graph2simplicial_lifting/"
+if os.path.isdir(path):
+    shutil.rmtree(path)
+cfg = compose(config_name="run.yaml", 
+              overrides=["dataset=graph/manual_dataset", "model=simplicial/san"], 
+              return_hydra_config=True)
+
+data = load_manual_graph()
+preprocessed_dataset = PreProcessor(data, path, cfg['transforms'])
+data = preprocessed_dataset[0]
+
+batch_size=2
+
+rank = 0
+n_cells = data[f'x_{rank}'].shape[0]
+train_prop = 0.5
+n_train = int(train_prop * n_cells)
+train_mask = torch.zeros(n_cells, dtype=torch.bool)
+train_mask[:n_train] = 1
+
+y = torch.zeros(n_cells, dtype=torch.long)
+data.y = y
+
+loader = NeighborCellsLoader(data,
+                             rank=rank,
+                             num_neighbors=[-1],
+                             input_nodes=train_mask,
+                             batch_size=batch_size,
+                             shuffle=False)
+train_nodes = []
+for batch in loader:
+    train_nodes += [n for n in batch.n_id[:batch_size]]
+for i in range(n_train):
+    assert i in train_nodes
+
+rank = 1
+n_cells = data[f'x_{rank}'].shape[0]
+train_prop = 0.5
+n_train = int(train_prop * n_cells)
+train_mask = torch.zeros(n_cells, dtype=torch.bool)
+train_mask[:n_train] = 1
+
+y = torch.zeros(n_cells, dtype=torch.long)
+data.y = y
+
+loader = NeighborCellsLoader(data,
+                             rank=rank,
+                             num_neighbors=[-1,-1],
+                             input_nodes=train_mask,
+                             batch_size=batch_size,
+                             shuffle=False)
+
+train_nodes = []
+for batch in loader:
+    train_nodes += [n for n in batch.n_id[:batch_size]]
+for i in range(n_train):
+    assert i in train_nodes
+shutil.rmtree(path)
+
+
+path = "./graph2hypergraph_lifting/"
+if os.path.isdir(path):
+    shutil.rmtree(path)
+cfg = compose(config_name="run.yaml", 
+              overrides=["dataset=graph/manual_dataset", "model=hypergraph/allsettransformer"], 
+              return_hydra_config=True)
+
+data = load_manual_graph()
+preprocessed_dataset = PreProcessor(data, path, cfg['transforms'])
+data = preprocessed_dataset[0]
+
+batch_size=2
+
+rank = 0
+n_cells = data[f'x_0'].shape[0]
+train_prop = 0.5
+n_train = int(train_prop * n_cells)
+train_mask = torch.zeros(n_cells, dtype=torch.bool)
+train_mask[:n_train] = 1
+
+y = torch.zeros(n_cells, dtype=torch.long)
+data.y = y
+
+loader = NeighborCellsLoader(data,
+                             rank=rank,
+                             num_neighbors=[-1],
+                             input_nodes=train_mask,
+                             batch_size=batch_size,
+                             shuffle=False)
+train_nodes = []
+for batch in loader:
+    train_nodes += [n for n in batch.n_id[:batch_size]]
+for i in range(n_train):
+    assert i in train_nodes
+
+rank = 1
+n_cells = data[f'x_hyperedges'].shape[0]
+train_prop = 0.5
+n_train = int(train_prop * n_cells)
+train_mask = torch.zeros(n_cells, dtype=torch.bool)
+train_mask[:n_train] = 1
+
+y = torch.zeros(n_cells, dtype=torch.long)
+data.y = y
+
+loader = NeighborCellsLoader(data,
+                             rank=rank,
+                             num_neighbors=[-1,-1],
+                             input_nodes=train_mask,
+                             batch_size=batch_size,
+                             shuffle=False)
+
+train_nodes = []
+for batch in loader:
+    train_nodes += [n for n in batch.n_id[:batch_size]]
+for i in range(n_train):
+    assert i in train_nodes
+shutil.rmtree(path)

test/data/dataload/test_Dataloaders.py

@@ -20,10 +20,6 @@ class TestCollateFunction:
 
     def setup_method(self):
         """Setup the test."""
-
-        hydra.initialize(
-        version_base="1.3", config_path="../../../configs", job_name="run"
-        )
         cfg = hydra.compose(config_name="run.yaml", overrides=["dataset=graph/NCI1"])
 
         graph_loader = hydra.utils.instantiate(cfg.dataset.loader, _recursive_=False)

topobenchmark/data/batching/__init__.py

@@ -0,0 +1,7 @@
+"""Init file for batching module."""
+
+from .neighbor_cells_loader import NeighborCellsLoader
+
+__all__ = [
+    "NeighborCellsLoader",
+]