implement marginals

Author: GiggleLiu

src/belief.jl

@@ -110,4 +110,8 @@ end
 # if BP is exact and converged (e.g. tree like), the result should be the same as the tensor network contraction
 function contraction_results(state::BPState{T}) where T
     return [sum(reduce((x, y) -> x .* y, mi)) for mi in state.message_in]
+end
+
+function marginals(state::BPState{T}) where T
+    return Dict([v] => normalize!(reduce((x, y) -> x .* y, mi), 1) for (v, mi) in enumerate(state.message_in))
 end

test/belief.jl

@@ -49,4 +49,9 @@ end
     tnet = TensorNetworkModel(mps_uai)
     expected_result = probability(tnet)[]
     @test all(r -> isapprox(r, expected_result), contraction_res)
+    mars = marginals(state)
+    mars_tnet = marginals(tnet)
+    for v in 1:TensorInference.num_variables(bp)
+        @test mars[[v]] ≈ mars_tnet[[v]]
+    end
 end