Added unit tests

Author: ntgbaoo

test.py

@@ -7,6 +7,7 @@
 from tgm.hooks import (
     HistoricalNegativeEdgeSamplerHook,
     HookManager,
+    NodeTypeNegativeSamplerHook,
     RandomNegativeEdgeSamplerHook,
 )
 
@@ -47,6 +48,7 @@ def test_hook_repre():
 def test_hook_reset_state():
     assert RandomNegativeEdgeSamplerHook.has_state == False
     assert HistoricalNegativeEdgeSamplerHook.has_state == True
+    assert NodeTypeNegativeSamplerHook.has_state == True
 
 
 def test_bad_negative_edge_sampler_init():
@@ -198,3 +200,168 @@ def test_hst_sampling(data_test_hst_sampling):
         sampler.reset_state()
         assert sampler._memory is None
         assert sampler._count == 0
+
+
+@pytest.fixture
+def data_test_node_type_sampling():
+    """Fixture for node type sampling tests.
+
+    Graph with 12 nodes (0-11), node types assigned as:
+        type 0: nodes [0, 2, 4, 6, 8, 10]  (even)
+        type 1: nodes [1, 3, 5, 7, 9, 11]  (odd)
+
+    Edges (src, dst):
+        # 1st batch
+        [1, 2],  dst=2  type 0
+        [3, 4],  dst=4  type 0
+        [5, 6],  dst=6  type 0
+        [7, 1],  dst=1  type 1
+        # 2nd batch
+        [2, 3],  dst=3  type 1  -> can sample node 1 (type 1, seen in batch 1)
+        [4, 8],  dst=8  type 0  -> can sample nodes 2, 4, 6 (type 0, seen in batch 1)
+        [6, 5],  dst=5  type 1  -> can sample node 1 (type 1, seen in batch 1)
+        [8, 9],  dst=9  type 1  -> can sample node 1 (type 1, seen in batch 1)
+        # 3rd batch
+        [1, 10], dst=10 type 0  -> can sample nodes 2, 4, 6, 8 (type 0, seen in batches 1&2)
+        [3, 7],  dst=7  type 1  -> can sample nodes 1, 3, 5, 9 (type 1, seen in batches 1&2)
+        [5, 11], dst=11 type 1  -> can sample nodes 1, 3, 5, 9 (type 1, seen in batches 1&2)
+        [7, 0],  dst=0  type 0  -> can sample nodes 2, 4, 6, 8 (type 0, seen in batches 1&2).
+    """
+    edge_index = torch.IntTensor(
+        [
+            # 1st batch
+            [1, 2],
+            [3, 4],
+            [5, 6],
+            [7, 1],
+            # 2nd batch
+            [2, 3],
+            [4, 8],
+            [6, 5],
+            [8, 9],
+            # 3rd batch
+            [1, 10],
+            [3, 7],
+            [5, 11],
+            [7, 0],
+        ]
+    )
+    edge_time = torch.arange(edge_index.size(0))
+    node_type = torch.IntTensor([i % 2 for i in range(12)])  # even=type 0, odd=type 1
+    node_type[-1] = (
+        2  # The last node (11) is of type 2, which is not present in the graph edges
+    )
+    return DGData.from_raw(
+        edge_time=edge_time, edge_index=edge_index, node_type=node_type
+    )
+
+
+def test_node_type_sampling(data_test_node_type_sampling):
+    dg = DGraph(data_test_node_type_sampling)
+
+    hm = HookManager(keys=['unit'])
+    sampler = NodeTypeNegativeSamplerHook(num_nodes=dg.num_nodes)
+
+    hm.register('unit', sampler)
+    loader = DGDataLoader(dg, batch_size=4, hook_manager=hm)
+
+    with hm.activate('unit'):
+        batch_iter = iter(loader)
+
+        # batch 1: no memory yet, all padded
+        batch_1 = next(batch_iter)
+        assert batch_1.neg.shape == (4,)
+        assert torch.equal(
+            batch_1.neg,
+            torch.full((4,), PADDED_NODE_ID, dtype=batch_1.neg.dtype),
+        )
+        assert torch.equal(
+            batch_1.valid_neg_mask,
+            torch.zeros(4, dtype=torch.bool),
+        )
+        # memory initialized with num_nodes entries
+        assert sampler._memory is not None
+        # print()
+        assert sampler._memory.shape == (dg.num_nodes,)
+        # dst nodes 1, 2, 4, 6 recorded in memory
+        assert sampler._memory[2] == 0  # type 0
+        assert sampler._memory[4] == 0  # type 0
+        assert sampler._memory[6] == 0  # type 0
+        assert sampler._memory[1] == 1  # type 1
+
+        # batch 2: memory has nodes 1(t1), 2(t0), 4(t0), 6(t0)
+        batch_2 = next(batch_iter)
+        assert batch_2.neg.shape == (4,)
+        assert torch.equal(
+            batch_2.valid_neg_mask,
+            torch.ones(4, dtype=torch.bool),
+        )
+        # dst=3(t1) -> neg must be type 1 -> only node 1
+        assert batch_2.neg[0].item() == 1
+        # dst=8(t0) -> neg must be type 0 -> one of {2, 4, 6}
+        assert batch_2.neg[1].item() in {2, 4, 6}
+        # dst=5(t1) -> neg must be type 1 -> only node 1
+        assert batch_2.neg[2].item() == 1
+        # dst=9(t1) -> neg must be type 1 -> only node 1
+        assert batch_2.neg[3].item() == 1
+
+        # batch 3: memory now also has nodes 3(t1), 5(t1), 8(t0), 9(t1)
+        batch_3 = next(batch_iter)
+        assert batch_3.neg.shape == (4,)
+        assert torch.equal(
+            batch_3.valid_neg_mask,
+            torch.Tensor([True, True, False, True]),
+        )
+        # dst=10(t0) -> neg must be type 0 -> one of {2, 4, 6, 8}
+        assert batch_3.neg[0].item() in {2, 4, 6, 8}
+        # dst=7(t1)  -> neg must be type 1 -> one of {1, 3, 5, 9}
+        assert batch_3.neg[1].item() in {1, 3, 5, 9}
+        # dst=11(t1) -> neg must be type 2 -> no nodes of type 2 have been seen yet, so neg is padded
+        assert batch_3.neg[2].item() == PADDED_NODE_ID
+        # dst=0(t0)  -> neg must be type 0 -> one of {2, 4, 6, 8}
+        assert batch_3.neg[3].item() in {2, 4, 6, 8}
+
+        sampler.reset_state()
+        assert sampler._memory is None
+
+
+def test_node_type_sampling_no_node_type():
+    """Hook should raise ValueError when dg.node_type is None."""
+    edge_index = torch.IntTensor([[1, 2], [3, 4]])
+    edge_time = torch.arange(2)
+    data = DGData.from_raw(edge_time=edge_time, edge_index=edge_index)
+    dg = DGraph(data)
+
+    sampler = NodeTypeNegativeSamplerHook(num_nodes=dg.num_nodes)
+    with pytest.raises(ValueError, match='dg.node_type is None'):
+        sampler(dg, dg.materialize())
+
+
+def test_node_type_sampling_with_id(data_test_node_type_sampling):
+    """Hook with id should suffix all produced attributes."""
+    dg = DGraph(data_test_node_type_sampling)
+
+    sampler = NodeTypeNegativeSamplerHook(num_nodes=dg.num_nodes, id='foo')
+    batch = sampler(dg, dg.materialize())
+
+    assert hasattr(batch, 'neg_foo')
+    assert hasattr(batch, 'neg_time_foo')
+    assert hasattr(batch, 'valid_neg_mask_foo')
+
+
+def test_node_type_sampling_dependencies():
+    """Check requires/produces sets."""
+    hook = NodeTypeNegativeSamplerHook(num_nodes=10)
+    assert hook.requires == {'edge_src', 'edge_dst', 'edge_time'}
+    assert hook.produces == {'neg', 'neg_time', 'valid_neg_mask'}
+
+    hook_with_id = NodeTypeNegativeSamplerHook(num_nodes=10, id='foo')
+    assert hook_with_id.produces == {'neg_foo', 'neg_time_foo', 'valid_neg_mask_foo'}
+
+
+def test_node_type_sampling_bad_init():
+    """Hook should raise ValueError when num_nodes <= 0."""
+    with pytest.raises(ValueError, match='num_nodes must be a positive integer'):
+        NodeTypeNegativeSamplerHook(num_nodes=0)
+    with pytest.raises(ValueError, match='num_nodes must be a positive integer'):
+        NodeTypeNegativeSamplerHook(num_nodes=-1)

