refactor UAI file reading

Author: GiggleLiu

docs/src/api/public.md

@@ -41,6 +41,7 @@ TreeSA
 MMAPModel
 RescaledArray
 TensorNetworkModel
+ArtifactProblemSpec
 UAIInstance
 ```
 
@@ -55,13 +56,15 @@ marginals
 maximum_logp
 most_probable_config
 probability
-read_evidence_file
+dataset_from_artifact
+problem_from_artifact
 read_instance
-read_instance_from_artifact
-read_model_file
+read_evidence
+read_solution
+read_queryvars
+read_instance_from_file
+read_evidence_file
 read_solution_file
 read_td_file
 sample
-set_evidence!
-set_query!
 ```

examples/asia/main.jl

@@ -76,13 +76,9 @@ get_vars(tn)
 # ---
 
 # Set an evidence: Assume that the "X-ray" result (variable 7) is positive.
-set_evidence!(instance, 7 => 0)
-
-# ---
-
 # Since setting an evidence may affect the contraction order of the tensor
 # network, recompute it.
-tn = TensorNetworkModel(instance)
+tn = TensorNetworkModel(instance, evidence=Dict(7=>0))
 
 # ---
 
@@ -107,8 +103,7 @@ logp, cfg = most_probable_config(tn)
 # Compute the most probable values of certain variables (e.g., 4 and 7) while
 # marginalizing over others. This is known as Maximum a Posteriori (MAP)
 # estimation.
-set_query!(instance, [4, 7])
-mmap = MMAPModel(instance)
+mmap = MMAPModel(instance, queryvars=[4,7])
 
 # ---
 

src/Core.jl

@@ -22,62 +22,23 @@ $(TYPEDEF)
 * `nclique` is the number of cliques,
 * `cards` is a vector of cardinalities for variables,
 * `factors` is a vector of factors,
-
-* `obsvars` is a vector of observed variables,
-* `obsvals` is a vector of observed values,
-* `queryvars` is a vector of query variables,
-* `reference_solution` is a vector with the reference solution.
 """
 struct UAIInstance{ET, FT <: Factor{ET}}
     nvars::Int
     nclique::Int
     cards::Vector{Int}
     factors::Vector{FT}
-
-    obsvars::Vector{Int}
-    obsvals::Vector{Int}
-    queryvars::Vector{Int}
-    reference_solution
 end
 
 Base.show(io::IO, ::MIME"text/plain", uai::UAIInstance) = Base.show(io, uai)
 function Base.show(io::IO, uai::UAIInstance)
     println(io, "UAIInstance(nvars = $(uai.nvars), nclique = $(uai.nclique))")
     println(io, " variables :")
-    for (var, card) in zip(1:uai.nvars, uai.cards)
-        println(io, string_var("  $var of size $card", var ∈ uai.queryvars, Dict(zip(uai.obsvars, uai.obsvals))))
-    end
     println(io, " factors : ")
-    for f in uai.factors
-        println(io, "  $(summary(f))")
+    for (k, f) in enumerate(uai.factors)
+        print(io, "  $(summary(f))")
+        k == length(uai.factors) || println(io)
     end
-    print(io, " reference_solution : $(uai.reference_solution)")
-end
-
-"""
-$TYPEDSIGNATURES
-
-Set the evidence of an UAI instance.
-"""
-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
-
-"""
-$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
 
 """
@@ -89,32 +50,32 @@ Probabilistic modeling with a tensor network.
 * `vars` is the degree of freedoms in the tensor network.
 * `code` is the tensor network contraction pattern.
 * `tensors` is the tensors fed into the tensor network.
-* `fixedvertices` is a dictionary to specifiy degree of freedoms fixed to certain values.
+* `evidence` is a dictionary to specifiy degree of freedoms fixed to certain values.
 """
 struct TensorNetworkModel{LT, ET, MT <: AbstractArray}
     vars::Vector{LT}
     code::ET
     tensors::Vector{MT}
-    fixedvertices::Dict{LT, Int}
+    evidence::Dict{LT, Int}
 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], ", ")
