new asia example

GiggleLiu · GiggleLiu · commit 3145d8fab9fe · 2023-06-30T04:39:21.000+08:00
diff --git a/benchmark/bench_mmap.jl b/benchmark/bench_mmap.jl
@@ -10,15 +10,15 @@ problem = read_uai_problem("Promedus_14")
 optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40)
 
 # Does not marginalize any var
-mmap1 = MMAPModel(problem; marginalizedvertices = Int[], optimizer)
+mmap1 = MMAPModel(problem; marginalized = Int[], optimizer)
 SUITE["mmap-1"] = @benchmarkable maximum_logp(mmap1)
 
 # Marginalizes all vars
-mmap2 = MMAPModel(problem; marginalizedvertices = collect(1:(problem.nvars)), optimizer)
+mmap2 = MMAPModel(problem; marginalized = collect(1:(problem.nvars)), optimizer)
 SUITE["mmap-2"] = @benchmarkable maximum_logp(mmap2)
 
 # Does not optimize over open vertices
-mmap3 = MMAPModel(problem; marginalizedvertices = [2, 4, 6], optimizer)
+mmap3 = MMAPModel(problem; marginalized = [2, 4, 6], optimizer)
 SUITE["mmap-3"] = @benchmarkable most_probable_config(mmap3)
 
 end  # module
diff --git a/example/asia/README.md b/example/asia/README.md
@@ -1,6 +1,12 @@
 # The "Asia" Bayesian network
 
-The variables and factors in `.uai` files are labelled as below.
+Please check the Julia code [asia.jl](asia.jl).
+
+The variables and factors for the asia model is described in the [asia.uai](asia.uai) file.
+The UAI file format is detailed in:
+https://personal.utdallas.edu/~vibhav.gogate/uai16-evaluation/uaiformat.html
+
+The meanings of variables and factors as listed bellow.
 
 ## Variables
 index 0 is mapped to yes, 1 is mapped to no.
diff --git a/example/asia/asia.jl b/example/asia/asia.jl
@@ -1,5 +1,33 @@
 using TensorInference
 
-problem = uai_problem_from_file(joinpath(@__DIR__, "data/asia.uai"))
-tnet = TensorNetworkModel(problem)
-marginals(problem)
+# Load the model that detailed in the README and `asia.uai`.
+instance = uai_problem_from_file(joinpath(@__DIR__, "asia.uai"))
+tnet = TensorNetworkModel(instance)
+
+# Get the probabilities (PR)
+probability(tnet)
+
+# Get the marginal probabilities (MAR)
+marginals(tnet) .|> first
+
+# Set the evidence variables "X-ray" (7) to be positive.
+set_evidence!(instance, 7=>0)
+
+# Since the evidence variable may change the contraction order, we re-compute the tensor network.
+tnet = TensorNetworkModel(instance)
+
+# Get the maximum log-probabilities (MAP)
+maximum_logp(tnet)
+
+# Get not only the maximum log-probability, but also the most probable conifguration
+# In the most probable configuration, the most probable one is the patient smoke (3) and has lung cancer (4)
+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])
+# We get the most probable configurations on [4, 7]
+most_probable_config(mmap)
+# The total probability of having lung cancer is roughly half.
+log_probability(mmap, [1, 0])
+log_probability(mmap, [0, 0])
diff --git a/example/asia/asia.uai b/example/asia/asia.uai
@@ -0,0 +1,44 @@
+MARKOV
+8
+ 2 2 2 2 2 2 2 2
+8
+ 1 0
+ 2 1 0
+ 1 2
+ 2 3 2
+ 2 4 2
+ 3 5 3 1
+ 2 6 5
+ 3 7 5 4
+
+2
+ 0.01
+ 0.99
+
+4
+ 0.05 0.01
+ 0.95 0.99
+
+2
+ 0.5
+ 0.5
+
+4
+ 0.1 0.01
+ 0.9 0.99
+
+4
+ 0.6 0.3
+ 0.4 0.7 
+
+8
+ 1 1 1 0
+ 0 0 0 1
+
+4
+ 0.98 0.05
+ 0.02 0.95
+
+8
+ 0.9 0.7 0.8 0.1
+ 0.1 0.3 0.2 0.9
diff --git a/example/asia/data/asia.uai b/example/asia/data/asia.uai
diff --git a/example/asia/data/asia.uai.evid b/example/asia/data/asia.uai.evid
diff --git a/src/Core.jl b/src/Core.jl
@@ -37,6 +37,30 @@ struct UAIInstance{ET, FT <: Factor{ET}}
     reference_marginals::Vector{Vector{ET}}
 end
 
