update

Author: GiggleLiu

example/asia/README.md

@@ -0,0 +1,25 @@
+# The "Asia" Bayesian network
+
+The variables and factors in `.uai` files are labelled as below.
+
+## 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,5 @@
+using TensorInference
+
+problem = uai_problem_from_file(joinpath(@__DIR__, "data/asia.uai"))
+tnet = TensorNetworkModel(problem)
+marginals(problem)

example/asia/data/asia.uai

@@ -0,0 +1,36 @@
+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

example/asia/data/asia.uai.evid

@@ -0,0 +1 @@
+2 0 0 6 0

src/Core.jl

@@ -58,7 +58,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)

src/TensorInference.jl

@@ -7,10 +7,10 @@ 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
 
 # marginals
 export TensorNetworkModel, get_vars, get_cards, log_probability, probability, marginals

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))")
@@ -99,7 +99,7 @@ function MMAPModel(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; margi
     return MMAPModel(setdiff(vars, marginalizedvertices), 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)...,

src/utils.jl

@@ -71,8 +71,7 @@ https://personal.utdallas.edu/~vibhav.gogate/uai16-evaluation/uaiformat.html
 function read_uai_evid_file(uai_evid_filepath::AbstractString)
     if isempty(uai_evid_filepath)
         # No evidence
-        obsvars = Int64[]
-        obsvals = Int64[]
+        return Int64[], Int64[]
     else
         # Read the last line of the uai evid file
         line = open(uai_evid_filepath) do file
@@ -90,10 +89,6 @@ function read_uai_evid_file(uai_evid_filepath::AbstractString)
         @assert nobsvars == length(obsvars)
     end
 
-    # # DEBUG:
-    # print("  "); @show obsvars
-    # print("  "); @show obsvals
-
     return obsvars, obsvals
 end
 
@@ -179,14 +174,21 @@ $(TYPEDSIGNATURES)
 
 Read a UAI problem from an artifact.
 """
-function read_uai_problem(problem::AbstractString)::UAIInstance
+function read_uai_problem(problem::AbstractString; eltype=Float64)::UAIInstance
     uai_filepath = joinpath(artifact"MAR_prob", problem * ".uai")
     uai_evid_filepath = joinpath(artifact"MAR_prob", problem * ".uai.evid")
     uai_mar_filepath = joinpath(artifact"MAR_sol", problem * ".uai.MAR")
+    return uai_problem_from_file(uai_filepath; uai_evid_filepath, uai_mar_filepath, eltype)
+end
 
-    nvars, cards, ncliques, factors = read_uai_file(uai_filepath; factor_eltype = Float64)
-    obsvars, obsvals = read_uai_evid_file(uai_evid_filepath)
-    reference_marginals = read_uai_mar_file(uai_mar_filepath)
+"""
+$(TYPEDSIGNATURES)
 
+Read a UAI problem from a file.
+"""
+function uai_problem_from_file(uai_filepath::String; uai_evid_filepath="", uai_mar_filepath="", eltype=Float64)::UAIInstance
+    nvars, cards, ncliques, factors = read_uai_file(uai_filepath; factor_eltype = eltype)
+    obsvars, obsvals = read_uai_evid_file(uai_evid_filepath)
+    reference_marginals = isempty(uai_mar_filepath) ? Vector{eltype}[] : read_uai_mar_file(uai_mar_filepath)
     return UAIInstance(nvars, ncliques, cards, factors, obsvars, obsvals, reference_marginals)
 end

test/inference.jl

@@ -75,7 +75,7 @@ end
                             error("space complexity too large! got $(sc)")
                         end
                         # @info(tn)
-                        # @info timespace_complexity(tn)
+                        @info contraction_complexity(tn)
                         marginals2 = marginals(tn)
                         # for dangling vertices, the output size is 1.
                         npass = 0