Merge pull request #78 from TensorBFS/jg/fix-overflow

Fix the overflow issue in probability

Author: mroavi

docs/make.jl

@@ -67,6 +67,7 @@ makedocs(;
         "Contributing" => "contributing.md",
     ],
     doctest = false,
+    warnonly = :missing_docs,
 )
 
 deploydocs(;

docs/src/api/public.md

@@ -66,4 +66,6 @@ read_model_file
 read_evidence_file
 read_td_file
 sample
+update_evidence!
+update_temperature
 ```

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)

src/mar.jl

@@ -140,7 +140,7 @@ tensor network.
 - `rescale`: Specifies whether to rescale the tensors during contraction.
 
 ### Example
-The following example is taken from [`examples/asia-network/main.jl`](@ref).
+The following example is taken from [`examples/asia-network/main.jl`](https://tensorbfs.github.io/TensorInference.jl/dev/generated/asia-network/main/).
 
 ```jldoctest; setup = :(using TensorInference, Random; Random.seed!(0))
 julia> model = read_model_file(pkgdir(TensorInference, "examples", "asia-network", "model.uai"));

test/pr.jl

@@ -6,25 +6,35 @@ using TensorInference
 @testset "UAI Reference Solution Comparison" begin
     problems = dataset_from_artifact("uai2014")["PR"]
     problem_sets = [
-        #("Alchemy", TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)), # fails
+        #("Alchemy", TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)),
         #("CSP", TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)),
         #("DBN", KaHyParBipartite(sc_target = 25)),
-        #("Grids", TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)), # fails
-        #("linkage", TreeSA(ntrials = 3, niters = 20, βs = 0.1:0.1:40)), # fails
+        #("Grids", TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)),
+        #("linkage", TreeSA(ntrials = 3, niters = 20, βs = 0.1:0.1:40)),
         #("ObjectDetection", TreeSA(ntrials = 1, niters = 5, βs = 1:0.1:100)),
         ("Pedigree", TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)),
         #("Promedus", TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)),
-        #("relational", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)), # fails
+        #("relational", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)), tw too large
         ("Segmentation", TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100))
     ]
     for (problem_set_name, optimizer) in problem_sets
         @testset "$(problem_set_name) problem set" begin
             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
-                @test isapprox(solution, read_solution(problem); atol = 1e-3)
+                solution = log_probability(tn) / log(10) |> first
+                @test isapprox(solution, read_solution(problem); atol = 1e-3, rtol=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