DAG to a maximal PAG by adding rule5-10 and adding selection variables

Author: EvieQ01

causallearn/utils/DAG2PAG.py

@@ -10,10 +10,11 @@
 from causallearn.graph.Endpoint import Endpoint
 from causallearn.graph.GeneralGraph import GeneralGraph
 from causallearn.graph.Node import Node
-from causallearn.search.ConstraintBased.FCI import rule0, rulesR1R2cycle, ruleR3, ruleR4B
+from causallearn.search.ConstraintBased.FCI import rule0, rulesR1R2cycle, ruleR3, ruleR4B, ruleR5, ruleR6, ruleR7, rule8, rule9, rule10
 from causallearn.utils.cit import CIT, d_separation
 
-def dag2pag(dag: Dag, islatent: List[Node]) -> GeneralGraph:
+
+def dag2pag(dag: Dag, islatent: List[Node], isselection: List[Node] = []) -> GeneralGraph:
     """
     Convert a DAG to its corresponding PAG
     Parameters
@@ -27,8 +28,8 @@ def dag2pag(dag: Dag, islatent: List[Node]) -> GeneralGraph:
     dg = nx.DiGraph()
     true_dag = nx.DiGraph()
     nodes = dag.get_nodes()
-    observed_nodes = list(set(nodes) - set(islatent))
-    mod_nodes = observed_nodes + islatent
+    observed_nodes = list(set(nodes) - set(islatent) - set(isselection))
+    mod_nodes = observed_nodes + islatent + isselection
     nodes = dag.get_nodes()
     nodes_ids = {node: i for i, node in enumerate(nodes)}
     mod_nodeids = {node: i for i, node in enumerate(mod_nodes)}
@@ -65,7 +66,7 @@ def dag2pag(dag: Dag, islatent: List[Node]) -> GeneralGraph:
             for Z in combinations(observed_nodes, l):
                 if nodex in Z or nodey in Z:
                     continue
-                if d_separated(dg, {nodes_ids[nodex]}, {nodes_ids[nodey]}, set(nodes_ids[z] for z in Z)):
+                if d_separated(dg, {nodes_ids[nodex]}, {nodes_ids[nodey]}, set(nodes_ids[z] for z in Z) | set([nodes_ids[s] for s in isselection])):
                     if edge:
                         PAG.remove_edge(edge)
                     sepset[(nodes_ids[nodex], nodes_ids[nodey])] |= set(Z)
@@ -105,6 +106,13 @@ def dag2pag(dag: Dag, islatent: List[Node]) -> GeneralGraph:
         change_flag = ruleR4B(PAG, -1, data, independence_test_method, 0.05, sep_sets=sepset_reindexed,
                           change_flag=change_flag,
                           bk=None, verbose=False)
+        change_flag = ruleR5(PAG, changeFlag=change_flag, verbose=True)
+        change_flag = ruleR6(PAG, changeFlag=change_flag)
+        change_flag = ruleR7(PAG, changeFlag=change_flag)
+        change_flag = rule8(PAG, nodes=observed_nodes, changeFlag=change_flag)
+        change_flag = rule9(PAG, nodes=observed_nodes, changeFlag=change_flag)
+        change_flag = rule10(PAG, changeFlag=change_flag)
+
     return PAG
 
 

tests/TestDAG2PAG.py

@@ -72,3 +72,17 @@ def test_case3(self):
         print(pag)
         graphviz_pag = GraphUtils.to_pgv(pag)
         graphviz_pag.draw("pag.png", prog='dot', format='png')
+
+    def test_case_selection(self):
+        nodes = []
+        for i in range(5):
+            nodes.append(GraphNode(str(i)))
+        dag = Dag(nodes)
+        dag.add_directed_edge(nodes[0], nodes[1])
+        dag.add_directed_edge(nodes[1], nodes[2])
+        dag.add_directed_edge(nodes[2], nodes[3])
+        # Selection nodes
+        dag.add_directed_edge(nodes[3], nodes[4])
+        dag.add_directed_edge(nodes[0], nodes[4])
+        pag = dag2pag(dag, islatent=[], isselection=[nodes[4]])
+        print(pag)