tgraphx 0.3.0__tar.gz → 0.4.0__tar.gz

{tgraphx-0.3.0 → tgraphx-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tgraphx
-Version: 0.3.0
+Version: 0.4.0
 Summary: Tensor-aware graph neural networks preserving spatial node feature layouts
 Author-email: Arash Sajjadi <arash.sajjadi@usask.ca>
 Maintainer-email: Arash Sajjadi <arash.sajjadi@usask.ca>
@@ -18,7 +18,6 @@ Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -26,7 +25,7 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Scientific/Engineering :: Image Recognition
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=1.13
@@ -34,6 +33,7 @@ Requires-Dist: torchvision>=0.14
 Requires-Dist: pyyaml>=5.4
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: matplotlib>=3.7; extra == "dev"
 Requires-Dist: build>=0.10; extra == "dev"
 Requires-Dist: twine>=4.0; extra == "dev"
 Provides-Extra: monitoring

{tgraphx-0.3.0 → tgraphx-0.4.0}/pyproject.toml

@@ -5,10 +5,10 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "tgraphx"
 # Keep this in sync with tgraphx/__init__.py::__version__
-version = "0.3.0"
+version = "0.4.0"
 description = "Tensor-aware graph neural networks preserving spatial node feature layouts"
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 license = {text = "MIT"}
 authors = [
     {name = "Arash Sajjadi", email = "arash.sajjadi@usask.ca"},
@@ -35,7 +35,6 @@ classifiers = [
     "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
@@ -53,6 +52,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "pytest>=7.0",
+    "matplotlib>=3.7",
     "build>=0.10",
     "twine>=4.0",
 ]

tgraphx-0.4.0/tests/test_algorithms.py

@@ -0,0 +1,275 @@
+"""Tests for tgraphx.algorithms — v0.3.2 beta primitives.
+
+Covers connectivity (connected_components, is_connected, count) and
+traversal (bfs_layers, bfs_edges, shortest_path_length).
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+from tgraphx.algorithms import (
+    bfs_edges,
+    bfs_layers,
+    connected_components,
+    is_connected,
+    number_connected_components,
+    shortest_path_length,
+    weakly_connected_components,
+)
+
+
+# ── connectivity ─────────────────────────────────────────────────────────────
+
+
+class TestConnectedComponents:
+    def test_two_disjoint_paths(self):
+        # Components: {0, 1, 2} and {3, 4}.
+        edge_index = torch.tensor([[0, 1, 3], [1, 2, 4]], dtype=torch.long)
+        labels = connected_components(edge_index, num_nodes=5)
+        assert labels.tolist() == [0, 0, 0, 1, 1]
+        assert number_connected_components(edge_index, 5) == 2
+
+    def test_one_connected_graph(self):
+        # 0-1-2-0 triangle.
+        edge_index = torch.tensor([[0, 1, 2], [1, 2, 0]], dtype=torch.long)
+        labels = connected_components(edge_index, 3)
+        assert labels.tolist() == [0, 0, 0]
+        assert is_connected(edge_index, 3)
+        assert number_connected_components(edge_index, 3) == 1
+
+    def test_isolated_nodes(self):
+        edge_index = torch.zeros((2, 0), dtype=torch.long)
+        labels = connected_components(edge_index, num_nodes=4)
+        # 4 isolated nodes ⇒ 4 components, deterministic 0..3.
+        assert labels.tolist() == [0, 1, 2, 3]
+        assert number_connected_components(edge_index, 4) == 4
+        assert not is_connected(edge_index, 4)
+
+    def test_single_node_is_connected(self):
+        assert is_connected(torch.zeros((2, 0), dtype=torch.long), num_nodes=1)
+
+    def test_empty_graph_not_connected(self):
+        assert not is_connected(torch.zeros((2, 0), dtype=torch.long), num_nodes=0)
+        assert number_connected_components(
+            torch.zeros((2, 0), dtype=torch.long), num_nodes=0,
+        ) == 0
+
+    def test_directed_treated_as_undirected(self):
+        # Edge 0→1 only (not 1→0).
+        edge_index = torch.tensor([[0], [1]], dtype=torch.long)
+        labels = connected_components(edge_index, num_nodes=2)
+        assert labels.tolist() == [0, 0]
+
+    def test_weakly_connected_alias(self):
+        edge_index = torch.tensor([[0, 1, 2, 4], [1, 2, 3, 5]], dtype=torch.long)
+        a = connected_components(edge_index, 6)
+        b = weakly_connected_components(edge_index, 6)
+        assert torch.equal(a, b)
+
+    def test_compact_label_ids(self):
+        # Components are returned as contiguous [0, K).
+        edge_index = torch.tensor([[1, 2, 5], [2, 3, 6]], dtype=torch.long)
+        labels = connected_components(edge_index, num_nodes=7)
+        unique = torch.unique(labels)
+        # Labels must form a contiguous range starting at 0.
+        assert unique.tolist() == list(range(unique.numel()))
+
+    def test_validates_edge_index_shape(self):
+        with pytest.raises(ValueError, match="\\[2, E\\]"):
+            connected_components(torch.zeros(3, 4, dtype=torch.long), num_nodes=4)
+
+    def test_validates_edge_id_range(self):
+        bad = torch.tensor([[0, 5], [1, 2]], dtype=torch.long)
+        with pytest.raises(ValueError, match="out of range"):
+            connected_components(bad, num_nodes=4)
+
+    def test_inferred_num_nodes(self):
+        edge_index = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
+        labels = connected_components(edge_index)  # infers num_nodes=3
+        assert labels.numel() == 3
+
+
+# ── BFS layers ───────────────────────────────────────────────────────────────
+
+
+class TestBFSLayers:
+    def test_chain(self):
+        # 0 → 1 → 2 → 3
+        edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)
+        layers = bfs_layers(edge_index, source=0, num_nodes=4)
+        assert [l.tolist() for l in layers] == [[0], [1], [2], [3]]
+
+    def test_branching(self):
+        # 0 → {1, 2}; 1 → 3
+        edge_index = torch.tensor([[0, 0, 1], [1, 2, 3]], dtype=torch.long)
+        layers = bfs_layers(edge_index, source=0, num_nodes=4)
+        assert layers[0].tolist() == [0]
+        assert sorted(layers[1].tolist()) == [1, 2]
+        assert layers[2].tolist() == [3]
+
+    def test_isolated_source(self):
+        layers = bfs_layers(torch.zeros((2, 0), dtype=torch.long),
+                            source=0, num_nodes=3)
+        assert [l.tolist() for l in layers] == [[0]]
+
+    def test_max_hops(self):
+        edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]], dtype=torch.long)
+        layers = bfs_layers(edge_index, source=0, num_nodes=5, max_hops=2)
+        assert sum(l.numel() for l in layers) == 3  # source + 2 hops
+
+    def test_invalid_source(self):
+        with pytest.raises(ValueError, match="source"):
+            bfs_layers(torch.zeros((2, 0), dtype=torch.long),
+                       source=10, num_nodes=4)
+
+
+# ── BFS edges ────────────────────────────────────────────────────────────────
+
+
+class TestBFSEdges:
+    def test_chain_edges(self):
+        edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)
+        bfs = bfs_edges(edge_index, source=0, num_nodes=4)
+        # Expect exactly the chain.
+        assert bfs.tolist() == [[0, 1, 2], [1, 2, 3]]
+
+    def test_no_edges_returns_empty(self):
+        bfs = bfs_edges(torch.zeros((2, 0), dtype=torch.long),
+                        source=0, num_nodes=3)
+        assert bfs.shape == (2, 0)
+
+    def test_branching_count(self):
+        # 0 → {1, 2}; 1 → 3
+        edge_index = torch.tensor([[0, 0, 1], [1, 2, 3]], dtype=torch.long)
+        bfs = bfs_edges(edge_index, source=0, num_nodes=4)
+        # Exactly 3 reachable non-source nodes ⇒ 3 BFS edges.
+        assert bfs.size(1) == 3
+        # Each child node appears once (BFS tree property).
+        assert sorted(bfs[1].tolist()) == [1, 2, 3]
+
+    def test_no_revisit(self):
+        # Cycle: 0 → 1 → 2 → 0
+        edge_index = torch.tensor([[0, 1, 2], [1, 2, 0]], dtype=torch.long)
+        bfs = bfs_edges(edge_index, source=0, num_nodes=3)
+        # Only 2 BFS edges — node 0 is the source and is not revisited.
+        assert bfs.size(1) == 2
+
+    def test_max_hops(self):
+        edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)
+        bfs = bfs_edges(edge_index, source=0, num_nodes=4, max_hops=1)
+        # With 1 hop only, BFS tree edges = 1.
+        assert bfs.size(1) == 1
+
+
+# ── Shortest path length ─────────────────────────────────────────────────────
+
+
+class TestShortestPathLength:
+    def test_chain(self):
+        edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)
+        d = shortest_path_length(edge_index, source=0, num_nodes=4)
+        assert d.tolist() == [0, 1, 2, 3]
+
+    def test_unreachable_node(self):
+        # 0 → 1 ; 2 isolated (in directed sense)
+        edge_index = torch.tensor([[0], [1]], dtype=torch.long)
+        d = shortest_path_length(edge_index, source=0, num_nodes=3)
+        assert d[0].item() == 0
+        assert d[1].item() == 1
+        assert d[2].item() == -1
+
+    def test_self_distance(self):
+        edge_index = torch.zeros((2, 0), dtype=torch.long)
+        d = shortest_path_length(edge_index, source=0, num_nodes=3)
+        assert d[0].item() == 0
+        assert d[1].item() == -1
+        assert d[2].item() == -1
+
+    def test_branching_distances(self):
+        # 0 → {1, 2}; 1 → 3
+        edge_index = torch.tensor([[0, 0, 1], [1, 2, 3]], dtype=torch.long)
+        d = shortest_path_length(edge_index, source=0, num_nodes=4)
+        assert d.tolist() == [0, 1, 1, 2]
+
+    def test_directed_one_way(self):
+        # 0 → 1, 1 → 2, no reverse edges.  From 2 nothing is reachable.
+        edge_index = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
+        d = shortest_path_length(edge_index, source=2, num_nodes=3)
+        assert d[2].item() == 0
+        assert d[0].item() == -1
+        assert d[1].item() == -1
+
+
+# ── Edge cases and robustness ─────────────────────────────────────────────────
+
+
+class TestEdgeCasesAndValidation:
+    def test_self_loop_edges_do_not_crash_bfs(self):
+        # Self-loops should be ignored by BFS (node already visited).
+        edge_index = torch.tensor([[0, 0, 1], [0, 1, 2]], dtype=torch.long)
+        layers = bfs_layers(edge_index, source=0, num_nodes=3)
+        all_nodes = set()
+        for l in layers:
+            all_nodes.update(l.tolist())
+        assert 0 in all_nodes and 1 in all_nodes and 2 in all_nodes
+
+    def test_self_loop_edges_do_not_crash_components(self):
+        edge_index = torch.tensor([[0, 1, 1], [1, 2, 1]], dtype=torch.long)  # node 1 self-loop
+        labels = connected_components(edge_index, num_nodes=3)
+        assert labels.tolist() == [0, 0, 0]
+
+    def test_duplicate_edges_do_not_duplicate_components(self):
+        # Duplicate edge 0→1 should not change component structure.
+        edge_index = torch.tensor([[0, 0, 1], [1, 1, 2]], dtype=torch.long)
+        labels = connected_components(edge_index, num_nodes=3)
+        assert labels.tolist() == [0, 0, 0]
+
+    def test_bfs_out_of_range_raises(self):
+        edge_index = torch.tensor([[0, 1], [1, 5]], dtype=torch.long)
+        with pytest.raises(ValueError, match="out of range"):
+            bfs_layers(edge_index, source=0, num_nodes=3)
+
+    def test_shortest_path_out_of_range_raises(self):
+        edge_index = torch.tensor([[0, 1], [1, 5]], dtype=torch.long)
+        with pytest.raises(ValueError, match="out of range"):
+            shortest_path_length(edge_index, source=0, num_nodes=3)
+
+    def test_large_sparse_graph_smoke(self):
+        # 1000-node ring, both directions.
+        N = 1000
+        src = torch.arange(N)
+        dst = (src + 1) % N
+        edge_index = torch.stack([
+            torch.cat([src, dst]),
+            torch.cat([dst, src]),
+        ], dim=0)
+        labels = connected_components(edge_index, num_nodes=N)
+        assert int(number_connected_components(edge_index, N)) == 1
+        d = shortest_path_length(edge_index, source=0, num_nodes=N)
+        # Max distance in a ring is N//2.
+        assert d.min().item() == 0
+        assert d.max().item() == N // 2
+
+    @pytest.mark.skipif(
+        not pytest.importorskip("networkx", reason="networkx not installed"),
+        reason="networkx not installed",
+    )
+    def test_parity_with_networkx_components(self):
+        """Optional parity test if networkx is installed."""
+        import networkx as nx
+        G = nx.Graph()
+        edges = [(0, 1), (1, 2), (4, 5)]
+        G.add_edges_from(edges)
+        # All 7 nodes must be present in both graphs for a fair comparison.
+        for n in range(7):
+            G.add_node(n)
+        edge_index = torch.tensor(
+            [[u for u, v in edges] + [v for u, v in edges],
+             [v for u, v in edges] + [u for u, v in edges]],
+            dtype=torch.long,
+        )
+        labels = connected_components(edge_index, num_nodes=7)
+        nx_count = nx.number_connected_components(G)
+        tgx_count = int(labels.max().item()) + 1
+        assert tgx_count == nx_count

tgraphx-0.4.0/tests/test_graph_utils.py

@@ -0,0 +1,118 @@
+"""Tests for tgraphx.algorithms.structural — degree utilities (v0.3.2 beta).
+
+Invariants tested
+-----------------
+- out-degree, in-degree, total-degree all correct on directed graph.
+- Isolated nodes have degree 0.
+- Self-loops count correctly.
+- Empty edge_index returns all-zero degrees.
+- dtype and device preserved.
+- degree_features shape and log_scale behavior.
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+from tgraphx.algorithms import degree, degree_features
+
+
+class TestDegree:
+    def test_out_degree_chain(self):
+        # 0 → 1 → 2 → 3
+        ei = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)
+        d = degree(ei, num_nodes=4, mode="out")
+        assert d.tolist() == [1, 1, 1, 0]
+
+    def test_in_degree_chain(self):
+        ei = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)
+        d = degree(ei, num_nodes=4, mode="in")
+        assert d.tolist() == [0, 1, 1, 1]
+
+    def test_total_degree_chain(self):
+        ei = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)
+        d = degree(ei, num_nodes=4, mode="both")
+        assert d.tolist() == [1, 2, 2, 1]
+
+    def test_isolated_nodes_zero(self):
+        ei = torch.zeros((2, 0), dtype=torch.long)
+        d = degree(ei, num_nodes=5, mode="out")
+        assert d.tolist() == [0, 0, 0, 0, 0]
+
+    def test_empty_edge_index_all_zero(self):
+        ei = torch.zeros((2, 0), dtype=torch.long)
+        d = degree(ei, num_nodes=3)
+        assert d.shape == (3,)
+        assert d.sum().item() == 0
+
+    def test_self_loop_counts_out_and_in(self):
+        # Self-loop at node 0.
+        ei = torch.tensor([[0], [0]], dtype=torch.long)
+        assert degree(ei, num_nodes=2, mode="out").tolist() == [1, 0]
+        assert degree(ei, num_nodes=2, mode="in").tolist() == [1, 0]
+        assert degree(ei, num_nodes=2, mode="both").tolist() == [2, 0]
+
+    def test_dtype_override(self):
+        ei = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
+        d = degree(ei, num_nodes=3, dtype=torch.float32)
+        assert d.dtype == torch.float32
+
+    def test_inferred_num_nodes(self):
+        ei = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
+        d = degree(ei)  # infers num_nodes=3
+        assert d.shape == (3,)
+
+    def test_out_degree_star(self):
+        # Hub at node 0 connected to 1,2,3.
+        ei = torch.tensor([[0, 0, 0], [1, 2, 3]], dtype=torch.long)
+        d = degree(ei, num_nodes=4, mode="out")
+        assert d.tolist() == [3, 0, 0, 0]
+
+    def test_invalid_mode(self):
+        ei = torch.tensor([[0], [1]], dtype=torch.long)
+        with pytest.raises(ValueError, match="mode"):
+            degree(ei, num_nodes=2, mode="all")
+
+    def test_invalid_edge_index_shape(self):
+        with pytest.raises(ValueError, match="\\[2, E\\]"):
+            degree(torch.zeros(3, 4, dtype=torch.long), num_nodes=4)
+
+    def test_edge_id_out_of_range(self):
+        ei = torch.tensor([[0, 5], [1, 2]], dtype=torch.long)
+        with pytest.raises(ValueError):
+            degree(ei, num_nodes=4, mode="out")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_cuda_device_preserved(self):
+        ei = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long).cuda()
+        d = degree(ei, num_nodes=4, mode="out")
+        assert d.device.type == "cuda"
+        assert d.tolist() == [1, 1, 1, 0]
+
+
+class TestDegreeFeatures:
+    def test_shape(self):
+        ei = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
+        feats = degree_features(ei, num_nodes=3)
+        assert feats.shape == (3, 3)
+
+    def test_columns(self):
+        # 0 → 1 → 2
+        ei = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
+        feats = degree_features(ei, num_nodes=3)
+        # out_deg: [1,1,0]; in_deg: [0,1,1]; total: [1,2,1]
+        assert feats[:, 0].tolist() == [1, 1, 0]
+        assert feats[:, 1].tolist() == [0, 1, 1]
+        assert feats[:, 2].tolist() == [1, 2, 1]
+
+    def test_log_scale(self):
+        ei = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
+        feats = degree_features(ei, num_nodes=3, log_scale=True)
+        assert feats.dtype == torch.float32
+        # All values should be >= 0 (log1p of non-negative integers).
+        assert (feats >= 0).all()
+
+    def test_empty_graph(self):
+        feats = degree_features(torch.zeros((2, 0), dtype=torch.long), num_nodes=4)
+        assert feats.shape == (4, 3)
+        assert feats.sum().item() == 0

