replace marginalized with queryvars

Author: GiggleLiu

examples/asia/main.jl

@@ -31,7 +31,7 @@ logp, cfg = most_probable_config(tnet)
 
 # Get the maximum log-probabilities (MMAP)
 # To get the probability of lung cancer, we need to marginalize out other variables.
-mmap = MMAPModel(instance; marginalized=[1,2,3,5,6,8])
+mmap = MMAPModel(instance; queryvars=[4,7])
 # We get the most probable configurations on [4, 7]
 most_probable_config(mmap)
 # The total probability of having lung cancer is roughly half.

src/mmap.jl

@@ -17,11 +17,11 @@ Computing the most likely assignment to the query variables,  Xₘ ⊆ X after m
 ```
 
 ### Fields
-* `vars` is the remaining (or not marginalized) degree of freedoms in the tensor network.
+* `vars` is the query variables in the tensor network.
 * `code` is the tropical tensor network contraction pattern.
 * `tensors` is the tensors fed into the tensor network.
 * `clusters` is the clusters, each element of this cluster is a [`TensorNetworkModel`](@ref) instance for marginalizing certain variables.
-* `fixedvertices` is a dictionary to specifiy degree of freedoms fixed to certain values, which should not have overlap with the marginalized variables.
+* `fixedvertices` is a dictionary to specifiy degree of freedoms fixed to certain values, which should not have overlap with the query variables.
 """
 struct MMAPModel{LT, AT <: AbstractArray}
     vars::Vector{LT}
@@ -37,7 +37,7 @@ function Base.show(io::IO, mmap::MMAPModel)
     tc, sc, rw = contraction_complexity(mmap)
     println(io, "$(typeof(mmap))")
     println(io, "variables: $variables")
-    println(io, "marginalized variables: $(map(x->x.eliminated_vars, mmap.clusters))")
+    println(io, "query variables: $(map(x->x.eliminated_vars, mmap.clusters))")
     print_tcscrw(io, tc, sc, rw)
 end
 Base.show(io::IO, ::MIME"text/plain", mmap::MMAPModel) = Base.show(io, mmap)
@@ -58,24 +58,26 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function MMAPModel(instance::UAIInstance; marginalized, openvertices = (), optimizer = GreedyMethod(), simplifier = nothing)::MMAPModel
+function MMAPModel(instance::UAIInstance; queryvars, openvertices = (), optimizer = GreedyMethod(), simplifier = nothing)::MMAPModel
     return MMAPModel(
-        1:(instance.nvars), instance.cards, instance.factors; marginalized, fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)), optimizer, simplifier, openvertices
+        1:(instance.nvars), instance.cards, instance.factors; queryvars, fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)), optimizer, simplifier, openvertices
     )
 end
 
 """
 $(TYPEDSIGNATURES)
 """
-function MMAPModel(vars::AbstractVector{LT}, cards::AbstractVector{Int}, factors::Vector{<:Factor{T}}; marginalized, openvertices = (),
-    fixedvertices = Dict{LT, Int}(),
-    optimizer = GreedyMethod(), simplifier = nothing,
-    marginalize_optimizer = GreedyMethod(), marginalize_simplifier = nothing
-)::MMAPModel where {T, LT}
+function MMAPModel(vars::AbstractVector{LT}, cards::AbstractVector{Int}, factors::Vector{<:Factor{T}}; queryvars, openvertices = (),
+            fixedvertices = Dict{LT, Int}(),
+            optimizer = GreedyMethod(), simplifier = nothing,
+            marginalize_optimizer = GreedyMethod(), marginalize_simplifier = nothing
+        )::MMAPModel where {T, LT}
     all_ixs = [[[var] for var in vars]..., [[factor.vars...] for factor in factors]...]  # labels for vertex tensors (unity tensors) and edge tensors
     iy = collect(LT, openvertices)
-    if !isempty(setdiff(iy, vars))
-        error("Marginalized variables should not contain any output variable.")
+    evidencevars = collect(LT, keys(fixedvertices))
+    marginalized = setdiff(vars, iy ∪ queryvars ∪ evidencevars)
+    if !isempty(setdiff(iy, marginalized))
+        error("Marginalized variables should not contain any output variable, got $(marginalized) and $iy.")
     end
     all_tensors = [[ones(T, cards[i]) for i in 1:length(vars)]..., getfield.(factors, :vals)...]
     size_dict = OMEinsum.get_size_dict(all_ixs, all_tensors)

test/cuda.jl

@@ -42,22 +42,22 @@ 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; marginalized = Int[], optimizer)
+    tn = MMAPModel(instance; queryvars = collect(1:instance.nvars), 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; marginalized = collect(1:(instance.nvars)), optimizer)
+    tn2 = MMAPModel(instance; queryvars = Int[], 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; marginalized = [2, 4, 6], optimizer)
+    tn3 = MMAPModel(instance; queryvars = setdiff(1:instance.nvars, [2, 4, 6]), optimizer)
     logp, config = most_probable_config(tn3; usecuda = true)
     @test log_probability(tn3, config) ≈ logp
 end

test/mmap.jl

@@ -15,17 +15,17 @@ end
     tn_ref = TensorNetworkModel(instance; optimizer)
 
     # Does not marginalize any var
-    mmap = MMAPModel(instance; marginalized = Int[], optimizer)
+    mmap = MMAPModel(instance; queryvars = collect(1:instance.nvars), optimizer)
     @debug(mmap)
     @test maximum_logp(tn_ref) ≈ maximum_logp(mmap)
 
     # Marginalize all vars
-    mmap2 = MMAPModel(instance; marginalized = collect(1:(instance.nvars)), optimizer)
+    mmap2 = MMAPModel(instance; queryvars = Int[], optimizer)
     @debug(mmap2)
     @test Array(probability(tn_ref))[] ≈ exp(maximum_logp(mmap2)[])
 
     # Does not optimize over open vertices
-    mmap3 = MMAPModel(instance; marginalized = [2, 4, 6], optimizer)
+    mmap3 = MMAPModel(instance; queryvars = setdiff(1:instance.nvars, [2, 4, 6]), optimizer)
     @debug(mmap3)
     logp, config = most_probable_config(mmap3)
     @test log_probability(mmap3, config) ≈ logp
@@ -42,17 +42,8 @@ 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; marginalized = setdiff(1:(instance.nvars), instance.queryvars), optimizer)
+      model = MMAPModel(instance; queryvars = instance.queryvars, optimizer)
       _, solution = most_probable_config(model)
       @test solution == instance.reference_solution
     end
-end
-
-using Artifacts
-include("utils.jl")
-model_filepath, evidence_filepath, query_filepath, solution_filepath = get_instance_filepaths("Segmentation_11", "MMAP")
-instance = read_instance(model_filepath; evidence_filepath, query_filepath, solution_filepath)
-ref_sol = read_solution_file(solution_filepath)[2:end]
-
-optimizer = TreeSA(ntrials=1, niters=2, βs=1:0.1:40)
-mmap = MMAPModel(instance; marginalized=setdiff(1:instance.nvars, instance.queryvars), optimizer)
+end