Merge pull request #110 from MarkDana/fix-shd-undirected-edge-comparison

Fixed an SHD bug: for the cases other than directed inversion

Author: kunwuz

causallearn/graph/SHD.py

@@ -1,4 +1,3 @@
-from causallearn.graph.Endpoint import Endpoint
 from causallearn.graph.Graph import Graph
 
 
@@ -7,8 +6,6 @@ class SHD:
     Compute the Structural Hamming Distance (SHD) between two graphs. In simple terms, this is the number of edge
     insertions, deletions or flips in order to transform one graph to another graph.
     """
-    __SHD = 0
-
     def __init__(self, truth: Graph, est: Graph):
         """
         Compute and store the Structural Hamming Distance (SHD) between two graphs.
@@ -20,30 +17,18 @@ def __init__(self, truth: Graph, est: Graph):
         est :
             Estimated graph.
         """
-        nodes = truth.get_nodes()
-        nodes_name = [node.get_name() for node in nodes]
-        self.__SHD: int = 0
+        truth_node_map = {node.get_name(): node_id for node, node_id in truth.node_map.items()}
+        est_node_map = {node.get_name(): node_id for node, node_id in est.node_map.items()}
+        assert set(truth_node_map.keys()) == set(est_node_map.keys()), "The two graphs have different sets of node names."
 
-        # Assumes the list of nodes for the two graphs are the same.
-        for i in list(range(0, len(nodes))):
-            for j in list(range(i + 1, len(nodes))):
-                if truth.get_edge(truth.get_node(nodes_name[i]), truth.get_node(nodes_name[j])) and (
-                        not est.get_edge(est.get_node(nodes_name[i]), est.get_node(nodes_name[j]))):
-                    self.__SHD += 1
-                if (not truth.get_edge(truth.get_node(nodes_name[i]), truth.get_node(nodes_name[j]))) and est.get_edge(
-                        est.get_node(nodes_name[i]), est.get_node(nodes_name[j])):
-                    self.__SHD += 1
-
-        for i in list(range(0, len(nodes))):
-            for j in list(range(0, len(nodes))):
-                if not truth.get_edge(truth.get_node(nodes_name[i]), truth.get_node(nodes_name[j])):
-                    continue
-                if not est.get_edge(est.get_node(nodes_name[i]), est.get_node(nodes_name[j])):
-                    continue
-                if truth.get_endpoint(truth.get_node(nodes_name[i]),
-                                      truth.get_node(nodes_name[j])) == Endpoint.ARROW and est.get_endpoint(
-                    est.get_node(nodes_name[j]), est.get_node(nodes_name[i])) == Endpoint.ARROW:
-                    self.__SHD += 1
+        self.__SHD: int = 0
+        for node_i_name, truth_node_i_id in truth_node_map.items():
+            for node_j_name, truth_node_j_id in truth_node_map.items():
+                if truth_node_j_id < truth_node_i_id: continue  # we allow `==' to care about the possibly self-loops.
+                est_node_i_id, est_node_j_id = est_node_map[node_i_name], est_node_map[node_j_name]
+                truth_ij_edge_endpoints = (truth.graph[truth_node_i_id, truth_node_j_id], truth.graph[truth_node_j_id, truth_node_i_id])
+                est_ij_edge_endpoints = (est.graph[est_node_i_id, est_node_j_id], est.graph[est_node_j_id, est_node_i_id])
+                if truth_ij_edge_endpoints != est_ij_edge_endpoints: self.__SHD += 1
 
     def get_shd(self) -> int:
         return self.__SHD