k3-node 1.0.0__py3-none-any.whl

k3_node/loader/temporal_dataloader.py

@@ -0,0 +1,98 @@
+from typing import List
+
+import numpy as np
+
+try:
+    import torch
+    import torch.utils.data
+    BaseDataLoader = torch.utils.data.DataLoader
+except ImportError:
+    torch = None
+    BaseDataLoader = object
+
+from k3_node.data import TemporalData
+
+
+class TemporalDataLoader(BaseDataLoader):
+    r"""A data loader which merges successive events of a
+    :class:`k3_node.data.TemporalData` to a mini-batch.
+
+    Args:
+        data (TemporalData): The :obj:`~k3_node.data.TemporalData` from which to load.
+        batch_size (int, optional): How many samples per batch to load. (default: :obj:`1`)
+        neg_sampling_ratio (float, optional): The ratio of sampled negative
+            destination nodes to the number of positive destination nodes. (default: :obj:`0.0`)
+        **kwargs (optional): Additional arguments of :class:`torch.utils.data.DataLoader`.
+    """
+    def __init__(
+        self,
+        data: TemporalData,
+        batch_size: int = 1,
+        neg_sampling_ratio: float = 0.0,
+        **kwargs,
+    ):
+        kwargs.pop('dataset', None)
+        kwargs.pop('collate_fn', None)
+        kwargs.pop('shuffle', None)
+
+        self.data = data
+        self.events_per_batch = batch_size
+        self.neg_sampling_ratio = neg_sampling_ratio
+
+        if neg_sampling_ratio > 0:
+            dst = data.dst
+            if torch is not None and isinstance(dst, torch.Tensor):
+                self.min_dst = int(dst.min())
+                self.max_dst = int(dst.max())
+            else:
+                self.min_dst = int(np.min(np.asarray(dst)))
+                self.max_dst = int(np.max(np.asarray(dst)))
+
+        if kwargs.get('drop_last', False) and len(data) % batch_size != 0:
+            arange = list(range(0, len(data) - batch_size, batch_size))
+        else:
+            arange = list(range(0, len(data), batch_size))
+
+        if torch is not None:
+            super().__init__(arange, 1, shuffle=False, collate_fn=self, **kwargs)
+        else:
+            self.dataset = arange
+            self.batch_size = 1
+            self.shuffle = False
+            self.collate_fn = self
+
+    def __call__(self, arange: List[int]) -> TemporalData:
+        start = arange[0]
+        end = start + self.events_per_batch
+        batch = self.data[start:end]
+
+        is_torch = torch is not None and isinstance(batch.dst, torch.Tensor)
+
+        n_ids = [batch.src, batch.dst]
+
+        if self.neg_sampling_ratio > 0:
+            num_neg = round(self.neg_sampling_ratio * (batch.dst.size(0) if is_torch else len(batch.dst)))
+            if is_torch:
+                batch.neg_dst = torch.randint(
+                    low=self.min_dst,
+                    high=self.max_dst + 1,
+                    size=(num_neg,),
+                    dtype=batch.dst.dtype,
+                    device=batch.dst.device,
+                )
+            else:
+                batch.neg_dst = np.random.randint(
+                    low=self.min_dst,
+                    high=self.max_dst + 1,
+                    size=(num_neg,),
+                    dtype=np.int64,
+                )
+            n_ids.append(batch.neg_dst)
+
+        if is_torch:
+            batch.n_id = torch.cat(n_ids, dim=0).unique()
+        else:
+            batch.n_id = np.unique(np.concatenate([np.asarray(x) for x in n_ids], axis=0))
+
+        return batch
+

k3_node/loader/test_dataloader.py