+    variables = join([string_var(var, open, tn.evidence) for var in tn.vars], ", ")
     tc, sc, rw = contraction_complexity(tn)
     println(io, "$(typeof(tn))")
     println(io, "variables: $variables")
     print_tcscrw(io, tc, sc, rw)
 end
 Base.show(io::IO, ::MIME"text/plain", tn::TensorNetworkModel) = Base.show(io, tn)
 
-function string_var(var, open, fixedvertices)
-    if var ∈ open && haskey(fixedvertices, var)
-        "$var (open, fixed to $(fixedvertices[var]))"
+function string_var(var, open, evidence)
+    if var ∈ open && haskey(evidence, var)
+        "$var (open, fixed to $(evidence[var]))"
     elseif var ∈ open
         "$var (open)"
-    elseif haskey(fixedvertices, var)
-        "$var (evidence → $(fixedvertices[var]))"
+    elseif haskey(evidence, var)
+        "$var (evidence → $(evidence[var]))"
     else
         "$var"
     end
@@ -129,19 +90,17 @@ $(TYPEDSIGNATURES)
 """
 function TensorNetworkModel(
     instance::UAIInstance;
-    openvertices = (),
+    openvars = (),
+    evidence = Dict{Int,Int}(),
     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,
         instance.factors;
-        openvertices,
-        fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)),
+        openvars,
+        evidence,
         optimizer,
         simplifier
     )
@@ -154,18 +113,18 @@ function TensorNetworkModel(
     vars::AbstractVector{LT},
     cards::AbstractVector{Int},
     factors::Vector{<:Factor{T}};
-    openvertices = (),
-    fixedvertices = Dict{LT, Int}(),
+    openvars = (),
+    evidence = Dict{LT, Int}(),
     optimizer = GreedyMethod(),
     simplifier = nothing
 )::TensorNetworkModel where {T, LT}
     # The 1st argument of `EinCode` is a vector of vector of labels for specifying the input tensors, 
     # The 2nd argument of `EinCode` is a vector of labels for specifying the output tensor,
     # e.g.
     # `EinCode([[1, 2], [2, 3]], [1, 3])` is the EinCode for matrix multiplication.
-    rawcode = EinCode([[[var] for var in vars]..., [[factor.vars...] for factor in factors]...], collect(LT, openvertices))  # labels for vertex tensors (unity tensors) and edge tensors
+    rawcode = EinCode([[[var] for var in vars]..., [[factor.vars...] for factor in factors]...], collect(LT, openvars))  # labels for vertex tensors (unity tensors) and edge tensors
     tensors = Array{T}[[ones(T, cards[i]) for i in 1:length(vars)]..., [t.vals for t in factors]...]
-    return TensorNetworkModel(collect(LT, vars), rawcode, tensors; fixedvertices, optimizer, simplifier)
+    return TensorNetworkModel(collect(LT, vars), rawcode, tensors; evidence, optimizer, simplifier)
 end
 
 """
@@ -175,7 +134,7 @@ function TensorNetworkModel(
     vars::AbstractVector{LT},
     rawcode::EinCode,
     tensors::Vector{<:AbstractArray};
-    fixedvertices = Dict{LT, Int}(),
+    evidence = Dict{LT, Int}(),
     optimizer = GreedyMethod(),
     simplifier = nothing
 )::TensorNetworkModel where {LT}
