Ancestral Sampling and Bayes Ball Algorithm

src/experimental/ProbabilisticGraphicalModels/bayesnet.jl

@@ -4,7 +4,7 @@
 A structure representing a Bayesian Network.
 """
 struct BayesianNetwork{V,T,F}
-    graph::SimpleGraph{T}
+    graph::SimpleDiGraph{T}
     "names of the variables in the network"
     names::Vector{V}
     "mapping from variable names to ids"
@@ -25,7 +25,7 @@ end
 
 function BayesianNetwork{V}() where {V}
     return BayesianNetwork(
-        SimpleGraph{Int}(), # by default, vertex ids are integers
+        SimpleDiGraph{Int}(), # by default, vertex ids are integers
         V[],
         Dict{V,Int}(),
         Dict{V,Any}(),
@@ -166,11 +166,26 @@ Ancestral sampling works by:
 2. Sampling from each node in order, using the already-sampled parent values for conditional distributions
 """
 function ancestral_sampling(bn::BayesianNetwork{V}) where {V}
-    ordered_vertices = Graphs.topological_sort(bn.graph)
-
+    ordered_vertices = Graphs.topological_sort_by_dfs(bn.graph)
     samples = Dict{V,Any}()
 
-    # TODO: Implement sampling logic
+    for vertex_id in ordered_vertices
+        vertex_name = bn.names[vertex_id]
+        if bn.is_observed[vertex_id]
+            samples[vertex_name] = bn.values[vertex_name]
+            continue
+        end
+        if bn.is_stochastic[vertex_id]
+            dist_idx = findfirst(id -> id == vertex_id, bn.stochastic_ids)
+            samples[vertex_name] = rand(bn.distributions[dist_idx])
+        else
+            # deterministic node
+            parent_ids = Graphs.inneighbors(bn.graph, vertex_id)
+            parent_values = [samples[bn.names[pid]] for pid in parent_ids]
+            func_idx = findfirst(id -> id == vertex_id, bn.deterministic_ids)
+            samples[vertex_name] = bn.deterministic_functions[func_idx](parent_values...)
+        end
+    end
 
     return samples
 end
@@ -184,11 +199,98 @@ If Z is provided, the conditioning information in `bn` will be ignored.
 function is_conditionally_independent end
 
 function is_conditionally_independent(bn::BayesianNetwork{V}, X::V, Y::V) where {V}
-    # TODO: Implement
+    # Use currently observed variables as Z
+    Z = V[v for (v, is_obs) in zip(bn.names, bn.is_observed) if is_obs]
+    return is_conditionally_independent(bn, X, Y, Z)
 end
 
 function is_conditionally_independent(
     bn::BayesianNetwork{V}, X::V, Y::V, Z::Vector{V}
 ) where {V}
-    # TODO: Implement
+    # Get vertex IDs
+    x_id = bn.names_to_ids[X]
+    y_id = bn.names_to_ids[Y]
+    z_ids = Set([bn.names_to_ids[z] for z in Z])
+
+    # Track visited nodes and their states
+    n_vertices = nv(bn.graph)
+    visited = falses(n_vertices)
+
+    # Queue entries are (node_id, from_parent)
+    queue = Tuple{Int,Bool}[]
+
+    # Start from X
+    push!(queue, (x_id, true))   # As if coming from parent
+    push!(queue, (x_id, false))  # As if coming from child
+
+    while !isempty(queue)
+        current_id, from_parent = popfirst!(queue)
+
+        if visited[current_id]
+            continue
+        end
+        visited[current_id] = true
+
+        # If we reached Y, path is active
+        if current_id == y_id
+            return false
+        end
+
+        is_conditioned = current_id in z_ids
+        parents = inneighbors(bn.graph, current_id)
+        children = outneighbors(bn.graph, current_id)
+
+        # Case 1: Node is not conditioned
+        if !is_conditioned
+            # Can go to children if coming from parent or at start node
+            if from_parent || current_id == x_id
+                for child in children
+                    push!(queue, (child, true))
+                end
+            end
+
+            # Can go to parents if coming from child or at start node
+            if !from_parent || current_id == x_id
+                for parent in parents
+                    push!(queue, (parent, false))
+                end
+            end
+        end
+
+        # Case 2: Node is conditioned or has conditioned descendants
+        if is_conditioned || has_conditioned_descendant(bn, current_id, z_ids)
+            # If this is a collider or descendant of collider
+            if length(parents) > 1 || !isempty(children)
+                # Can go to parents regardless of direction
+                for parent in parents
+                    push!(queue, (parent, false))
+                end
+            end
+        end
+    end
+
+    return true
+end
+
+function has_conditioned_descendant(bn::BayesianNetwork, node_id::Int, z_ids::Set{Int})
+    visited = falses(nv(bn.graph))
+    queue = Int[node_id]
+
+    while !isempty(queue)
+        current = popfirst!(queue)
+
+        if visited[current]
+            continue
+        end
+        visited[current] = true
+
+        if current in z_ids
+            return true
+        end
+
+        # Add all unvisited children
+        append!(queue, filter(c -> !visited[c], outneighbors(bn.graph, current)))
+    end
+
+    return false
 end