@@ -0,0 +1,113 @@
+import numpy as np
+import pytest
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+from k3_node.data import Data, HeteroData, TemporalData
+from k3_node.loader import (
+    DataListLoader,
+    DataLoader,
+    DenseDataLoader,
+    RandomNodeLoader,
+    TemporalDataLoader,
+)
+
+
+def create_dummy_data(num_nodes=4, num_features=8):
+    x = np.random.randn(num_nodes, num_features).astype(np.float32)
+    edge_index = np.array([[0, 1, 2, 3], [1, 2, 3, 0]], dtype=np.int64)
+    y = np.array([0, 1, 0, 1], dtype=np.int64)
+    if torch is not None:
+        x = torch.from_numpy(x)
+        edge_index = torch.from_numpy(edge_index)
+        y = torch.from_numpy(y)
+    return Data(x=x, edge_index=edge_index, y=y)
+
+
+def test_data_loader_basic():
+    dataset = [create_dummy_data(num_nodes=i + 2) for i in range(4)]
+    loader = DataLoader(dataset, batch_size=2, shuffle=False)
+
+    batches = list(loader)
+    assert len(batches) == 2
+
+    batch0 = batches[0]
+    assert batch0.num_graphs == 2
+    assert hasattr(batch0, 'batch')
+    assert hasattr(batch0, 'ptr')
+    assert batch0.num_nodes == dataset[0].num_nodes + dataset[1].num_nodes
+
+
+def test_data_loader_follow_batch_and_exclude():
+    dataset = [create_dummy_data(num_nodes=3) for _ in range(3)]
+    loader = DataLoader(dataset, batch_size=2, follow_batch=['y'], exclude_keys=['edge_index'])
+
+    batch = next(iter(loader))
+    assert 'edge_index' not in batch
+    assert hasattr(batch, 'y_batch')
+
+
+def test_data_list_loader():
+    dataset = [create_dummy_data(num_nodes=3) for _ in range(4)]
+    loader = DataListLoader(dataset, batch_size=2, shuffle=False)
+
+    batches = list(loader)
+    assert len(batches) == 2
+    assert isinstance(batches[0], list)
+    assert len(batches[0]) == 2
+    assert isinstance(batches[0][0], Data)
+
+
+def test_dense_data_loader():
+    def create_dense_graph(num_nodes=4, num_features=6):
+        x = np.random.randn(num_nodes, num_features).astype(np.float32)
+        adj = np.random.randn(num_nodes, num_nodes).astype(np.float32)
+        if torch is not None:
+            x = torch.from_numpy(x)
+            adj = torch.from_numpy(adj)
+        return Data(x=x, adj=adj)
+
+    dataset = [create_dense_graph() for _ in range(4)]
+    loader = DenseDataLoader(dataset, batch_size=2, shuffle=False)
+
+    batch = next(iter(loader))
+    assert batch.x.shape == (2, 4, 6)
+    assert batch.adj.shape == (2, 4, 4)
+
+
+def test_temporal_data_loader():
+    src = np.array([0, 1, 0, 2, 1, 3], dtype=np.int64)
+    dst = np.array([1, 2, 2, 3, 3, 0], dtype=np.int64)
+    t = np.array([1, 2, 3, 4, 5, 6], dtype=np.int64)
+    msg = np.random.randn(6, 4).astype(np.float32)
+
+    if torch is not None:
+        src = torch.from_numpy(src)
+        dst = torch.from_numpy(dst)
+        t = torch.from_numpy(t)
+        msg = torch.from_numpy(msg)
+
+    data = TemporalData(src=src, dst=dst, t=t, msg=msg)
+    loader = TemporalDataLoader(data, batch_size=3, neg_sampling_ratio=1.0)
+
+    batches = list(loader)
+    assert len(batches) == 2
+    assert len(batches[0]) == 3
+    assert hasattr(batches[0], 'neg_dst')
+    assert hasattr(batches[0], 'n_id')
+
+
+def test_random_node_loader():
+    data = create_dummy_data(num_nodes=10)
+    loader = RandomNodeLoader(data, num_parts=2)
+
+    parts = list(loader)
+    assert len(parts) == 2
+    for part in parts:
+        assert isinstance(part, Data)
+        assert part.num_nodes <= 10
+        assert hasattr(part, 'x')
+

k3_node/loader/test_keras_dataset.py

