update node/hyedge count functions

Author: yifanfeng97

.idea/HyperG_package.iml

@@ -1,6 +1,6 @@
-def count_node(H):
-    return H[0].max().item() + 1
+def count_node(H, node_num=None):
+    return H[0].max().item() + 1 if node_num is None else node_num
 
 
-def count_hyedge(H):
-    return H[1].max().item() + 1
+def count_hyedge(H, hyedge_num=None):
+    return H[1].max().item() + 1 if hyedge_num is None else hyedge_num

.idea/misc.xml

@@ -3,18 +3,18 @@
 from .count import count_hyedge, count_node
 
 
-def degree_node(H):
+def degree_node(H, node_num=None):
     node_idx, edge_idx = H
-    node_num = count_node(H)
+    node_num = count_node(H, node_num)
     src = torch.ones_like(node_idx).float().to(H.device)
     out = torch.zeros(node_num).to(H.device)
     return out.scatter_add(0, node_idx, src).long()
     # return torch.zeros(node_num).scatter_add(0, node_idx, torch.ones_like(node_idx).float()).long()
 
 
-def degree_hyedge(H: torch.Tensor):
+def degree_hyedge(H: torch.Tensor, hyedge_num=None):
     node_idx, hyedge_idx = H
-    edge_num = count_hyedge(H)
+    hyedge_num = count_hyedge(H, hyedge_num=hyedge_num)
     src = torch.ones_like(hyedge_idx).float().to(H.device)
-    out = torch.zeros(edge_num).to(H.device)
+    out = torch.zeros(hyedge_num).to(H.device)
     return out.scatter_add(0, hyedge_idx, src).long()

SuperMoon/hyedge/utils/count.py

@@ -5,9 +5,9 @@
 from .verify import contiguous_hyedge_idx
 
 
-def self_loop_remove(H):
+def self_loop_remove(H, hyedge_num=None):
     node_idx, hyedge_idx = H
-    DE = degree_hyedge(H)
+    DE = degree_hyedge(H, hyedge_num)
     loop_edge_idx = torch.where(DE == 1)[0]
 
     mask = torch.ones_like(hyedge_idx).bool()
@@ -18,9 +18,9 @@ def self_loop_remove(H):
     return contiguous_hyedge_idx(H)
 
 
-def self_loop_add(H):
+def self_loop_add(H, node_num=None):
     H = self_loop_remove(H)
-    node_num = count_node(H)
+    node_num = count_node(H, node_num=node_num)
 
     loop_node_idx = torch.arange(node_num)
     loop_hyedge_idx = torch.arange(node_num)

SuperMoon/hyedge/utils/degree.py

@@ -3,9 +3,9 @@
 from .degree import degree_hyedge
 
 
-def contiguous_hyedge_idx(H):
+def contiguous_hyedge_idx(H, hyedge_num=None):
     node_idx, hyedge_idx = H
-    DE = degree_hyedge(H)
+    DE = degree_hyedge(H, hyedge_num)
     zero_idx = torch.where(DE == 0)[0]
 
     bias = torch.zeros_like(hyedge_idx)