fix the overflow issue in probability

Author: GiggleLiu

src/Core.jl

@@ -203,6 +203,16 @@ function log_probability(tn::TensorNetworkModel, config::Union{Dict, AbstractVec
     assign = config isa AbstractVector ? Dict(zip(get_vars(tn), config)) : config
     return sum(x -> log(x[2][(getindex.(Ref(assign), x[1]) .+ 1)...]), zip(getixsv(tn.code), tn.tensors))
 end
+"""
+$(TYPEDSIGNATURES)
+
+Evaluate the log probability (or partition function).
+It is the logged version of [`probability`](@ref), which is less likely to overflow.
+"""
+function log_probability(tn::TensorNetworkModel; usecuda = false)::AbstractArray
+    res = probability(tn; usecuda, rescale=true)
+    return asarray(res.log_factor .+ log.(res.normalized_value), res.normalized_value)
+end
 
 """
 $(TYPEDSIGNATURES)

test/pr.jl

@@ -22,9 +22,19 @@ using TensorInference
             for (id, problem) in problems[problem_set_name]
                 @info "Testing: $(problem_set_name)_$id"
                 tn = TensorNetworkModel(read_model(problem); optimizer, evidence=read_evidence(problem))
-                solution = probability(tn) |> first |> log10
+                solution = log_probability(tn) / log(10) |> first
                 @test isapprox(solution, read_solution(problem); atol = 1e-3)
             end
         end
     end
 end
+
+@testset "issue 77" begin
+    problems = dataset_from_artifact("uai2014")["PR"]
+    problem_set_name = "Alchemy"
+    optimizer = TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)
+    id, problem = problems[problem_set_name] |> first
+    tn = TensorNetworkModel(read_model(problem); optimizer, evidence=read_evidence(problem))
+    solution = log_probability(tn) / log(10) |> first
+    @test isapprox(solution, read_solution(problem); atol=1e-3)
+end