test/experimental/ProbabilisticGraphicalModels/bayesnet.jl

@@ -7,8 +7,9 @@ using JuliaBUGS.ProbabilisticGraphicalModels:
     add_deterministic_vertex!,
     add_edge!,
     condition,
-    decondition
-
+    decondition,
+    ancestral_sampling,
+    is_conditionally_independent
 @testset "BayesianNetwork" begin
     @testset "Adding vertices" begin
         bn = BayesianNetwork{Symbol}()
@@ -96,7 +97,169 @@ using JuliaBUGS.ProbabilisticGraphicalModels:
         @test bn_cond2.values[:B] == 2.0
     end
 
-    @testset "Simple ancestral sampling" begin end
+    @testset "Simple ancestral sampling" begin
+        bn = BayesianNetwork{Symbol}()
+        # Add stochastic vertices
+        add_stochastic_vertex!(bn, :A, Normal(0, 1), false)
+        add_stochastic_vertex!(bn, :B, Normal(1, 2), false)
+        # Add deterministic vertex C = A + B
+        add_deterministic_vertex!(bn, :C, (a, b) -> a + b)
+        add_edge!(bn, :A, :C)
+        add_edge!(bn, :B, :C)
+        samples = ancestral_sampling(bn)
+        @test haskey(samples, :A)
+        @test haskey(samples, :B)
+        @test haskey(samples, :C)
+        @test samples[:A] isa Number
+        @test samples[:B] isa Number
+        @test samples[:C] ≈ samples[:A] + samples[:B]
+    end
+
+    @testset "Complex ancestral sampling" begin
+        bn = BayesianNetwork{Symbol}()
+        add_stochastic_vertex!(bn, :μ, Normal(0, 2), false)
+        add_stochastic_vertex!(bn, :σ, LogNormal(0, 0.5), false)
+        add_stochastic_vertex!(bn, :X, Normal(0, 1), false)
+        add_stochastic_vertex!(bn, :Y, Normal(0, 1), false)
+        add_deterministic_vertex!(bn, :X_scaled, (μ, σ, x) -> x * σ + μ)
+        add_deterministic_vertex!(bn, :Y_scaled, (μ, σ, y) -> y * σ + μ)
+        add_deterministic_vertex!(bn, :Sum, (x, y) -> x + y)
+        add_deterministic_vertex!(bn, :Product, (x, y) -> x * y)
+        add_deterministic_vertex!(bn, :N, () -> 2.0)
+        add_deterministic_vertex!(bn, :Mean, (s, n) -> s / n)
+        add_edge!(bn, :μ, :X_scaled)
+        add_edge!(bn, :σ, :X_scaled)
+        add_edge!(bn, :X, :X_scaled)
+        add_edge!(bn, :μ, :Y_scaled)
+        add_edge!(bn, :σ, :Y_scaled)
+        add_edge!(bn, :Y, :Y_scaled)
+        add_edge!(bn, :X_scaled, :Sum)
+        add_edge!(bn, :Y_scaled, :Sum)
+        add_edge!(bn, :X_scaled, :Product)
+        add_edge!(bn, :Y_scaled, :Product)
+        add_edge!(bn, :Sum, :Mean)
+        add_edge!(bn, :N, :Mean)
+        samples = ancestral_sampling(bn)
+
+        @test all(
+            haskey(samples, k) for
+            k in [:μ, :σ, :X, :Y, :X_scaled, :Y_scaled, :Sum, :Product, :Mean, :N]
+        )
+
+        @test all(samples[k] isa Number for k in keys(samples))
+        @test samples[:X_scaled] ≈ samples[:X] * samples[:σ] + samples[:μ]
+        @test samples[:Y_scaled] ≈ samples[:Y] * samples[:σ] + samples[:μ]
+        @test samples[:Sum] ≈ samples[:X_scaled] + samples[:Y_scaled]
+        @test samples[:Product] ≈ samples[:X_scaled] * samples[:Y_scaled]
+        @test samples[:Mean] ≈ samples[:Sum] / samples[:N]
+        @test samples[:N] ≈ 2.0
+        @test samples[:σ] > 0
+        # Multiple samples test
+        n_samples = 1000
+        means = zeros(n_samples)
+        for i in 1:n_samples
+            samples = ancestral_sampling(bn)
+            means[i] = samples[:Mean]
+        end
+
+        @test mean(means) ≈ 0 atol = 0.5
+        @test std(means) > 0
+    end
+
+    @testset "Bayes Ball" begin
+        @testset "Chain Structure (A → B → C)" begin
+            bn = BayesianNetwork{Symbol}()
+
+            add_stochastic_vertex!(bn, :A, Normal(), false)
+            add_stochastic_vertex!(bn, :B, Normal(), false)
+            add_stochastic_vertex!(bn, :C, Normal(), false)
+
+            add_edge!(bn, :A, :B)
+            add_edge!(bn, :B, :C)
+
+            @test is_conditionally_independent(bn, :A, :C, [:B])
+            @test !is_conditionally_independent(bn, :A, :C, Symbol[])
+        end
+
+        @testset "Fork Structure (A ← B → C)" begin
+            println("\nTesting Fork Structure")
+            bn = BayesianNetwork{Symbol}()
+
+            add_stochastic_vertex!(bn, :A, Normal(), false)
+            add_stochastic_vertex!(bn, :B, Normal(), false)
+            add_stochastic_vertex!(bn, :C, Normal(), false)
+
+            add_edge!(bn, :B, :A)
+            add_edge!(bn, :B, :C)
+
+            println("Graph structure:")
+            println("Edges: ", collect(edges(bn.graph)))
 