tgraphx-0.4.0/tests/test_hard_negative_sampling.py

@@ -0,0 +1,155 @@
+"""Tests for hard_negative_sampling (v0.3.2 beta).
+
+Invariants tested
+-----------------
+- No false negatives (no positive edge returned).
+- No self-loops.
+- No duplicates in output.
+- Hard negatives have higher mean similarity than random negatives on a
+  toy embedding where similarity is strongly correlated with id proximity.
+- Deterministic with seed.
+- No global RNG pollution.
+- force_undirected semantics.
+- candidate_pool_size exhaustion produces warning.
+- Device preserved (CUDA optional).
+"""
+from __future__ import annotations
+
+import warnings
+
+import pytest
+import torch
+
+from tgraphx import hard_negative_sampling, negative_sampling
+
+
+def _chain_graph(N: int):
+    ei = torch.tensor([[i for i in range(N - 1)], [i + 1 for i in range(N - 1)]], dtype=torch.long)
+    emb = torch.zeros(N, 4)
+    for i in range(N):
+        emb[i, 0] = float(i) / N
+    return ei, emb, N
+
+
+class TestHardNegativeSampling:
+    def test_no_false_negatives(self):
+        ei, emb, N = _chain_graph(8)
+        neg = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, seed=0)
+        pos = {(int(ei[0, c]), int(ei[1, c])) for c in range(ei.size(1))}
+        for c in range(neg.size(1)):
+            assert (int(neg[0, c]), int(neg[1, c])) not in pos
+
+    def test_no_self_loops(self):
+        ei, emb, N = _chain_graph(8)
+        neg = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, seed=0)
+        for c in range(neg.size(1)):
+            assert int(neg[0, c]) != int(neg[1, c])
+
+    def test_no_duplicates(self):
+        ei, emb, N = _chain_graph(10)
+        neg = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=8, seed=0)
+        pairs = [(int(neg[0, c]), int(neg[1, c])) for c in range(neg.size(1))]
+        assert len(pairs) == len(set(pairs))
+
+    def test_harder_than_random(self):
+        ei, emb, N = _chain_graph(10)
+        emb_norm = emb / (emb.norm(dim=1, keepdim=True) + 1e-8)
+
+        neg_hard = hard_negative_sampling(
+            ei, emb, num_nodes=N, num_neg_samples=4,
+            candidate_pool_size=512, seed=0,
+        )
+        neg_rand = negative_sampling(ei, N, num_neg_samples=4, seed=0)
+
+        def mean_sim(ne):
+            if ne.size(1) == 0:
+                return 0.0
+            return float((emb_norm[ne[0]] * emb_norm[ne[1]]).sum(dim=1).mean().item())
+
+        # Hard negatives should have at least as high mean similarity as random.
+        assert mean_sim(neg_hard) >= mean_sim(neg_rand) - 0.05
+
+    def test_determinism(self):
+        ei, emb, N = _chain_graph(8)
+        a = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, seed=7)
+        b = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, seed=7)
+        assert torch.equal(a, b)
+
+    def test_different_seeds_different_output(self):
+        ei, emb, N = _chain_graph(12)
+        a = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, seed=1)
+        b = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, seed=2)
+        # Different seeds should (very likely) produce different results.
+        assert not torch.equal(a, b)
+
+    def test_no_global_rng_pollution(self):
+        ei, emb, N = _chain_graph(8)
+        torch.manual_seed(99)
+        before = torch.rand(3)
+        torch.manual_seed(99)
+        hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, seed=42)
+        after = torch.rand(3)
+        assert torch.equal(before, after)
+
+    def test_output_dtype_and_device(self):
+        ei, emb, N = _chain_graph(6)
+        neg = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=2, seed=0)
+        assert neg.dtype == torch.long
+        assert neg.device == ei.device
+
+    def test_force_undirected_no_reverse_positives(self):
+        ei = torch.tensor([[0, 1], [1, 0]], dtype=torch.long)
+        emb = torch.randn(4, 4)
+        neg = hard_negative_sampling(
+            ei, emb, num_nodes=4, num_neg_samples=4,
+            force_undirected=True, seed=0,
+        )
+        forbidden = {(0, 1), (1, 0)}
+        for c in range(neg.size(1)):
+            assert (int(neg[0, c]), int(neg[1, c])) not in forbidden
+
+    def test_cosine_vs_dot_different(self):
+        ei, emb, N = _chain_graph(8)
+        c = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, method="cosine", seed=0)
+        d = hard_negative_sampling(ei, emb, num_nodes=N, num_neg_samples=4, method="dot", seed=0)
+        # Both should be valid (no false negatives).
+        pos = {(int(ei[0, i]), int(ei[1, i])) for i in range(ei.size(1))}
+        for ne in (c, d):
+            for col in range(ne.size(1)):
+                assert (int(ne[0, col]), int(ne[1, col])) not in pos
+
+    def test_small_pool_warns(self):
+        ei, emb, N = _chain_graph(8)
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            neg = hard_negative_sampling(
+                ei, emb, num_nodes=N, num_neg_samples=100,
+                candidate_pool_size=4, seed=0,
+            )
+        # Should warn about returning fewer than requested.
+        assert any("candidate_pool_size" in str(warning.message) or
+                   "hard_negative_sampling" in str(warning.message)
+                   for warning in w) or neg.size(1) < 100
+
+    def test_invalid_method(self):
+        ei, emb, N = _chain_graph(4)
+        with pytest.raises(ValueError, match="method"):
+            hard_negative_sampling(ei, emb, num_nodes=N, method="euclidean")
+
+    def test_invalid_pool_size(self):
+        ei, emb, N = _chain_graph(4)
+        with pytest.raises(ValueError, match="candidate_pool_size"):
+            hard_negative_sampling(ei, emb, num_nodes=N, candidate_pool_size=0)
+
+    def test_embeddings_too_small(self):
+        ei = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
+        emb = torch.zeros(2, 4)  # only 2 embeddings for 3 nodes
+        with pytest.raises(ValueError, match="node_embeddings"):
+            hard_negative_sampling(ei, emb, num_nodes=3)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_cuda_device_preserved(self):
+        ei = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long).cuda()
+        emb = torch.randn(5, 4).cuda()
+        neg = hard_negative_sampling(ei, emb, num_nodes=5, num_neg_samples=4, seed=0)
+        assert neg.device.type == "cuda"

{tgraphx-0.3.0 → tgraphx-0.4.0}/tests/test_imports.py

@@ -113,6 +113,7 @@ def test_import_performance_from_top():
 
 def test_import_model_classes_from_top():
     from tgraphx import (  # noqa: F401
+        CNNEncoder,
         GraphClassifier,
         NodeClassifier,
         EdgePredictor,