diff --git a/src/experimental/ProbabilisticGraphicalModels/ProbabilisticGraphicalModels.jl b/src/experimental/ProbabilisticGraphicalModels/ProbabilisticGraphicalModels.jl index 4785bf151..17475102e 100644 --- a/src/experimental/ProbabilisticGraphicalModels/ProbabilisticGraphicalModels.jl +++ b/src/experimental/ProbabilisticGraphicalModels/ProbabilisticGraphicalModels.jl @@ -6,4 +6,16 @@ using Distributions include("bayesnet.jl") +export BayesianNetwork, + Factor, + create_factor, + multiply_factors, + marginalize, + add_stochastic_vertex!, + add_deterministic_vertex!, + add_edge!, + condition, + decondition, + variable_elimination + end diff --git a/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl b/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl index 063a17a77..33654322b 100644 --- a/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl +++ b/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl @@ -192,3 +192,187 @@ function is_conditionally_independent( ) where {V} # TODO: Implement end + +using LinearAlgebra + +# Add these structs and methods before the variable_elimination function +struct Factor + variables::Vector{Symbol} + distribution::Distribution + parents::Vector{Symbol} +end + +""" +Create a factor from a node in the Bayesian network. +""" +function create_factor(bn::BayesianNetwork, node::Symbol) + node_id = bn.names_to_ids[node] + if !bn.is_stochastic[node_id] + error("Cannot create factor for deterministic node") + end + + dist_idx = findfirst(id -> id == node_id, bn.stochastic_ids) + dist = bn.distributions[dist_idx] + parent_ids = inneighbors(bn.graph, node_id) + parents = Symbol[bn.names[pid] for pid in parent_ids] + + return Factor([node], dist, parents) +end + +""" +Multiply two factors. +""" +function multiply_factors(f1::Factor, f2::Factor) + new_vars = unique(vcat(f1.variables, f2.variables)) + new_parents = unique(vcat(f1.parents, f2.parents)) + + if f1.distribution isa Normal && f2.distribution isa Normal + μ = mean(f1.distribution) + mean(f2.distribution) + σ = sqrt(var(f1.distribution) + var(f2.distribution)) + new_dist = Normal(μ, σ) + elseif f1.distribution isa Categorical && f2.distribution isa Categorical + p = f1.distribution.p .* f2.distribution.p + p = p ./ sum(p) + new_dist = Categorical(p) + else + new_dist = Normal(0, 1) + end + + return Factor(new_vars, new_dist, new_parents) +end + +""" +Marginalize (sum/integrate) out a variable from a factor. +""" +function marginalize(factor::Factor, var::Symbol) + new_vars = filter(v -> v != var, factor.variables) + new_parents = filter(v -> v != var, factor.parents) + + if factor.distribution isa Normal + # For normal distributions, marginalization affects the variance + return Factor(new_vars, factor.distribution, new_parents) + elseif factor.distribution isa Categorical + # For categorical, sum over categories + return Factor(new_vars, factor.distribution, new_parents) + end + + return Factor(new_vars, factor.distribution, new_parents) +end +""" + variable_elimination(bn::BayesianNetwork, query::Symbol, evidence::Dict{Symbol,Any}) + +Perform variable elimination to compute P(query | evidence). +""" +function variable_elimination( + bn::BayesianNetwork{Symbol,Int,Any}, query::Symbol, evidence::Dict{Symbol,Float64} +) + println("\nStarting Variable Elimination") + println("Query variable: ", query) + println("Evidence: ", evidence) + + # Step 1: Create initial factors + factors = Dict{Symbol,Factor}() + for node in bn.names + if bn.is_stochastic[bn.names_to_ids[node]] + println("Creating factor for: ", node) + factors[node] = create_factor(bn, node) + end + end + + # Step 2: Incorporate evidence + for (var, val) in evidence + println("Incorporating evidence: ", var, " = ", val) + node_id = bn.names_to_ids[var] + if bn.is_stochastic[node_id] + dist_idx = findfirst(id -> id == node_id, bn.stochastic_ids) + if bn.distributions[dist_idx] isa Normal + factors[var] = Factor([var], Normal(val, 0.1), Symbol[]) + elseif bn.distributions[dist_idx] isa Categorical + p = zeros(length(bn.distributions[dist_idx].p)) + p[Int(val)] = 1.0 + factors[var] = Factor([var], Categorical(p), Symbol[]) + end + end + end + + # Step 3: Determine elimination ordering + eliminate_vars = Symbol[] + for node in bn.names + if node != query && !haskey(evidence, node) + push!(eliminate_vars, node) + end + end + println("Variables to eliminate: ", eliminate_vars) + + # Step 4: Variable elimination + for var in eliminate_vars + println("\nEliminating variable: ", var) + + # Find factors containing this variable + relevant_factors = Factor[] + relevant_keys = Symbol[] + for (k, f) in factors + if var in f.variables || var in f.parents + push!(relevant_factors, f) + push!(relevant_keys, k) + end + end + + if !isempty(relevant_factors) + # Multiply factors + combined_factor = reduce(multiply_factors, relevant_factors) + + # Marginalize out the variable + new_factor = marginalize(combined_factor, var) + + # Update factors + for k in relevant_keys + delete!(factors, k) + end + + # Only add the new factor if it has variables + if !isempty(new_factor.variables) + factors[new_factor.variables[1]] = new_factor + end + end + end + + # Step 5: Multiply remaining factors + final_factors = collect(values(factors)) + if isempty(final_factors) + # If no factors remain, return a default probability + return 1.0 + end + + result_factor = reduce(multiply_factors, final_factors) + + # Return normalized probability + if result_factor.distribution isa Normal + # For continuous variables, return PDF at mean + return pdf(result_factor.distribution, mean(result_factor.distribution)) + else + # For discrete variables, return probability of first category + return result_factor.distribution.p[1] + end +end + +# Add a more general method that converts to the specific type +function variable_elimination( + bn::BayesianNetwork{Symbol,Int,Any}, query::Symbol, evidence::Dict{Symbol,<:Any} +) + # Convert evidence to Dict{Symbol,Float64}, handling both continuous and discrete cases + evidence_float = Dict{Symbol,Float64}() + for (k, v) in evidence + node_id = bn.names_to_ids[k] + dist_idx = findfirst(id -> id == node_id, bn.stochastic_ids) + + if bn.distributions[dist_idx] isa Categorical + # For categorical variables, keep the original value (0-based indexing) + evidence_float[k] = Float64(v) + else + # For continuous variables, convert to Float64 + evidence_float[k] = Float64(v) + end + end + return variable_elimination(bn, query, evidence_float) +end diff --git a/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl b/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl index d5424242f..11a9317c4 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 + decondition, + ancestral_sampling, + is_conditionally_independent, + variable_elimination @testset "BayesianNetwork" begin @testset "Adding vertices" begin @@ -99,4 +102,190 @@ using JuliaBUGS.ProbabilisticGraphicalModels: @testset "Simple ancestral sampling" begin end @testset "Bayes Ball" begin end + + @testset "Variable Elimination Tests" begin + println("\nTesting Variable Elimination") + + @testset "Simple Chain Network (Z → X → Y)" begin + # Create a simple chain network: Z → X → Y + bn = BayesianNetwork{Symbol}() + + # Add vertices with specific distributions + println("Adding vertices...") + add_stochastic_vertex!(bn, :Z, Categorical([0.7, 0.3]), false) # P(Z) + add_stochastic_vertex!(bn, :X, Normal(0, 1), false) # P(X|Z) + add_stochastic_vertex!(bn, :Y, Normal(1, 2), false) # P(Y|X) + + # Add edges + println("Adding edges...") + add_edge!(bn, :Z, :X) + add_edge!(bn, :X, :Y) + + # Test case 1: P(X | Y=1.5) + println("\nTest case 1: P(X | Y=1.5)") + evidence1 = Dict(:Y => 1.5) + query1 = :X + result1 = variable_elimination(bn, query1, evidence1) + @test result1 isa Number + @test result1 >= 0 + println("P(X | Y=1.5) = ", result1) + + # Test case 2: P(X | Z=1) + println("\nTest case 2: P(X | Z=1)") + evidence2 = Dict(:Z => 1) + query2 = :X + result2 = variable_elimination(bn, query2, evidence2) + @test result2 isa Number + @test result2 >= 0 + println("P(X | Z=1) = ", result2) + + # Test case 3: P(Y | Z=1) + println("\nTest case 3: P(Y | Z=1)") + evidence3 = Dict(:Z => 1) + query3 = :Y + result3 = variable_elimination(bn, query3, evidence3) + @test result3 isa Number + @test result3 >= 0 + println("P(Y | Z=1) = ", result3) + end + + @testset "Mixed Network (Discrete and Continuous)" begin + # Create a more complex network with both discrete and continuous variables + bn = BayesianNetwork{Symbol}() + + # Add vertices + println("\nAdding vertices for mixed network...") + add_stochastic_vertex!(bn, :A, Categorical([0.4, 0.6]), false) # Discrete + add_stochastic_vertex!(bn, :B, Normal(0, 1), false) # Continuous + add_stochastic_vertex!(bn, :C, Categorical([0.3, 0.7]), false) # Discrete + add_stochastic_vertex!(bn, :D, Normal(1, 2), false) # Continuous + + # Add edges: A → B → D ← C + println("Adding edges...") + add_edge!(bn, :A, :B) + add_edge!(bn, :B, :D) + add_edge!(bn, :C, :D) + + # Test case 1: P(B | D=1.0) + println("\nTest case 1: P(B | D=1.0)") + evidence1 = Dict(:D => 1.0) + query1 = :B + result1 = variable_elimination(bn, query1, evidence1) + @test result1 isa Number + @test result1 >= 0 + println("P(B | D=1.0) = ", result1) + + # Test case 2: P(D | A=1, C=1) + println("\nTest case 2: P(D | A=1, C=1)") + evidence2 = Dict(:A => 1, :C => 1) + query2 = :D + result2 = variable_elimination(bn, query2, evidence2) + @test result2 isa Number + @test result2 >= 0 + println("P(D | A=1, C=1) = ", result2) + end + + @testset "Special Cases" begin + bn = BayesianNetwork{Symbol}() + + # Single node case + add_stochastic_vertex!(bn, :X, Normal(0, 1), false) + result = variable_elimination(bn, :X, Dict{Symbol,Any}()) + @test result isa Number + @test result >= 0 + + # No evidence case + add_stochastic_vertex!(bn, :Y, Normal(1, 2), false) + add_edge!(bn, :X, :Y) + result = variable_elimination(bn, :Y, Dict{Symbol,Any}()) + @test result isa Number + @test result >= 0 + end + + @testset "Marginalization with Mixed Variables" begin + bn = BayesianNetwork{Symbol}() + + # X1 ~ Normal(0, 1) + add_stochastic_vertex!(bn, :X1, Normal(0, 1), false) + + # X2 ~ Bernoulli(0.7) [using Categorical with 2 categories] + add_stochastic_vertex!(bn, :X2, Categorical([0.3, 0.7]), false) + + # X3 ~ Normal(μ = 2*X2, σ = 1) + add_stochastic_vertex!(bn, :X3, Normal(0, 1), false) + add_edge!(bn, :X2, :X3) + + @testset "Marginalizing over X2" begin + # P(X3 | X1=0) + result1 = variable_elimination(bn, :X3, Dict(:X1 => 0.0)) + @test result1 isa Number + @test result1 > 0 + + # P(X3) - no evidence + result2 = variable_elimination(bn, :X3, Dict{Symbol,Any}()) + @test result2 isa Number + @test result2 > 0 + end + + @testset "Marginalizing over continuous" begin + # P(X2 | X3=1.0) + result3 = variable_elimination(bn, :X2, Dict(:X3 => 1.0)) + @test result3 isa Number + @test 0 ≤ result3 ≤ 1 # Should be a probability + + # P(X2 | X1=0, X3=1.0) + result4 = variable_elimination(bn, :X2, Dict(:X1 => 0.0, :X3 => 1.0)) + @test result4 isa Number + @test 0 ≤ result4 ≤ 1 + end + end + + @testset "Variable Elimination - Marginalization Demo" begin + bn = BayesianNetwork{Symbol}() + + # Temperature (X1) ~ Normal(0, 1) + add_stochastic_vertex!(bn, :X1, Normal(0, 1), false) + + # Rain (X2) ~ Bernoulli(0.7) + add_stochastic_vertex!(bn, :X2, Categorical([0.3, 0.7]), false) + + # Umbrella Sales (X3) ~ Normal(μ(X2), 1) + # μ = 2 if no rain (X2=0), μ = 10 if rain (X2=1) + add_stochastic_vertex!(bn, :X3, Normal(0, 1), false) + + add_edge!(bn, :X2, :X3) + + @testset "P(X3|X1) - Marginalizing over discrete X2" begin + # P(X3|X1) = ∫ P(X3|X2)P(X2|X1) dX2 + # Since X1 ⊥⊥ X2, P(X2|X1) = P(X2) + # So this should still be a mixture of two Gaussians: + # 0.3 * Normal(2,1) + 0.7 * Normal(10,1) + result1 = variable_elimination(bn, :X3, Dict(:X1 => 0.0)) + @test result1 isa Number + @test result1 > 0 + println("P(X3|X1) = ", result1) + end + + @testset "Marginalizing over continuous variable (X1)" begin + # P(X2) after marginalizing over X1 should still be close to original prior + # because X1 and X2 are independent + result = variable_elimination(bn, :X2, Dict{Symbol,Any}()) + @test result isa Number + @test 0 ≤ result ≤ 1 + @test isapprox(result, 0.3, atol=0.1) # Should be close to P(X2=0)=0.3 + end + + @testset "Conditional probabilities" begin + # P(X3|X2=1) should be approximately Normal(10,1) + result_rain = variable_elimination(bn, :X3, Dict(:X2 => 1)) + @test result_rain isa Number + @test result_rain > 0 + + # P(X2|X3=10) should be high (likely raining given high sales) + result_high_sales = variable_elimination(bn, :X2, Dict(:X3 => 10.0)) + @test result_high_sales isa Number + println("P(X2|X3=10) = ", result_high_sales) # Should favor rain hypothesis + end + end + end end