+"""
+$TYPEDSIGNATURES
+
+Set the evidence of an UAI instance.
+
+### Examples
+```jldoctest; setup=:(using TensorInference)
+julia> problem = read_uai_problem("Promedus_14"); problem.obsvars, problem.obsvals
+([42, 48, 27, 30, 29, 15, 124, 5, 148], [1, 1, 1, 1, 1, 1, 1, 1, 1])
+
+julia> set_evidence!(problem, 2=>0, 4=>1); problem.obsvars, problem.obsvals
+([2, 4], [0, 1])
+```
+"""
+function set_evidence!(uai::UAIInstance, pairs::Pair{Int}...)
+    empty!(uai.obsvars)
+    empty!(uai.obsvals)
+    for (var, val) in pairs
+        push!(uai.obsvars, var)
+        push!(uai.obsvals, val)
+    end
+    return uai
+end
+
 """
 $(TYPEDEF)
 
@@ -95,7 +119,7 @@ function TensorNetworkModel(
         instance.cards,
         instance.factors;
         openvertices,
-        fixedvertices = Dict(zip(instance.obsvars, instance.obsvals .- 1)),
+        fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)),
         optimizer,
         simplifier
     )
diff --git a/src/RescaledArray.jl b/src/RescaledArray.jl
@@ -15,6 +15,7 @@ Base.show(io::IO, c::RescaledArray) = print(io, "exp($(c.log_factor)) * $(c.norm
 Base.show(io::IO, ::MIME"text/plain", c::RescaledArray) = Base.show(io, c)
 Base.Array(c::RescaledArray) = rmul!(Array(c.normalized_value), exp(c.log_factor))
 Base.copy(c::RescaledArray) = RescaledArray(c.log_factor, copy(c.normalized_value))
+Base.getindex(r::RescaledArray, indices...) = map(x->x * exp(r.log_factor), getindex(r.normalized_value, indices...))
 
 """
 $(TYPEDSIGNATURES)
diff --git a/src/TensorInference.jl b/src/TensorInference.jl
@@ -11,6 +11,7 @@ export contraction_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABiparti
 
 # read and load uai files
 export read_uai_file, read_td_file, read_uai_evid_file, read_uai_mar_file, read_uai_problem, uai_problem_from_file
+export set_evidence!
 
 # marginals
 export TensorNetworkModel, get_vars, get_cards, log_probability, probability, marginals
diff --git a/src/maxprob.jl b/src/maxprob.jl
@@ -19,7 +19,7 @@ The backward rule for tropical einsum.
 function backward_tropical(ixs, @nospecialize(xs::Tuple), iy, @nospecialize(y), @nospecialize(ymask), size_dict)
     y .= masked_inv.(y, ymask)
     masks = []
-    for i in 1:length(ixs)
+    for i in eachindex(ixs)
         nixs = OMEinsum._insertat(ixs, i, iy)
         nxs  = OMEinsum._insertat(xs, i, y)
         niy  = ixs[i]
