Merge pull request #22 from TensorBFS/jg/asia-example

Add Asia example and polish the interfaces

Author: mroavi

README.md

@@ -28,3 +28,6 @@ pkg> add TensorInference
 ```
 
 To update, just type `up` in the package mode.
+
+## Example
+Please check the [example](example) folder.

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

example/asia/README.md

@@ -0,0 +1,31 @@
+# The "Asia" Bayesian network
+
+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.
+
+1. visit to Asia (a)
+2. tuberculosis (t)
+3. smoking (s)
+4. lung cancer (l)
+5. bronchitis (b)
+6. either tub. or lung cancer (e)
+7. positive X-ray (x)
+8. dyspnoea (d)
+
+## Factors
+1. p(a)
+2. p(t|a)
+3. p(s)
+4. p(l|s)
+5. p(b|s)
+6. p(e|l,t)
+7. p(x|e)
+8. p(d|e,b)

example/asia/asia.jl

@@ -0,0 +1,33 @@
+using TensorInference
+
+# 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])

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

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)
 
@@ -58,7 +82,7 @@ end
 function Base.show(io::IO, tn::TensorNetworkModel)
     open = getiyv(tn.code)
     variables = join([string_var(var, open, tn.fixedvertices) for var in tn.vars], ", ")
-    tc, sc, rw = timespacereadwrite_complexity(tn)
+    tc, sc, rw = contraction_complexity(tn)
     println(io, "$(typeof(tn))")
     println(io, "variables: $variables")
     print_tcscrw(io, tc, sc, rw)
@@ -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
     )

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)

src/TensorInference.jl

@@ -7,10 +7,11 @@ using TropicalGEMM
 
 # reexport OMEinsum functions
 export RescaledArray
-export timespace_complexity, timespacereadwrite_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABipartite, MergeGreedy, MergeVectors
+export contraction_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABipartite, MergeGreedy, MergeVectors
 
 # read and load uai files
-export read_uai_file, read_td_file, read_uai_evid_file, read_uai_mar_file, read_uai_problem
+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

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

src/mmap.jl

@@ -34,7 +34,7 @@ end
 function Base.show(io::IO, mmap::MMAPModel)
     open = getiyv(mmap.code)
     variables = join([string_var(var, open, mmap.fixedvertices) for var in mmap.vars], ", ")
-    tc, sc, rw = timespacereadwrite_complexity(mmap)
+    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))")
@@ -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,10 +96,10 @@ 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.timespacereadwrite_complexity(mmap::MMAPModel{LT}) where {LT}
+function OMEinsum.contraction_complexity(mmap::MMAPModel{LT}) where {LT}
     # extract size
     size_dict = Dict(zip(get_vars(mmap), get_cards(mmap; fixedisone = true)))
     sc = -Inf
@@ -111,18 +111,17 @@ function OMEinsum.timespacereadwrite_complexity(mmap::MMAPModel{LT}) where {LT}
             # the head sector are for unity tensors.
             size_dict[cluster.eliminated_vars[k]] = length(cluster.tensors[k])
         end
-        tc, sci, rw = timespacereadwrite_complexity(cluster.code, size_dict)
+        tc, sci, rw = contraction_complexity(cluster.code, size_dict)
         push!(tcs, tc)
         push!(rws, rw)
         sc = max(sc, sci)
     end
 
-    tc, sci, rw = timespacereadwrite_complexity(mmap.code, size_dict)
+    tc, sci, rw = contraction_complexity(mmap.code, size_dict)
     push!(tcs, tc)
     push!(rws, tc)
     OMEinsum.OMEinsumContractionOrders.log2sumexp2(tcs), max(sc, sci), OMEinsum.OMEinsumContractionOrders.log2sumexp2(rws)
 end
-OMEinsum.timespace_complexity(mmap::MMAPModel) = timespacereadwrite_complexity(mmap)[1:2]
 
 function adapt_tensors(mmap::MMAPModel; usecuda, rescale)
     return [adapt_tensors(mmap.code, mmap.tensors, mmap.fixedvertices; usecuda, rescale)...,
@@ -174,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