From cae89147f059f779e53500323893fd7af0dc1824 Mon Sep 17 00:00:00 2001
From: Naseweisssss <ryanlin9309@gmail.com>
Date: Sun, 17 Nov 2024 08:33:37 +0000
Subject: [PATCH] added test case for the corner case X/Y in Z

---
 .../ProbabilisticGraphicalModels/bayesnet.jl  |  5 +++++
 .../ProbabilisticGraphicalModels/bayesnet.jl  | 21 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl b/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl
index 2f12e63c..82a4babe 100644
--- a/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl
+++ b/src/experimental/ProbabilisticGraphicalModels/bayesnet.jl
@@ -215,6 +215,11 @@ end
 function is_conditionally_independent(
     bn::BayesianNetwork{V}, X::V, Y::V, Z::Vector{V}
 ) where {V}
+    println("debugging: X: $X, Y: $Y, Z: $Z")
+    if X in Z || Y in Z
+        return true
+    end
+
     # Get vertex IDs
     x_id = bn.names_to_ids[X]
     y_id = bn.names_to_ids[Y]
diff --git a/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl b/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl
index a99035ae..2ce48b45 100644
--- a/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl
+++ b/test/experimental/ProbabilisticGraphicalModels/bayesnet.jl
@@ -214,6 +214,27 @@ using JuliaBUGS.ProbabilisticGraphicalModels:
             @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}()