Merge pull request #140 from mschauer/bayesball

Add Bayesball graph

Author: mschauer

docs/src/library.md

@@ -102,6 +102,11 @@ graph_to_text
 kwargs_pdag_graphmakie
 CausalInference.combinations_without
 ```
+## Bayes ball
+```@docs
+CausalInference.bayesball
+CausalInference.bayesball_graph
+```
 
 ## Adjustment
 ```@docs
@@ -152,6 +157,12 @@ CausalInference.precompute_semidirected
 CausalInference.next_CPDAG
 ```
 
+## DAG Zig-Zag
+```@docs
+CausalInference.dagzigzag
+CausalInference.count_dag_moves 
+```
+
 ## Exact
 ```@docs
 CausalInference.exactscorebased

src/gensearch.jl

@@ -76,10 +76,30 @@ function descendants(g, X, veto = no_veto)
     return gensearch(g, X, (pe, ne, v, w) -> ne == RIGHT && !veto(pe, ne, v, w))
 end
 
+"""
+    bayesball(g, X, S = Set{eltype(g)}())
+
+Return the set of vertices d-connected to the set of vertices X given set of vertices S in dag g.
+"""
 function bayesball(g, X, S = Set{eltype(g)}(), veto = no_veto)
     return gensearch(g, X, (pe, ne, v, w) -> !veto(pe, ne, v, w) && (pe == INIT || (v in S && pe == RIGHT && ne == LEFT) || (!(v in S) && !(pe == RIGHT && ne == LEFT))))
 end
 
+
+"""
+    bayesball_graph(g, X, S = Set{eltype(g)}())
+
+Return an undirected graph `b` containing edges for possible moves of the
+Bayes ball. Vertex `i` of `g` is vertex `2i-1` of `b`
+if entered forward, `2i` if entered backward.
+"""
+function bayesball_graph(g, X, S = Set{eltype(g)}())
+    ι(e, i) = e == RIGHT ? 2i-1 : 2i 
+    edges = Pair{Int,Int}[]
+    CausalInference.gensearch(g, X, (pe, ne, v, w) ->  (pe == INIT || (v in S && pe == RIGHT && ne == LEFT) || (!(v in S) && !(pe == RIGHT && ne == LEFT))) && (push!(edges, ι(pe, v)=>ι(ne, w)); true))
+    CausalInference.graph(edges, 2*nv(g))
+end
+
 """
     alt_test_dsep(g, X, Y, S, veto = no_veto)
 

src/misc.jl

@@ -17,8 +17,7 @@ end
 
 Create `Graph` from edge-list.
 """
-function graph(E)
-    d = maximum(flatten(E))
+function graph(E, d = maximum(flatten(E)))  
     g = Graph(d)
     for (i, j) in E
         add_edge!(g, i, j)

test/gensearch.jl

@@ -1,6 +1,7 @@
 using Graphs
 using CausalInference
 using Test
+using CausalInference: graph 
 
 # for some of the tests other results would not be wrong per se 
 # e.g. for the finding methods, other valid sets could be returned
@@ -93,4 +94,10 @@ end
     @test_throws ArgumentError dsep(g1, [1,2], [2,3], [4,5])
     @test_throws ArgumentError dsep(g1, [1,2], [3,4], [4,5])
     @test_throws ArgumentError dsep(g1, [1,2], [3,4], [5,1])
+end
+
+@testset "bayesball_graph" begin
+    g = digraph([1=>3, 2=>3, 3=>4, 2=>4, 1=>4])
+    g2 = CausalInference.bayesball_graph(g, 2, [3])
+    @test g2 == graph([(2, 5), (2, 7), (4, 5), (4, 7)], 8)
 end