@@ -185,7 +144,7 @@ function TensorNetworkModel(
     # The 3rd and 4th arguments are the optimizer and simplifier that configures which algorithm to use and simplify.
     size_dict = OMEinsum.get_size_dict(getixsv(rawcode), tensors)
     code = optimize_code(rawcode, size_dict, optimizer, simplifier)
-    TensorNetworkModel(collect(LT, vars), code, tensors, fixedvertices)
+    TensorNetworkModel(collect(LT, vars), code, tensors, evidence)
 end
 
 """
@@ -202,10 +161,10 @@ Get the cardinalities of variables in this tensor network.
 """
 function get_cards(tn::TensorNetworkModel; fixedisone = false)::Vector
     vars = get_vars(tn)
-    [fixedisone && haskey(tn.fixedvertices, vars[k]) ? 1 : length(tn.tensors[k]) for k in 1:length(vars)]
+    [fixedisone && haskey(tn.evidence, vars[k]) ? 1 : length(tn.tensors[k]) for k in 1:length(vars)]
 end
 
-chfixedvertices(tn::TensorNetworkModel, fixedvertices) = TensorNetworkModel(tn.vars, tn.code, tn.tensors, fixedvertices)
+chevidence(tn::TensorNetworkModel, evidence) = TensorNetworkModel(tn.vars, tn.code, tn.tensors, evidence)
 
 """
 $(TYPEDSIGNATURES)
@@ -221,7 +180,7 @@ end
 $(TYPEDSIGNATURES)
 
 Contract the tensor network and return a probability array with its rank specified in the contraction code `tn.code`.
-The returned array may not be l1-normalized even if the total probability is l1-normalized, because the evidence `tn.fixedvertices` may not be empty.
+The returned array may not be l1-normalized even if the total probability is l1-normalized, because the evidence `tn.evidence` may not be empty.
 """
 function probability(tn::TensorNetworkModel; usecuda = false, rescale = true)::AbstractArray
     return tn.code(adapt_tensors(tn; usecuda, rescale)...)

src/TensorInference.jl

@@ -18,9 +18,9 @@ export RescaledArray
 export contraction_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABipartite, MergeGreedy, MergeVectors
 
 # read and load uai files
-export read_model_file, read_td_file, read_evidence_file, read_solution_file
-export read_instance, read_instance_from_artifact, UAIInstance
-export set_evidence!, set_query!
+export read_instance_from_file, read_td_file, read_evidence_file, read_solution_file
+export problem_from_artifact, ArtifactProblemSpec
+export read_instance, UAIInstance, read_evidence, read_solution, read_queryvars, dataset_from_artifact
 
 # marginals
 export TensorNetworkModel, get_vars, get_cards, log_probability, probability, marginals
@@ -47,13 +47,13 @@ function __init__()
     @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
 end
 
-import PrecompileTools
-PrecompileTools.@setup_workload begin
-    # Putting some things in `@setup_workload` instead of `@compile_workload` can reduce the size of the
-    # precompile file and potentially make loading faster.
-    #PrecompileTools.@compile_workload begin
-    #    include("../example/asia/main.jl")
-    #end
-end
+# import PrecompileTools
+# PrecompileTools.@setup_workload begin
+#     # Putting some things in `@setup_workload` instead of `@compile_workload` can reduce the size of the
+#     # precompile file and potentially make loading faster.
+#     PrecompileTools.@compile_workload begin
+#         include("../example/asia/main.jl")
+#     end
+# end
 
 end # module

src/map.jl

@@ -49,7 +49,7 @@ function most_probable_config(tn::TensorNetworkModel; usecuda = false)::Tuple{Re
     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 content(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.evidence, vars[k]) ? tn.evidence[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
 end
 
 """

src/mar.jl

@@ -1,11 +1,11 @@
 # generate tensors based on which vertices are fixed.
-adapt_tensors(gp::TensorNetworkModel; usecuda, rescale) = adapt_tensors(gp.code, gp.tensors, gp.fixedvertices; usecuda, rescale)
-function adapt_tensors(code, tensors, fixedvertices; usecuda, rescale)
+adapt_tensors(gp::TensorNetworkModel; usecuda, rescale) = adapt_tensors(gp.code, gp.tensors, gp.evidence; usecuda, rescale)
+function adapt_tensors(code, tensors, evidence; usecuda, rescale)
     ixs = getixsv(code)
     # `ix` is the vector of labels (or a degree of freedoms) for a tensor,
     # if a label in `ix` is fixed to a value, do the slicing to the tensor it associates to.
     map(tensors, ixs) do t, ix
-        dims = map(ixi -> ixi ∉ keys(fixedvertices) ? Colon() : ((fixedvertices[ixi] + 1):(fixedvertices[ixi] + 1)), ix)
+        dims = map(ixi -> ixi ∉ keys(evidence) ? Colon() : ((evidence[ixi] + 1):(evidence[ixi] + 1)), ix)
         t2 = t[dims...]
         t3 = usecuda ? CuArray(t2) : t2
         rescale ? rescale_array(t3) : t3