@@ -0,0 +1,221 @@
+import numpy as np
+import keras
+from keras import ops
+
+from k3_node.data import Data
+from k3_node.layers import GCNConv
+from k3_node.loader import DataLoader, FullGraphDataset, NeighborLoader
+from k3_node.loader.keras_dataset import to_keras_batch
+
+
+def _graph(num_nodes=30, num_classes=3):
+    rng = np.random.default_rng(0)
+    y = rng.integers(0, num_classes, num_nodes)
+    x = (np.eye(num_classes)[y] + 0.3 * rng.standard_normal((num_nodes, num_classes))).astype("float32")
+    edge_index = rng.integers(0, num_nodes, (2, 90)).astype("int32")
+    train_mask = np.zeros(num_nodes, dtype=bool)
+    train_mask[:10] = True
+    return Data(x=x, edge_index=edge_index, y=y.astype("int32"), train_mask=train_mask, test_mask=~train_mask)
+
+
+class SmallGCN(keras.Model):
+    def __init__(self):
+        super().__init__()
+        self.conv = GCNConv(3, 3)
+
+    def call(self, data):
+        return self.conv(data.x, data.edge_index)
+
+
+def test_fit_and_evaluate_on_full_graph():
+    data = _graph()
+    model = SmallGCN()
+    model.compile(keras.optimizers.Adam(0.05), keras.losses.SparseCategoricalCrossentropy(from_logits=True))
+    history = model.fit(FullGraphDataset(data, mask="train_mask"), epochs=30, verbose=0)
+    assert history.history["loss"][-1] < history.history["loss"][0]
+    model.evaluate(FullGraphDataset(data, mask="test_mask"), verbose=0)
+    preds = model.predict(FullGraphDataset(data), verbose=0)
+    assert preds.shape == (30, 3)
+
+
+def test_normalized_mask_gives_mean_over_masked_nodes():
+    data = _graph()
+    model = SmallGCN()
+    loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=None)
+    model.compile(loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True))
+    reported = model.evaluate(FullGraphDataset(data, mask="train_mask"), verbose=0)
+    per_node = ops.convert_to_numpy(loss_fn(data.y, model(to_keras_batch(data)[0])))
+    np.testing.assert_allclose(reported, per_node[: 10].mean(), rtol=1e-5)
+
+
+def test_graph_loader_works_with_fit_evaluate_predict():
+    from k3_node.datasets import FakeDataset
+    from k3_node.layers import GCNConv, global_mean_pool
+
+    class GraphClassifier(keras.Model):
+        def __init__(self):
+            super().__init__()
+            self.conv = GCNConv(8, 16)
+            self.head = keras.layers.Dense(3)
+
+        def call(self, data):
+            x = self.conv(data.x, data.edge_index)
+            return self.head(global_mean_pool(x, data.batch, data.num_graphs))
+
+    dataset = FakeDataset(num_graphs=20, avg_num_nodes=8, num_channels=8, num_classes=3)
+    loader = DataLoader(dataset, batch_size=6, shuffle=True)
+    model = GraphClassifier()
+    model.compile(optimizer="adam", loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=["accuracy"])
+    history = model.fit(loader, epochs=2, verbose=0)  # compiled on TF/JAX: num_graphs must be static
+    assert len(history.history["loss"]) == 2
+    model.evaluate(DataLoader(dataset, batch_size=6), verbose=0)
+    assert model.predict(DataLoader(dataset, batch_size=6), verbose=0).shape == (20, 3)
+    # Iterating the loader directly still yields Batch objects
+    batch = next(iter(loader))
+    assert hasattr(batch, "edge_index") and batch.num_graphs == 6
+
+
+def test_neighbor_loader_weights_seed_nodes_only():
+    data = _graph()
+    loader = NeighborLoader(data, num_neighbors=[3], batch_size=4, input_nodes=np.arange(10))
+    inputs, y, weight = loader[0]
+    assert y.shape[0] == inputs.x.shape[0]
+    assert (weight[:4] > 0).all() and (weight[4:] == 0).all()
+
+
+def test_full_graph_dataset_index_split():
+    import numpy as np
+    from k3_node.data import Data
+    from k3_node.loader import FullGraphDataset
+
+    data = Data(x=np.ones((5, 2), "float32"), edge_index=np.array([[0, 1], [1, 2]]),
+                train_idx=np.array([3, 1]), train_y=np.array([2, 1]), num_nodes=5)
+    inputs, y, weight = FullGraphDataset(data, index="train_idx", target="train_y")[0]
+    assert "train_idx" not in inputs._fields and "train_y" not in inputs._fields
+    np.testing.assert_array_equal(y, [0, 1, 0, 2, 0])
+    np.testing.assert_allclose(weight, [0, 2.5, 0, 2.5, 0])  # mean over the 2 indexed nodes
+
+
+def test_full_graph_dataset_fresh_negatives_every_epoch():
+    import numpy as np
+    from k3_node.data import Data
+    from k3_node.loader import FullGraphDataset
+
+    edge_index = np.array([[0, 1, 2, 3], [1, 2, 3, 4]])
+    data = Data(x=np.ones((50, 2), "float32"), edge_index=edge_index, edge_label_index=edge_index,
+                edge_label=np.ones(4, "float32"), num_nodes=50)
+    dataset = FullGraphDataset(data, neg_sampling_ratio=2.0)
+    (inputs1, y1), (inputs2, _) = dataset[0], dataset[0]
+    assert inputs1.edge_label_index.shape == (2, 12)
+    np.testing.assert_array_equal(y1, [1] * 4 + [0] * 8)
+    assert not np.array_equal(inputs1.edge_label_index[:, 4:], inputs2.edge_label_index[:, 4:])
+
+
+def test_link_neighbor_loader_uses_local_ids():
+    import numpy as np
+    from k3_node.data import Data
+    from k3_node.loader import LinkNeighborLoader
+
+    rng = np.random.default_rng(0)
+    edge_index = rng.integers(0, 200, size=(2, 600))
+    data = Data(x=np.arange(200, dtype="float32")[:, None], edge_index=edge_index, num_nodes=200)
+    loader = LinkNeighborLoader(data, num_neighbors=[5], batch_size=16, neg_sampling_ratio=1.0)
+    batch = next(iter(loader))
+    eli = np.asarray(batch.edge_label_index)
+    n_id = np.asarray(batch.n_id)
+    assert eli.max() < n_id.shape[0]  # local ids into the sampled subgraph
+    label = np.asarray(batch.edge_label)
+    pos = n_id[eli[:, label == 1]]  # back to global ids: must be real edges
+    real = set(map(tuple, edge_index.T.tolist()))
+    assert all(tuple(e) in real for e in pos.T.tolist())
+
+
+def test_loader_with_mask():
+    import numpy as np
+    from k3_node.data import Data
+    from k3_node.loader import ClusterData, ClusterLoader
+
+    rng = np.random.default_rng(0)
+    data = Data(x=rng.random((60, 3)).astype("float32"), edge_index=rng.integers(0, 60, (2, 200)),
+                y=rng.integers(0, 3, 60), train_mask=np.arange(60) < 30, num_nodes=60)
+    loader = ClusterLoader(ClusterData(data, num_parts=4), batch_size=2).with_mask("train_mask")
+    inputs, y, weight = loader[0]
+    assert weight.shape == y.shape and "train_mask" not in inputs._fields
+
+
+def test_shadow_roots_and_labels():
+    import numpy as np
+    from k3_node.data import Data
+    from k3_node.loader import ShaDowKHopSampler
+
+    rng = np.random.default_rng(0)
+    data = Data(x=np.arange(50, dtype="float32")[:, None], edge_index=rng.integers(0, 50, (2, 200)),
+                y=np.arange(50), num_nodes=50)
+    loader = ShaDowKHopSampler(data, depth=2, num_neighbors=3, node_idx=np.arange(10, 20), batch_size=5)
+    batch = next(iter(loader))
+    root_x = np.asarray(batch.x)[np.asarray(batch.root_n_id), 0]
+    np.testing.assert_array_equal(root_x, np.arange(10, 15))  # the roots are the seed nodes
+    np.testing.assert_array_equal(np.asarray(batch.y), np.arange(10, 15))  # one label per subgraph
+
+
+def test_graph_saint_samplers():
+    import numpy as np
+    from k3_node.data import Data
+    from k3_node.loader import GraphSAINTEdgeSampler, GraphSAINTNodeSampler, GraphSAINTRandomWalkSampler
+
+    rng = np.random.default_rng(0)
+    edge_index = rng.integers(0, 100, (2, 400))
+    data = Data(x=np.arange(100, dtype="float32")[:, None], edge_index=edge_index,
+                edge_attr=np.arange(400, dtype="float32"), y=np.arange(100), num_nodes=100)
+    for loader in [GraphSAINTNodeSampler(data, batch_size=30, num_steps=3, sample_coverage=5),
+                   GraphSAINTEdgeSampler(data, batch_size=20, num_steps=3, sample_coverage=5),
+                   GraphSAINTRandomWalkSampler(data, batch_size=10, walk_length=2, num_steps=3, sample_coverage=5)]:
+        batches = list(loader)
+        assert len(batches) == 3
+        batch = batches[0]
+        x = np.asarray(batch.x)[:, 0].astype(int)
+        ei = np.asarray(batch.edge_index)
+        # every subgraph edge is a real edge between sampled nodes, carrying its own attribute
+        real = {tuple(e): i for i, e in enumerate(edge_index.T.tolist())}
+        for (a, b), attr in zip(ei.T.tolist(), np.asarray(batch.edge_attr).astype(int)):
+            assert (x[a], x[b]) in real and tuple(edge_index[:, attr]) == (x[a], x[b])
+        assert np.asarray(batch.node_norm).shape == (batch.num_nodes,)
+        assert np.asarray(batch.edge_norm).shape == (ei.shape[1],)
+        inputs, y = loader[0]  # Keras batch
+        assert inputs.x.shape[0] == y.shape[0]
+
+
+def test_neighbor_loader_disjoint():
+    import numpy as np
+    from k3_node.data import Data
+    from k3_node.loader import NeighborLoader
+
+    rng = np.random.default_rng(0)
+    edge_index = rng.integers(0, 30, (2, 150))
+    data = Data(x=np.arange(30, dtype="float32")[:, None], edge_index=edge_index, num_nodes=30)
+    batch = next(iter(NeighborLoader(data, num_neighbors=[3, 2], batch_size=8, disjoint=True)))
+    b, n_id, ei = np.asarray(batch.batch), np.asarray(batch.n_id), np.asarray(batch.edge_index)
+    np.testing.assert_array_equal(b[:8], np.arange(8))  # seeds first, one subgraph each
+    assert np.all(b[ei[0]] == b[ei[1]])  # edges never cross subgraphs
+    real = set(map(tuple, edge_index.T.tolist()))
+    assert all((n_id[s], n_id[t]) in real for s, t in ei.T.tolist())
+    for g in range(8):  # no node appears twice within one subgraph
+        assert len(set(n_id[b == g].tolist())) == int((b == g).sum())
+
+
+def test_full_graph_dataset_hetero():
+    import numpy as np
+    from k3_node.data import HeteroData
+    from k3_node.loader import FullGraphDataset
+
+    data = HeteroData()
+    data["user"].x = np.ones((4, 2), "float32")
+    data["user"].y = np.array([0, 1, 0, 1])
+    data["user"].train_mask = np.array([True, True, False, False])
+    data["item"].x = np.ones((3, 5), "float32")
+    data["user", "buys", "item"].edge_index = np.array([[0, 1, 3], [0, 2, 1]])
+    inputs, y, weight = FullGraphDataset(data, node_type="user", mask="train_mask")[0]
+    assert set(inputs.x_dict) == {"user", "item"}
+    assert ("user", "buys", "item") in inputs.edge_index_dict
+    np.testing.assert_array_equal(y, [0, 1, 0, 1])
+    np.testing.assert_allclose(weight, [2, 2, 0, 0])

