From f07d31b6a2f8ca7fcf13f8df3f213b3edd075453 Mon Sep 17 00:00:00 2001 From: Naseweisssss <89737029+naseweisssss@users.noreply.github.com> Date: Thu, 21 Nov 2024 12:15:55 +0000 Subject: [PATCH] Ancestral Sampling and Bayes Ball Algorithm (#233) Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Xianda Sun <5433119+sunxd3@users.noreply.github.com> Co-authored-by: Xianda Sun --- .../ProbabilisticGraphicalModels.jl | 11 + .../ProbabilisticGraphicalModels/bayesnet.jl | 161 +++++++++++++- .../ProbabilisticGraphicalModels/bayesnet.jl | 209 +++++++++++++++++- 3 files changed, 366 insertions(+), 15 deletions(-) diff --git a/src/experimental/ProbabilisticGraphicalModels/ProbabilisticGraphicalModels.jl b/src/experimental/ProbabilisticGraphicalModels/ProbabilisticGraphicalModels.jl index 4785bf151..29c261abc 100644 --- a/src/experimental/ProbabilisticGraphicalModels/ProbabilisticGraphicalModels.jl +++ b/src/experimental/ProbabilisticGraphicalModels/ProbabilisticGraphicalModels.jl @@ -6,4 +6,15 @@ using Distributions include("bayesnet.jl") +export BayesianNetwork, + add_stochastic_vertex!, + add_deterministic_vertex!, + add_edge!, + condition, + condition!, + decondition, + decondition!, + ancestral_sampling, + is_conditionally_independent + end diff --git a/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl b/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl index 063a17a77..3eca2da56 100644 --- a/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl +++ b/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}(), @@ -164,31 +164,168 @@ Perform ancestral sampling on a Bayesian network to generate one sample from the Ancestral sampling works by: 1. Finding a topological ordering of the nodes 2. Sampling from each node in order, using the already-sampled parent values for conditional distributions + +### Return Value +The function returns a `Dict{V, Any}` where: +- Each key is a variable name (of type `V`) in the Bayesian Network. +- Each value is the sampled value for that variable, which can be of any type (`Any`). + +This dictionary represents a single sample from the joint distribution of the Bayesian Network, capturing the dependencies and conditional relationships defined in the network structure. + """ 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 """ - is_conditionally_independent(bn::BayesianNetwork, X::V, Y::V[, Z::Vector{V}]) where {V} + is_conditionally_independent(bn::BayesianNetwork, X::Vector{V}, Y::Vector{V}, Z::Vector{V}) where {V} + +Test whether sets of variables X and Y are conditionally independent given set Z in a Bayesian Network using the Bayes Ball algorithm. + +# Arguments +- `bn::BayesianNetwork`: The Bayesian Network structure +- `X::Vector{V}`: First set of variables to test for independence +- `Y::Vector{V}`: Second set of variables to test for independence +- `Z::Vector{V}`: Set of conditioning variables (can be empty) -Determines if two variables X and Y are conditionally independent given the conditioning information already known. -If Z is provided, the conditioning information in `bn` will be ignored. +# Returns +- `true`: if X and Y are conditionally independent given Z (X ⊥ Y | Z) +- `false`: if X and Y are conditionally dependent given Z + +# Description +The Bayes Ball algorithm determines conditional independence by checking if there exists an active path between +variables in X and Y given Z. The algorithm follows these rules: +- In a chain (A → B → C): B blocks the path if conditioned +- In a fork (A ← B → C): B blocks the path if conditioned +- In a collider (A → B ← C): B opens the path if conditioned +# Examples +``` """ -function is_conditionally_independent end +function is_conditionally_independent( + bn::BayesianNetwork{V}, X::Vector{V}, Y::Vector{V}, Z::Vector{V} +) where {V} + isempty(X) && throw(ArgumentError("X cannot be empty")) + isempty(Y) && throw(ArgumentError("Y cannot be empty")) + + x_ids = Set([bn.names_to_ids[x] for x in X]) + y_ids = Set([bn.names_to_ids[y] for y in Y]) + z_ids = Set([bn.names_to_ids[z] for z in Z]) + + # Check if any variable in X or Y is in Z + if !isempty(intersect(x_ids, z_ids)) || !isempty(intersect(y_ids, z_ids)) + return true + end + + # Add observed variables to conditioning set + for (id, is_obs) in enumerate(bn.is_observed) + if is_obs + push!(z_ids, id) + end + end + + # Track visited nodes and their directions + n_vertices = nv(bn.graph) + visited_up = falses(n_vertices) # Visited going up (from child to parent) + visited_down = falses(n_vertices) # Visited going down (from parent to child) + + # Queue entries are (node_id, going_up) + queue = Tuple{Int,Bool}[] + + # Start from all X nodes + for x_id in x_ids + push!(queue, (x_id, true)) # Try going up + push!(queue, (x_id, false)) # Try going down + end + + while !isempty(queue) + current_id, going_up = popfirst!(queue) + + # Skip if we've visited this node in this direction + if (going_up && visited_up[current_id]) || (!going_up && visited_down[current_id]) + continue + end + + # Mark as visited in current direction + if going_up + visited_up[current_id] = true + else + visited_down[current_id] = true + end + + # If we reached a Y node, path is active + if current_id in y_ids + return false + end + + is_conditioned = current_id in z_ids + parents = inneighbors(bn.graph, current_id) + children = outneighbors(bn.graph, current_id) + + if is_conditioned + # If conditioned: + # - In a chain/fork: blocks the path + # - In a collider or descendant of collider: allows going up to parents + if length(parents) > 1 || !isempty(children) # Is collider or has children + for parent in parents + push!(queue, (parent, true)) # Can only go up to parents + end + end + else + # If not conditioned: + if going_up + # Going up: can visit parents + for parent in parents + push!(queue, (parent, true)) + end + else + # Going down: can visit children + for child in children + push!(queue, (child, false)) + end + end + + # At starting nodes (X), we can go both up and down + if current_id in x_ids + if going_up + for child in children + push!(queue, (child, false)) + end + else + for parent in parents + push!(queue, (parent, true)) + end + end + end + end + end -function is_conditionally_independent(bn::BayesianNetwork{V}, X::V, Y::V) where {V} - # TODO: Implement + return true end +# Single variable version with Z function is_conditionally_independent( bn::BayesianNetwork{V}, X::V, Y::V, Z::Vector{V} ) where {V} - # TODO: Implement + return is_conditionally_independent(bn, [X], [Y], Z) end diff --git a/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl b/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl index d5424242f..947a10655 100644 --- a/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl +++ b/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl @@ -7,7 +7,10 @@ using JuliaBUGS.ProbabilisticGraphicalModels: add_deterministic_vertex!, add_edge!, condition, - decondition + condition!, + decondition, + ancestral_sampling, + is_conditionally_independent @testset "BayesianNetwork" begin @testset "Adding vertices" begin @@ -96,7 +99,207 @@ 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, Symbol[]) + @test is_conditionally_independent(bn, :A, :C, [:B]) + 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) - @testset "Bayes Ball" begin end + 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 "Bayes Ball Algorithm Tests" begin + # Create a simple network: A → B → C + bn = BayesianNetwork{Symbol}() + add_stochastic_vertex!(bn, :A, Normal(0, 1), false) + add_stochastic_vertex!(bn, :B, Normal(0, 1), false) + add_stochastic_vertex!(bn, :C, Normal(0, 1), false) + add_edge!(bn, :A, :B) + add_edge!(bn, :B, :C) + @testset "Corner Case: X or Y in Z" begin + # Test case where X is in Z + @test is_conditionally_independent(bn, :A, :C, [:A]) # A ⊥ C | A + # Test case where Y is in Z + @test is_conditionally_independent(bn, :A, :C, [:C]) # A ⊥ C | C + # Test case where both X and Y are in Z + @test is_conditionally_independent(bn, :A, :C, [:A, :C]) # A ⊥ C | A, C + end + 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, Symbol[]) + 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, :B, [:NonExistent]) + end + end + + @testset "Conditional Independence (Single and Multiple Variables)" begin + bn = BayesianNetwork{Symbol}() + + # Create a complex network + # A → B → D + # ↓ ↓ ↑ + # C → E → F + for v in [:A, :B, :C, :D, :E, :F] + add_stochastic_vertex!(bn, v, Normal(), false) + end + + add_edge!(bn, :A, :B) + add_edge!(bn, :A, :C) + add_edge!(bn, :B, :D) + add_edge!(bn, :B, :E) + add_edge!(bn, :C, :E) + add_edge!(bn, :E, :F) + add_edge!(bn, :F, :D) + + # Test single variable independence + @test is_conditionally_independent(bn, :A, :F, [:B, :C]) + @test !is_conditionally_independent(bn, :A, :D, Symbol[]) + + # Test multiple variable independence + @test is_conditionally_independent(bn, [:A], [:F], [:B, :C]) + @test is_conditionally_independent(bn, [:A, :B], [:F], [:C, :E]) + @test !is_conditionally_independent(bn, [:A, :B], [:D, :F], Symbol[]) + + # Test when some variables ccdare in conditioning set + @test is_conditionally_independent(bn, [:A, :B], [:D, :F], [:A]) + @test is_conditionally_independent(bn, [:A, :B], [:D, :F], [:F]) + + # Test error handling + @test_throws KeyError is_conditionally_independent( + bn, [:A, :NonExistent], [:F], [:B] + ) + @test_throws KeyError is_conditionally_independent( + bn, [:A], [:F, :NonExistent], [:B] + ) + @test_throws KeyError is_conditionally_independent(bn, [:A], [:F], [:NonExistent]) + @test_throws ArgumentError is_conditionally_independent(bn, Symbol[], [:F], [:B]) + @test_throws ArgumentError is_conditionally_independent(bn, [:A], Symbol[], [:B]) + end end