test/unit/test_hooks/test_negative_edge_sampler_hook.py

@@ -7,6 +7,7 @@
     TGBNegativeEdgeSamplerHook,
     TGBTHGNegativeEdgeSamplerHook,
     TGBTKGNegativeEdgeSamplerHook,
+    NodeTypeNegativeSamplerHook,
 )
 from .neighbors import NeighborSamplerHook, RecencyNeighborHook
 from .hook_manager import HookManager

tgm/hooks/__init__.py

@@ -1,4 +1,8 @@
-from .sampler import RandomNegativeEdgeSamplerHook, HistoricalNegativeEdgeSamplerHook
+from .sampler import (
+    RandomNegativeEdgeSamplerHook,
+    HistoricalNegativeEdgeSamplerHook,
+    NodeTypeNegativeSamplerHook,
+)
 from .tgb_sampler import (
     TGBNegativeEdgeSamplerHook,
     TGBTHGNegativeEdgeSamplerHook,

tgm/hooks/negatives/__init__.py

@@ -271,11 +271,16 @@ class NodeTypeNegativeSamplerHook(StatefulHook):
 
     def __init__(
         self,
+        num_nodes: int,
         id: str | None = None,
     ) -> None:
         super().__init__()
+        if num_nodes <= 0:
+            raise ValueError('num_nodes must be a positive integer.')
+
         self._id = id
         self._memory: torch.Tensor | None = None
+        self._num_nodes = num_nodes
 
         self.__post_init__()
 
@@ -288,6 +293,9 @@ def __call__(self, dg: DGraph, batch: DGBatch) -> DGBatch:
                 'dg.node_type is None. NodeTypeNegativeSamplerHook requires node type information in the DGraph.'
             )
         if self._memory is None:
+            logger.debug(
+                f'NodeTypeNegativeSamplerHook: Empty node label memory on first batch. All negatives will be set to PADDED_NODE_ID ({PADDED_NODE_ID}).'
+            )
             neg = torch.full(
                 (batch.edge_dst.size(0),),
                 PADDED_NODE_ID,
@@ -370,8 +378,11 @@ def _update_memory(self, dg: DGraph, batch: DGBatch) -> None:
             has not yet appeared as a destination and is not eligible for sampling.
         """
         if self._memory is None:
+            logger.debug(
+                f'Initializing memory for NodeTypeNegativeSamplerHook with shape ({dg.num_nodes},) on device {dg.device}.'
+            )
             self._memory = torch.full(
-                (dg.num_nodes,),
+                (self._num_nodes,),
                 -1,
                 dtype=torch.int32,
                 device=dg.device,