k3_node/loader/test_neighbor_loader.py

@@ -0,0 +1,122 @@
+import numpy as np
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+from k3_node.data import Data, HeteroData
+from k3_node.loader import HGTLoader, LinkNeighborLoader, NeighborLoader
+
+
+def get_homo_graph():
+    # 6 nodes connected in a ring: 0->1->2->3->4->5->0 and some cross edges
+    edge_index = np.array([
+        [0, 1, 2, 3, 4, 5, 0, 2],
+        [1, 2, 3, 4, 5, 0, 3, 5],
+    ], dtype=np.int64)
+    x = np.random.randn(6, 16).astype(np.float32)
+    y = np.array([0, 1, 0, 1, 0, 1], dtype=np.int64)
+
+    if torch is not None:
+        edge_index = torch.from_numpy(edge_index)
+        x = torch.from_numpy(x)
+        y = torch.from_numpy(y)
+
+    return Data(x=x, edge_index=edge_index, y=y)
+
+
+def get_hetero_graph():
+    data = HeteroData()
+    data['paper'].x = np.random.randn(10, 8).astype(np.float32)
+    data['author'].x = np.random.randn(5, 8).astype(np.float32)
+
+    data['author', 'writes', 'paper'].edge_index = np.array([
+        [0, 1, 2, 3, 4, 0, 1],
+        [0, 1, 2, 3, 4, 5, 6],
+    ], dtype=np.int64)
+    data['paper', 'cites', 'paper'].edge_index = np.array([
+        [0, 1, 2, 3],
+        [1, 2, 3, 4],
+    ], dtype=np.int64)
+
+    if torch is not None:
+        data['paper'].x = torch.from_numpy(data['paper'].x)
+        data['author'].x = torch.from_numpy(data['author'].x)
+        data['author', 'writes', 'paper'].edge_index = torch.from_numpy(data['author', 'writes', 'paper'].edge_index)
+        data['paper', 'cites', 'paper'].edge_index = torch.from_numpy(data['paper', 'cites', 'paper'].edge_index)
+
+    return data
+
+
+def test_homo_neighbor_loader():
+    data = get_homo_graph()
+    input_nodes = [0, 1]
+    if torch is not None:
+        input_nodes = torch.tensor(input_nodes, dtype=torch.long)
+
+    loader = NeighborLoader(
+        data,
+        num_neighbors=[2, 2],
+        batch_size=2,
+        input_nodes=input_nodes,
+        shuffle=False,
+    )
+
+    batch = next(iter(loader))
+    assert batch.batch_size == 2
+    assert hasattr(batch, 'n_id')
+    assert hasattr(batch, 'e_id')
+    assert hasattr(batch, 'num_sampled_nodes')
+    assert hasattr(batch, 'num_sampled_edges')
+    assert batch.x.shape[0] == len(batch.n_id)
+    assert batch.edge_index.shape[0] == 2
+
+
+def test_hetero_neighbor_loader():
+    data = get_hetero_graph()
+    loader = NeighborLoader(
+        data,
+        num_neighbors=[2, 2],
+        batch_size=2,
+        input_nodes=('paper', [0, 1]),
+        shuffle=False,
+    )
+
+    batch = next(iter(loader))
+    assert batch['paper'].batch_size == 2
+    assert hasattr(batch['paper'], 'n_id')
+    assert hasattr(batch['author', 'writes', 'paper'], 'edge_index')
+
+
+def test_link_neighbor_loader():
+    data = get_homo_graph()
+    loader = LinkNeighborLoader(
+        data,
+        num_neighbors=[2, 2],
+        batch_size=2,
+        neg_sampling_ratio=1.0,
+        shuffle=False,
+    )
+
+    batch = next(iter(loader))
+    assert hasattr(batch, 'edge_label_index')
+    assert hasattr(batch, 'edge_label')
+    assert hasattr(batch, 'n_id')
+    assert batch.edge_label.shape[0] == 4  # 2 positive + 2 negative
+
+
+def test_hgt_loader():
+    data = get_hetero_graph()
+    loader = HGTLoader(
+        data,
+        num_samples=[4, 4],
+        input_nodes=('paper', [0, 1]),
+        batch_size=2,
+        shuffle=False,
+    )
+
+    batch = next(iter(loader))
+    assert batch['paper'].batch_size == 2
+    assert hasattr(batch['paper'], 'n_id')
+