-    @testset "Bayes Ball" begin end
+            result = is_conditionally_independent(bn, :A, :C, Symbol[])
+            println("Result for A ⊥ C | ∅: $result")
+        end
+
+        @testset "Collider Structure (A → B ← C)" begin
+            bn = BayesianNetwork{Symbol}()
+
+            add_stochastic_vertex!(bn, :A, Normal(), false)
+            add_stochastic_vertex!(bn, :B, Normal(), false)
+            add_stochastic_vertex!(bn, :C, Normal(), false)
+
+            add_edge!(bn, :A, :B)
+            add_edge!(bn, :C, :B)
+
+            @test is_conditionally_independent(bn, :A, :C, Symbol[])
+            @test !is_conditionally_independent(bn, :A, :C, [:B])
+        end
+
+        @testset "Complex Structure" begin
+            bn = BayesianNetwork{Symbol}()
+
+            for v in [:A, :B, :C, :D, :E]
+                add_stochastic_vertex!(bn, v, Normal(), false)
+            end
+
+            # Create structure:
+            #     A → B → D
+            #         ↓   ↑
+            #         C → E
+            add_edge!(bn, :A, :B)
+            add_edge!(bn, :B, :C)
+            add_edge!(bn, :B, :D)
+            add_edge!(bn, :C, :E)
+            add_edge!(bn, :E, :D)
+
+            @test is_conditionally_independent(bn, :A, :E, [:B, :C])
+            @test !is_conditionally_independent(bn, :A, :E, [:B])
+            @test !is_conditionally_independent(bn, :A, :E, Symbol[])
+        end
+
+        @testset "Using Observed Variables" begin
+            bn = BayesianNetwork{Symbol}()
+
+            add_stochastic_vertex!(bn, :A, Normal(), false)
+            add_stochastic_vertex!(bn, :B, Normal(), true)  # B is observed
+            add_stochastic_vertex!(bn, :C, Normal(), false)
+
+            add_edge!(bn, :A, :B)
+            add_edge!(bn, :B, :C)
+
+            @test is_conditionally_independent(bn, :A, :C)
+
+            bn_decond = decondition(bn)
+            @test !is_conditionally_independent(bn_decond, :A, :C)
+        end
+
+        @testset "Error Handling" begin
+            bn = BayesianNetwork{Symbol}()
+
+            add_stochastic_vertex!(bn, :A, Normal(), false)
+            add_stochastic_vertex!(bn, :B, Normal(), false)
+
+            @test_throws KeyError is_conditionally_independent(bn, :A, :NonExistent)
+            @test_throws KeyError is_conditionally_independent(bn, :NonExistent, :B)
+            @test_throws KeyError is_conditionally_independent(bn, :A, :B, [:NonExistent])
+        end
+    end
 end