@@ -44,21 +44,21 @@ $(TYPEDSIGNATURES)
 
 Returns the largest log-probability and the most probable configuration.
 """
-function most_probable_config(tn::TensorNetworkModel; usecuda = false)::Tuple{Tropical, Vector}
+function most_probable_config(tn::TensorNetworkModel; usecuda = false)::Tuple{Real, Vector}
     vars = get_vars(tn)
     tensors = map(t -> Tropical.(log.(t)), adapt_tensors(tn; usecuda, rescale = false))
     logp, grads = cost_and_gradient(tn.code, tensors)
     # use Array to convert CuArray to CPU arrays
-    return Array(logp)[], map(k -> haskey(tn.fixedvertices, vars[k]) ? tn.fixedvertices[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
+    return content(Array(logp)[]), map(k -> haskey(tn.fixedvertices, vars[k]) ? tn.fixedvertices[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
 end
 
 """
 $(TYPEDSIGNATURES)
 
 Returns an output array containing largest log-probabilities.
 """
-function maximum_logp(tn::TensorNetworkModel; usecuda = false)::AbstractArray{<:Tropical}
+function maximum_logp(tn::TensorNetworkModel; usecuda = false)::AbstractArray{<:Real}
     # generate tropical tensors with its elements being log(p).
     tensors = map(t -> Tropical.(log.(t)), adapt_tensors(tn; usecuda, rescale = false))
-    return tn.code(tensors...)
+    return map(content, tn.code(tensors...))
 end
diff --git a/src/mmap.jl b/src/mmap.jl
@@ -58,16 +58,16 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function MMAPModel(instance::UAIInstance; marginalizedvertices, openvertices = (), optimizer = GreedyMethod(), simplifier = nothing)::MMAPModel
+function MMAPModel(instance::UAIInstance; marginalized, openvertices = (), optimizer = GreedyMethod(), simplifier = nothing)::MMAPModel
     return MMAPModel(
-        1:(instance.nvars), instance.factors; marginalizedvertices, fixedvertices = Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices
+        1:(instance.nvars), instance.factors; marginalized, fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)), optimizer, simplifier, openvertices
     )
 end
 
 """
 $(TYPEDSIGNATURES)
 """
-function MMAPModel(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; marginalizedvertices, openvertices = (),
+function MMAPModel(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; marginalized, openvertices = (),
     fixedvertices = Dict{LT, Int}(),
     optimizer = GreedyMethod(), simplifier = nothing,
     marginalize_optimizer = GreedyMethod(), marginalize_simplifier = nothing
@@ -81,7 +81,7 @@ function MMAPModel(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; margi
     size_dict = OMEinsum.get_size_dict(all_ixs, all_tensors)
 
     # detect clusters for marginalize variables
-    subsets = connected_clusters(all_ixs, marginalizedvertices)
+    subsets = connected_clusters(all_ixs, marginalized)
     clusters = Cluster{LT}[]
     ixs = Vector{LT}[]
     for (contracted, cluster) in subsets
@@ -96,7 +96,7 @@ function MMAPModel(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; margi
     rem_indices = setdiff(1:length(all_ixs), vcat([c.second for c in subsets]...))
     remaining_tensors = all_tensors[rem_indices]
     code = optimize_code(EinCode([all_ixs[rem_indices]..., ixs...], iy), size_dict, optimizer, simplifier)
-    return MMAPModel(setdiff(vars, marginalizedvertices), code, remaining_tensors, clusters, fixedvertices)
+    return MMAPModel(setdiff(vars, marginalized), code, remaining_tensors, clusters, fixedvertices)
 end
 
 function OMEinsum.contraction_complexity(mmap::MMAPModel{LT}) where {LT}
@@ -173,35 +173,29 @@ function visit_var!(var, vars::AbstractVector{LT}, ixs, visited_ixs, visited_var
     end
 end
 
-"""
-$(TYPEDSIGNATURES)
-"""
-function most_probable_config(mmap::MMAPModel; usecuda = false)::Tuple{Tropical, Vector}
+function most_probable_config(mmap::MMAPModel; usecuda = false)::Tuple{Real, Vector}
     vars = get_vars(mmap)
     tensors = map(t -> OMEinsum.asarray(Tropical.(log.(t)), t), adapt_tensors(mmap; usecuda, rescale = false))
     logp, grads = cost_and_gradient(mmap.code, tensors)
     # use Array to convert CuArray to CPU arrays
-    return Array(logp)[], map(k -> haskey(mmap.fixedvertices, vars[k]) ? mmap.fixedvertices[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
+    return content(Array(logp)[]), map(k -> haskey(mmap.fixedvertices, vars[k]) ? mmap.fixedvertices[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
 end
 
-"""
-$(TYPEDSIGNATURES)
-"""
-function maximum_logp(mmap::MMAPModel; usecuda = false)::AbstractArray{<:Tropical}
+function maximum_logp(mmap::MMAPModel; usecuda = false)::AbstractArray{<:Real}
     tensors = map(t -> OMEinsum.asarray(Tropical.(log.(t)), t), adapt_tensors(mmap; usecuda, rescale = false))
-    return mmap.code(tensors...)
+    return map(content, mmap.code(tensors...))
 end
 
-"""
-$(TYPEDSIGNATURES)
-"""
 function log_probability(mmap::MMAPModel, config::Union{Dict, AbstractVector}; rescale = true, usecuda = false)::Real
     @assert length(get_vars(mmap)) == length(config)
     fixedvertices = config isa AbstractVector ? Dict(zip(get_vars(mmap), config)) : config
     assign = merge(mmap.fixedvertices, fixedvertices)
     # two contributions to the probability, not-clustered tensors and clusters.
     m1 = sum(x -> log(x[2][(getindex.(Ref(assign), x[1]) .+ 1)...]), zip(getixsv(mmap.code), mmap.tensors))
-    m2 = sum(cluster -> probability(cluster; fixedvertices, usecuda, rescale).log_factor, mmap.clusters)
+    m2 = sum(mmap.clusters) do cluster
+        p = probability(cluster; fixedvertices, usecuda, rescale)
+        rescale ? p.log_factor : log(p[])
+    end
     return m1 + m2
 end
 
diff --git a/src/utils.jl b/src/utils.jl
@@ -84,7 +84,7 @@ function read_uai_evid_file(uai_evid_filepath::AbstractString)
 
         # Convert to 1-based indexing
         obsvars = observations[1, :] .+ 1
-        obsvals = observations[2, :] .+ 1
+        obsvals = observations[2, :]
 
         @assert nobsvars == length(obsvars)
     end
diff --git a/test/cuda.jl b/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; marginalizedvertices = Int[], optimizer)
+    tn = MMAPModel(instance; marginalized = Int[], 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; marginalizedvertices = collect(1:(instance.nvars)), optimizer)
+    tn2 = MMAPModel(instance; marginalized = collect(1:(instance.nvars)), 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)[].n)
 
     # does not optimize over open vertices
-    tn3 = MMAPModel(instance; marginalizedvertices = [2, 4, 6], optimizer)
+    tn3 = MMAPModel(instance; marginalized = [2, 4, 6], optimizer)
     logp, config = most_probable_config(tn3; usecuda = true)
     @test log_probability(tn3, config) ≈ logp.n
 end
diff --git a/test/maxprob.jl b/test/maxprob.jl
@@ -11,6 +11,6 @@ using TensorInference
     @info contraction_complexity(tn)
     most_probable_config(tn)
     @time logp, config = most_probable_config(tn)
-    @test log_probability(tn, config) ≈ logp.n
+    @test log_probability(tn, config) ≈ logp
     @test maximum_logp(tn)[] ≈ logp
 end
diff --git a/test/mmap.jl b/test/mmap.jl
@@ -14,18 +14,18 @@ end
     optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40)
     tn_ref = TensorNetworkModel(instance; optimizer)
     # does not marginalize any var
-    mmap = MMAPModel(instance; marginalizedvertices = Int[], optimizer)
+    mmap = MMAPModel(instance; marginalized = Int[], optimizer)
     @info(mmap)
     @test maximum_logp(tn_ref) ≈ maximum_logp(mmap)
 
     # marginalize all vars
-    mmap2 = MMAPModel(instance; marginalizedvertices = collect(1:(instance.nvars)), optimizer)
+    mmap2 = MMAPModel(instance; marginalized = collect(1:(instance.nvars)), optimizer)
     @info(mmap2)
-    @test Array(probability(tn_ref))[] ≈ exp(maximum_logp(mmap2)[].n)
+    @test Array(probability(tn_ref))[] ≈ exp(maximum_logp(mmap2)[])
 
     # does not optimize over open vertices
-    mmap3 = MMAPModel(instance; marginalizedvertices = [2, 4, 6], optimizer)
+    mmap3 = MMAPModel(instance; marginalized = [2, 4, 6], optimizer)
     @info(mmap3)
     logp, config = most_probable_config(mmap3)
-    @test log_probability(mmap3, config) ≈ logp.n
+    @test log_probability(mmap3, config) ≈ logp
 end
