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,94 @@ 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, can go both up and down initially
+    push!(queue, (x_id, true))  # As if coming from a parent
+
+    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
+
+        # Get neighbors
+        parents = inneighbors(bn.graph, current_id)
+        children = outneighbors(bn.graph, current_id)
+
+        # Rule 1: If coming from parent and not conditioned, can go to children
+        if from_parent && !is_conditioned
+            append!(queue, [(child, true) for child in children])
+        end
+
+        # Rule 2: If coming from child and not conditioned, can go to parents
+        if !from_parent && !is_conditioned
+            append!(queue, [(parent, false) for parent in parents])
+        end
+
+        # Rule 3: If at a collider (or descendant of collider) and it's conditioned, 
+        # can go up to parents
+        if !from_parent && (is_conditioned || has_conditioned_descendant(bn, current_id, z_ids))
+            if length(parents) > 1  # Is a collider
+                append!(queue, [(parent, false) for parent in parents])
+            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
+
+        # Check if current node is conditioned
+        if current in z_ids
+            return true
+        end
+
+        # Add all unvisited children to queue
+        for child in outneighbors(bn.graph, current)
+            if !visited[child]
+                push!(queue, child)
+            end
+        end
+    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,165 @@ 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
+            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)
+
+            @test is_conditionally_independent(bn, :A, :C, [:B])
+            @test !is_conditionally_independent(bn, :A, :C, Symbol[])
+        end
 
-    @testset "Bayes Ball" begin 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