fix tests

Author: GiggleLiu

src/Core.jl

@@ -37,7 +37,7 @@ struct UAIInstance{ET, FT <: Factor{ET}}
     obsvars::Vector{Int}
     obsvals::Vector{Int}
     queryvars::Vector{Int}
-    reference_solution::Union{Vector{Vector{ET}}, Vector{Int}, Float64}
+    reference_solution
 end
 
 Base.show(io::IO, ::MIME"text/plain", uai::UAIInstance) = Base.show(io, uai)
@@ -69,6 +69,17 @@ function set_evidence!(uai::UAIInstance, pairs::Pair{Int}...)
     return uai
 end
 
+"""
+$TYPEDSIGNATURES
+
+Set the query variables of an UAI instance.
+"""
+function set_query!(uai::UAIInstance, vars::AbstractVector{Int})
+    empty!(uai.queryvars)
+    append!(uai.queryvars, vars)
+    return uai
+end
+
 """
 $(TYPEDEF)
 
@@ -122,6 +133,9 @@ function TensorNetworkModel(
     optimizer = GreedyMethod(),
     simplifier = nothing
 )::TensorNetworkModel
+    if !isempty(instance.queryvars)
+        @warn "The `queryvars` field of the input `UAIInstance` instance is designed for the `MMAPModel`, which is not respected by `TensorNetworkModel`. Got non-empty value: $(uai.queryvars)"
+    end
     return TensorNetworkModel(
         1:(instance.nvars),
         instance.cards,

src/TensorInference.jl

@@ -18,7 +18,7 @@ export contraction_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABiparti
 
 # read and load uai files
 export read_model_file, read_td_file, read_evidence_file, read_solution_file, read_instance, UAIInstance
-export set_evidence!
+export set_evidence!, set_query!
 
 # marginals
 export TensorNetworkModel, get_vars, get_cards, log_probability, probability, marginals

src/mmap.jl

@@ -58,9 +58,9 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function MMAPModel(instance::UAIInstance; queryvars, openvertices = (), optimizer = GreedyMethod(), simplifier = nothing)::MMAPModel
+function MMAPModel(instance::UAIInstance; openvertices = (), optimizer = GreedyMethod(), simplifier = nothing)::MMAPModel
     return MMAPModel(
-        1:(instance.nvars), instance.cards, instance.factors; queryvars, fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)), optimizer, simplifier, openvertices
+        1:(instance.nvars), instance.cards, instance.factors; queryvars=instance.queryvars, fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)), optimizer, simplifier, openvertices
     )
 end
 

src/utils.jl

@@ -246,7 +246,7 @@ end
 
 function read_instance_from_string(uai::AbstractString; eltype = Float64)::UAIInstance
     nvars, cards, ncliques, factors = read_model_string(uai; factor_eltype = eltype)
-    return UAIInstance(nvars, ncliques, cards, factors, Int[], Int[], Vector{eltype}[])
+    return UAIInstance(nvars, ncliques, cards, factors, Int[], Int[], Int[], nothing)
 end
 
 # patch to get content by broadcasting into array, while keep array size unchanged.

test/cuda.jl

@@ -42,22 +42,25 @@ end
     optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40)
     tn_ref = TensorNetworkModel(instance; optimizer)
     # does not marginalize any var
-    tn = MMAPModel(instance; queryvars = collect(1:instance.nvars), optimizer)
+    set_query!(instance, collect(1:instance.nvars))
+    tn = MMAPModel(instance; optimizer)
     r1, r2 = maximum_logp(tn_ref; usecuda = true), maximum_logp(tn; usecuda = true)
     @test r1 isa CuArray
     @test r2 isa CuArray
     @test r1 ≈ r2
 
     # marginalize all vars
-    tn2 = MMAPModel(instance; queryvars = Int[], optimizer)
+    set_query!(instance, Int[])
+    tn2 = MMAPModel(instance; optimizer)
     cup = probability(tn_ref; usecuda = true)
     culogp = maximum_logp(tn2; usecuda = true)
     @test cup isa RescaledArray{T, N, <:CuArray} where {T, N}
     @test culogp isa CuArray
     @test Array(cup)[] ≈ exp(Array(culogp)[])
 
     # does not optimize over open vertices
-    tn3 = MMAPModel(instance; queryvars = setdiff(1:instance.nvars, [2, 4, 6]), optimizer)
+    set_query!(instance, setdiff(1:instance.nvars, [2, 4, 6]))
+    tn3 = MMAPModel(instance; optimizer)
     logp, config = most_probable_config(tn3; usecuda = true)
     @test log_probability(tn3, config) ≈ logp
 end

test/mmap.jl

@@ -15,17 +15,20 @@ end
     tn_ref = TensorNetworkModel(instance; optimizer)
 
     # Does not marginalize any var
-    mmap = MMAPModel(instance; queryvars = collect(1:instance.nvars), optimizer)
+    set_query!(instance, collect(1:instance.nvars))
+    mmap = MMAPModel(instance; optimizer)
     @debug(mmap)
     @test maximum_logp(tn_ref) ≈ maximum_logp(mmap)
 
     # Marginalize all vars
-    mmap2 = MMAPModel(instance; queryvars = Int[], optimizer)
+    set_query!(instance, Int[])
+    mmap2 = MMAPModel(instance; optimizer)
     @debug(mmap2)
     @test Array(probability(tn_ref))[] ≈ exp(maximum_logp(mmap2)[])
 
     # Does not optimize over open vertices
-    mmap3 = MMAPModel(instance; queryvars = setdiff(1:instance.nvars, [2, 4, 6]), optimizer)
+    set_query!(instance, setdiff(1:instance.nvars, [2, 4, 6]))
+    mmap3 = MMAPModel(instance; optimizer)
     @debug(mmap3)
     logp, config = most_probable_config(mmap3)
     @test log_probability(mmap3, config) ≈ logp
@@ -42,7 +45,7 @@ end
       @info "Testing: $problem_name"
       model_filepath, evidence_filepath, query_filepath, solution_filepath = get_instance_filepaths(problem_name, "MMAP")
       instance = read_instance(model_filepath; evidence_filepath, query_filepath, solution_filepath)
-      model = MMAPModel(instance; queryvars = instance.queryvars, optimizer)
+      model = MMAPModel(instance; optimizer)
       _, solution = most_probable_config(model)
       @test solution == instance.reference_solution
     end

test/runtests.jl

@@ -21,6 +21,10 @@ end
     include("pr.jl")
 end
 
+@testset "sampling" begin
+    include("sampling.jl")
+end
+
 using CUDA
 if CUDA.functional()
     include("cuda.jl")