enhance nbr and N_v of graph with specified k-hop

Author: yifanfeng97

dhg/_global.py

@@ -12,6 +12,7 @@ def get_dhg_cache_root():
 # global paths
 CACHE_ROOT = get_dhg_cache_root()
 DATASETS_ROOT = CACHE_ROOT / "datasets"
-REMOTE_ROOT = "https://data.deephypergraph.com/"
+# REMOTE_ROOT = "https://data.deephypergraph.com/"
+REMOTE_ROOT = "https://download.moon-lab.tech:28501/"
 REMOTE_DATASETS_ROOT = REMOTE_ROOT + "datasets/"
 # REMOTE_DATASETS_ROOT = "https://data.shrec22.moon-lab.tech:18443/DHG/datasets/"

dhg/data/dblp.py

@@ -6,10 +6,9 @@
 
 
 class DBLP8k(BaseData):
-    r"""The DBLP-8k dataset is a citation network dataset for link prediction task.
-    The dataset is a part of the dataset crawled according to DBLP API,
-     and we have selected each item based on some conditions, 
-     such as the venue and publication year (from 2018 to 2022). It contains 6498 authors and 2603 papers.
+    r"""The DBLP-8k dataset is a citation network dataset for link prediction task. 
+    The dataset is a part of the dataset crawled according to DBLP API, and we have selected each item based on some conditions, such as the venue and publication year (from 2018 to 2022). It contains 6498 authors and 2603 papers.
+    
     The content of the DBLP-8k dataset includes the following:
 
     - ``num_vertices``: The number of vertices: :math:`8,657`.

dhg/structure/graphs/graph.py

@@ -421,13 +421,14 @@ def deg_v(self) -> List[int]:
         """
         return self.D_v._values().cpu().numpy().tolist()
 
-    def nbr_v(self, v_idx: int) -> Tuple[List[int], List[float]]:
-        r""" Return a vertex list of the neighbors of the vertex ``v_idx``.
+    def nbr_v(self, v_idx: int, hop: int = 1) -> List[int]:
+        r""" Return a vertex list of the ``k``-hop neighbors of the vertex ``v_idx``.
 
         Args:
             ``v_idx`` (``int``): The index of the vertex.
+            ``hop`` (``int``): The number of the hop.
         """
-        return self.N_v(v_idx).cpu().numpy().tolist()
+        return self.N_v(v_idx, hop).cpu().numpy().tolist()
 
     # =====================================================================================
     # properties for deep learning
@@ -521,13 +522,25 @@ def D_v_neg_1_2(self,) -> torch.Tensor:
             ).coalesce()
         return self.cache["D_v_neg_1_2"]
 
-    def N_v(self, v_idx: int) -> Tuple[List[int], List[float]]:
-        r""" Return the neighbors of the vertex ``v_idx`` with ``torch.Tensor`` format.
+    def N_v(self, v_idx: int, hop: int = 1) -> List[int]:
+        r""" Return the ``k``-hop neighbors of the vertex ``v_idx`` with ``torch.Tensor`` format.
 
         Args:
             ``v_idx`` (``int``): The index of the vertex.
+            ``hop`` (``int``): The number of the hop.
         """
-        sub_v_set = self.A[v_idx]._indices()[0].clone()
+        assert hop >= 1, "``hop`` must be a number larger than or equal to 1."
+        if hop == 1:
+            A_k = self.A
+        else:
+            if self.cache.get(f"A_{hop}") is None:
+                A_1, A_k = self.A.clone(), self.A.clone()
+                for _ in range(hop - 1):
+                    A_k = torch.sparse.mm(A_k, A_1)
+                self.cache[f"A_{hop}"] = A_k
+            else:
+                A_k = self.cache[f"A_{hop}"]
+        sub_v_set = A_k[v_idx]._indices()[0].clone()
         return sub_v_set
 
     @property

docs/source/api/data.rst

@@ -61,6 +61,7 @@ Hypergraph Datasets
     dhg.data.WalmartTrips
     dhg.data.HouseCommittees
     dhg.data.News20
+    dhg.data.DBLP8k
 
 
 **Welcome to contribute datasets!**

tests/structure/test_graph.py

@@ -174,6 +174,11 @@ def test_nbr(g1):
     assert g1.nbr_v(0) == [1, 2, 3]
     g1.remove_edges((0, 2))
     assert g1.nbr_v(2) == []
+    # hop k
+    g3 = Graph(5, [(0, 1), (0, 3), (1, 4), (2, 3)])
+    assert sorted(g3.nbr_v(3, 1)) == [0, 2]
+    assert sorted(g3.nbr_v(3, 2)) == [1, 3]
+    assert sorted(g3.nbr_v(3, 3)) == [0, 2, 4]
 
 
 # test deep learning