k3_node/loader/test_sampler_utils.py

@@ -0,0 +1,82 @@
+import numpy as np
+import pytest
+
+from k3_node.loader.sampler_utils import FastGraph, sample_neighbors_homo
+
+
+def _reference_full_sampling(edge_index, seeds, num_hops, subgraph_type):
+    """The original per-node Python algorithm, restricted to k=-1 where it is deterministic."""
+    graph = FastGraph(edge_index)
+    nodes, visited = [], {}
+    for s in seeds:
+        if int(s) not in visited:
+            visited[int(s)] = len(nodes)
+            nodes.append(int(s))
+    frontier, sampled = list(nodes), []
+    for _ in range(num_hops):
+        next_frontier = []
+        for target in frontier:
+            srcs, e_ids = graph.get_neighbors(target)
+            for s, e in zip(srcs, e_ids):
+                sampled.append((int(s), target, int(e)))
+                if int(s) not in visited:
+                    visited[int(s)] = len(nodes)
+                    nodes.append(int(s))
+                    next_frontier.append(int(s))
+        frontier = next_frontier
+    if subgraph_type == "induced":
+        edges = [(visited[u], visited[v], e) for e, (u, v) in enumerate(edge_index.T) if u in visited and v in visited]
+    elif subgraph_type == "bidirectional":
+        edges = [x for u, v, e in sampled for x in ((visited[u], visited[v], e), (visited[v], visited[u], e))]
+    else:
+        edges = [(visited[u], visited[v], e) for u, v, e in sampled]
+    edges = np.array(edges, dtype=np.int64).reshape(-1, 3)
+    return np.array(nodes), edges[:, 0], edges[:, 1], edges[:, 2]
+
+
+def _random_graph(num_nodes=40, num_edges=160, seed=0):
+    rng = np.random.default_rng(seed)
+    return rng.integers(0, num_nodes, (2, num_edges)).astype(np.int64)
+
+
+@pytest.mark.parametrize("subgraph_type", ["directional", "bidirectional", "induced"])
+def test_full_neighborhood_matches_reference(subgraph_type):
+    edge_index = _random_graph()
+    seeds = np.array([3, 7, 3, 11])
+    expected = _reference_full_sampling(edge_index, seeds, 2, subgraph_type)
+    nodes, row, col, edge, _, _ = sample_neighbors_homo(edge_index, seeds, [-1, -1], subgraph_type=subgraph_type)
+    for got, want in zip((nodes, row, col, edge), expected):
+        np.testing.assert_array_equal(got, want)
+
+
+@pytest.mark.parametrize("replace", [False, True])
+def test_sampled_neighborhood_properties(replace):
+    edge_index = _random_graph()
+    graph = FastGraph(edge_index)
+    seeds = np.array([0, 5, 9])
+    nodes, row, col, edge, n_counts, e_counts = sample_neighbors_homo(
+        edge_index, seeds, [3, 2], replace=replace, graph=graph
+    )
+    assert len(np.unique(nodes)) == len(nodes)
+    np.testing.assert_array_equal(nodes[:3], seeds)
+    assert sum(n_counts) == len(nodes) and sum(e_counts) == len(edge)
+    # Every sampled edge is a real edge, with endpoints mapped to their local ids.
+    np.testing.assert_array_equal(nodes[row], edge_index[0, edge])
+    np.testing.assert_array_equal(nodes[col], edge_index[1, edge])
+    # First hop: each seed gets min(3, in-degree) edges, distinct when sampling without replacement.
+    in_degree = np.bincount(edge_index[1], minlength=40)
+    first_hop = slice(0, e_counts[0])
+    for i, s in enumerate(seeds):
+        picked = edge[first_hop][col[first_hop] == i]
+        assert len(picked) == min(3, in_degree[s])
+        if not replace:
+            assert len(np.unique(picked)) == len(picked)
+    # The reusable id map is left clean for the next batch.
+    assert (graph.local_map(40) == -1).all()
+
+
+def test_disjoint_and_out_of_range_seeds():
+    edge_index = np.array([[1, 2], [0, 0]])
+    nodes, row, col, edge, _, _ = sample_neighbors_homo(edge_index, np.array([0, 5]), [-1], num_nodes=3, disjoint=True)
+    np.testing.assert_array_equal(nodes, [0, 5, 1, 2])
+    np.testing.assert_array_equal(edge, [0, 1])