Merge pull request #48 from TensorBFS/jg/revert-tropicalgemm

remove tropicalgemm as a dependency

Author: mroavi

Project.toml

@@ -11,7 +11,6 @@ OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
-TropicalGEMM = "a4ad3063-64a7-4bad-8738-34ed09bc0236"
 TropicalNumbers = "b3a74e9c-7526-4576-a4eb-79c0d4c32334"
 
 [compat]
@@ -21,6 +20,5 @@ OMEinsum = "0.7"
 PrecompileTools = "1"
 Requires = "1"
 StatsBase = "0.34"
-TropicalGEMM = "0.1"
 TropicalNumbers = "0.5.4"
 julia = "1.3"

docs/src/api/public.md

@@ -62,4 +62,5 @@ read_solution_file
 read_td_file
 sample
 set_evidence!
+set_query!
 ```

examples/asia/main.jl

@@ -107,7 +107,8 @@ 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.
-mmap = MMAPModel(instance; queryvars = [4, 7])
+set_query!(instance, [4, 7])
+mmap = MMAPModel(instance)
 
 # ---
 

src/Core.jl

@@ -37,7 +37,7 @@ struct UAIInstance{ET, FT <: Factor{ET}}
     obsvars::Vector{Int}
     obsvals::Vector{Int}
     queryvars::Vector{Int}
-    reference_solution::Union{Vector{Vector{ET}}, Vector{Int}, Float64}
+    reference_solution
 end
 
 Base.show(io::IO, ::MIME"text/plain", uai::UAIInstance) = Base.show(io, uai)
@@ -69,6 +69,17 @@ function set_evidence!(uai::UAIInstance, pairs::Pair{Int}...)
     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
+
 """
 $(TYPEDEF)
 
@@ -122,6 +133,9 @@ function TensorNetworkModel(
     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,

src/TensorInference.jl

@@ -10,7 +10,6 @@ module TensorInference
 using OMEinsum, LinearAlgebra
 using DocStringExtensions, TropicalNumbers
 # The Tropical GEMM support
-using TropicalGEMM
 using StatsBase
 
 # reexport OMEinsum functions
@@ -19,7 +18,7 @@ export contraction_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABiparti
 
 # read and load uai files
 export read_model_file, read_td_file, read_evidence_file, read_solution_file, read_instance, UAIInstance
-export set_evidence!
+export set_evidence!, set_query!
 
 # marginals
 export TensorNetworkModel, get_vars, get_cards, log_probability, probability, marginals

src/mmap.jl

@@ -58,9 +58,9 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function MMAPModel(instance::UAIInstance; queryvars, openvertices = (), optimizer = GreedyMethod(), simplifier = nothing)::MMAPModel
+function MMAPModel(instance::UAIInstance; openvertices = (), optimizer = GreedyMethod(), simplifier = nothing)::MMAPModel
     return MMAPModel(
-        1:(instance.nvars), instance.cards, instance.factors; queryvars, fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)), optimizer, simplifier, openvertices
+        1:(instance.nvars), instance.cards, instance.factors; queryvars=instance.queryvars, fixedvertices = Dict(zip(instance.obsvars, instance.obsvals)), optimizer, simplifier, openvertices
     )
 end
 

src/sampling.jl

@@ -128,6 +128,13 @@ function sample(tn::TensorNetworkModel, n::Int; usecuda = false)::AbstractMatrix
     return samples.samples
 end
 
+function generate_samples(se::SlicedEinsum, cache::CacheTree{T}, samples, size_dict::Dict) where {T}
+    # slicing is not supported yet.
+    if length(se.slicing) != 0
+        @warn "Slicing is not supported for caching, got nslices = $(length(se.slicing))! Fallback to `NestedEinsum`."
+    end
+    return generate_samples(se.eins, cache, samples, size_dict)
+end
 function generate_samples(code::NestedEinsum, cache::CacheTree{T}, samples, size_dict::Dict) where {T}
     if !OMEinsum.isleaf(code)
         xs = ntuple(i -> cache.siblings[i].content, length(cache.siblings))

src/utils.jl

@@ -246,7 +246,7 @@ end
 
 function read_instance_from_string(uai::AbstractString; eltype = Float64)::UAIInstance
     nvars, cards, ncliques, factors = read_model_string(uai; factor_eltype = eltype)
-    return UAIInstance(nvars, ncliques, cards, factors, Int[], Int[], Vector{eltype}[])
+    return UAIInstance(nvars, ncliques, cards, factors, Int[], Int[], Int[], nothing)
 end
 
 # patch to get content by broadcasting into array, while keep array size unchanged.

test/cuda.jl

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

test/mmap.jl

@@ -15,17 +15,20 @@ end
     tn_ref = TensorNetworkModel(instance; optimizer)
 
     # Does not marginalize any var
-    mmap = MMAPModel(instance; queryvars = collect(1:instance.nvars), optimizer)
+    set_query!(instance, collect(1:instance.nvars))
+    mmap = MMAPModel(instance; optimizer)
     @debug(mmap)
     @test maximum_logp(tn_ref) ≈ maximum_logp(mmap)
 
     # Marginalize all vars
-    mmap2 = MMAPModel(instance; queryvars = Int[], optimizer)
+    set_query!(instance, Int[])
+    mmap2 = MMAPModel(instance; optimizer)
     @debug(mmap2)
     @test Array(probability(tn_ref))[] ≈ exp(maximum_logp(mmap2)[])
 
     # Does not optimize over open vertices
-    mmap3 = MMAPModel(instance; queryvars = setdiff(1:instance.nvars, [2, 4, 6]), optimizer)
+    set_query!(instance, setdiff(1:instance.nvars, [2, 4, 6]))
+    mmap3 = MMAPModel(instance; optimizer)
     @debug(mmap3)
     logp, config = most_probable_config(mmap3)
     @test log_probability(mmap3, config) ≈ logp
@@ -42,7 +45,7 @@ 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; queryvars = instance.queryvars, optimizer)
+      model = MMAPModel(instance; optimizer)
       _, solution = most_probable_config(model)
       @test solution == instance.reference_solution
     end