diff --git a/Project.toml b/Project.toml index 90c66c5..280abc7 100644 --- a/Project.toml +++ b/Project.toml @@ -9,7 +9,6 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" -PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" ProblemReductions = "899c297d-f7d2-4ebf-8815-a35996def416" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" @@ -28,7 +27,6 @@ DocStringExtensions = "0.8.6, 0.9" LinearAlgebra = "1" OMEinsum = "0.8" Pkg = "1" -PrecompileTools = "1" PrettyTables = "2" ProblemReductions = "0.3" StatsBase = "0.34" diff --git a/examples/hard-core-lattice-gas/main.jl b/examples/hard-core-lattice-gas/main.jl index 14cc289..0739442 100644 --- a/examples/hard-core-lattice-gas/main.jl +++ b/examples/hard-core-lattice-gas/main.jl @@ -62,7 +62,7 @@ mars = marginals(pmodel) show_graph(SimpleGraph(graph), sites; vertex_colors=[(b = mars[[i]][2]; (1-b, 1-b, 1-b)) for i in 1:nv(graph)], texts=fill("", nv(graph))) # The can see the sites at the corner is more likely to be occupied. # To obtain two-site correlations, one can set the variables to query marginal probabilities manually. -pmodel2 = TensorNetworkModel(problem, β; mars=[[e.src, e.dst] for e in edges(graph)]) +pmodel2 = TensorNetworkModel(problem, β; unity_tensors_labels = [[e.src, e.dst] for e in edges(graph)]) mars = marginals(pmodel2); # We show the probability that both sites on an edge are not occupied diff --git a/src/Core.jl b/src/Core.jl index dcaf26f..39866fe 100644 --- a/src/Core.jl +++ b/src/Core.jl @@ -45,18 +45,18 @@ $(TYPEDEF) Probabilistic modeling with a tensor network. ### Fields -* `vars` are the degrees of freedom in the tensor network. +* `nvars` are the number of variables in the tensor network. * `code` is the tensor network contraction pattern. -* `tensors` are the tensors fed into the tensor network, the leading tensors are unity tensors associated with `mars`. +* `tensors` are the tensors fed into the tensor network, the leading tensors are unity tensors associated with `unity_tensors_labels`. * `evidence` is a dictionary used to specify degrees of freedom that are fixed to certain values. -* `mars` is a vector, each element is a vector of variables to compute marginal probabilities. +* `unity_tensors_idx` is a vector of indices of the unity tensors in the `tensors` array. Unity tensors are dummy tensors used to obtain the marginal probabilities. """ -struct TensorNetworkModel{LT, ET, MT <: AbstractArray} - vars::Vector{LT} +struct TensorNetworkModel{ET, MT <: AbstractArray} + nvars::Int code::ET tensors::Vector{MT} - evidence::Dict{LT, Int} - mars::Vector{Vector{LT}} + evidence::Dict{Int, Int} + unity_tensors_idx::Vector{Int} end """ @@ -78,7 +78,7 @@ end function Base.show(io::IO, tn::TensorNetworkModel) open = getiyv(tn.code) - variables = join([string_var(var, open, tn.evidence) for var in tn.vars], ", ") + variables = join([string_var(var, open, tn.evidence) for var in get_vars(tn)], ", ") tc, sc, rw = contraction_complexity(tn) println(io, "$(typeof(tn))") println(io, "variables: $variables") @@ -110,84 +110,25 @@ $(TYPEDSIGNATURES) * `evidence` is a dictionary of evidences, the values are integers start counting from 0. * `optimizer` is the tensor network contraction order optimizer, please check the package [`OMEinsumContractionOrders.jl`](https://github.com/TensorBFS/OMEinsumContractionOrders.jl) for available algorithms. * `simplifier` is some strategies for speeding up the `optimizer`, please refer the same link above. -* `mars` is a list of marginal probabilities. It is all single variables by default, i.e. `[[1], [2], ..., [n]]`. One can also specify multi-variables, which may increase the computational complexity. +* `unity_tensors_labels` is a list of labels for the unity tensors. It is all single variables by default, i.e. `[[1], [2], ..., [n]]`. One can also specify multi-variables, which may increase the computational complexity. """ function TensorNetworkModel( - model::UAIModel; + model::UAIModel{ET, FT}; openvars = (), evidence = Dict{Int,Int}(), optimizer = GreedyMethod(), simplifier = nothing, - mars = [[i] for i=1:model.nvars] -)::TensorNetworkModel - return TensorNetworkModel( - 1:(model.nvars), - model.cards, - model.factors; - openvars, - evidence, - optimizer, - simplifier, - mars - ) -end - -""" -$(TYPEDSIGNATURES) -""" -function TensorNetworkModel( - vars::AbstractVector{LT}, - cards::AbstractVector{Int}, - factors::Vector{<:Factor{T}}; - openvars = (), - evidence = Dict{LT, Int}(), - optimizer = GreedyMethod(), - simplifier = nothing, - mars = [[v] for v in vars] -)::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([mars..., [[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 mar]...) for mar in mars]..., [t.vals for t in factors]...] - return TensorNetworkModel(collect(LT, vars), rawcode, tensors; evidence, optimizer, simplifier, mars) -end - -""" -$(TYPEDSIGNATURES) -""" -function TensorNetworkModel( - vars::AbstractVector{LT}, - rawcode::EinCode, - tensors::Vector{<:AbstractArray}; - evidence = Dict{LT, Int}(), - optimizer = GreedyMethod(), - simplifier = nothing, - mars = [[v] for v in vars] -)::TensorNetworkModel where {LT} + unity_tensors_labels = [[i] for i=1:model.nvars] +) where {ET, FT} # `optimize_code` optimizes the contraction order of a raw tensor network without a contraction order specified. # The 1st argument is the contraction pattern to be optimized (without contraction order). # The 2nd arugment is the size dictionary, which is a label-integer dictionary. # The 3rd and 4th arguments are the optimizer and simplifier that configures which algorithm to use and simplify. + rawcode = EinCode([unity_tensors_labels..., [[factor.vars...] for factor in model.factors]...], collect(Int, openvars)) # labels for vertex tensors (unity tensors) and edge tensors + tensors = Array{ET}[[ones(ET, [model.cards[i] for i in lb]...) for lb in unity_tensors_labels]..., [t.vals for t in model.factors]...] size_dict = OMEinsum.get_size_dict(getixsv(rawcode), tensors) code = optimize_code(rawcode, size_dict, optimizer, simplifier) - TensorNetworkModel(collect(LT, vars), code, tensors, evidence, mars) -end - -""" -$(TYPEDSIGNATURES) -""" -function TensorNetworkModel( - model::UAIModel{T}, code; - evidence = Dict{Int,Int}(), - mars = [[i] for i=1:model.nvars], - vars = [1:model.nvars...] -)::TensorNetworkModel where{T} - @debug "constructing tensor network model from code" - tensors = Array{T}[[ones(T, [model.cards[i] for i in mar]...) for mar in mars]..., [t.vals for t in model.factors]...] - - return TensorNetworkModel(vars, code, tensors, evidence, mars) + return TensorNetworkModel(model.nvars, code, tensors, evidence, collect(Int, 1:length(unity_tensors_labels))) end """ @@ -195,17 +136,16 @@ $(TYPEDSIGNATURES) Get the variables in this tensor network, they are also known as legs, labels, or degree of freedoms. """ -get_vars(tn::TensorNetworkModel)::Vector = tn.vars +get_vars(tn::TensorNetworkModel)::Vector = 1:tn.nvars """ $(TYPEDSIGNATURES) -Get the cardinalities of variables in this tensor network. +Get the ardinalities of variables in this tensor network. """ function get_cards(tn::TensorNetworkModel; fixedisone = false)::Vector - vars = get_vars(tn) size_dict = OMEinsum.get_size_dict(getixsv(tn.code), tn.tensors) - [fixedisone && haskey(tn.evidence, vars[k]) ? 1 : size_dict[vars[k]] for k in eachindex(vars)] + [fixedisone && haskey(tn.evidence, k) ? 1 : size_dict[k] for k in 1:tn.nvars] end chevidence(tn::TensorNetworkModel, evidence) = TensorNetworkModel(tn.vars, tn.code, tn.tensors, evidence) diff --git a/src/TensorInference.jl b/src/TensorInference.jl index 19125ba..31d53bd 100644 --- a/src/TensorInference.jl +++ b/src/TensorInference.jl @@ -52,13 +52,4 @@ include("mmap.jl") include("sampling.jl") include("cspmodels.jl") -# 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-network/main.jl") -# end -# end - end # module diff --git a/src/cspmodels.jl b/src/cspmodels.jl index 128f684..dbd7b58 100644 --- a/src/cspmodels.jl +++ b/src/cspmodels.jl @@ -28,10 +28,11 @@ Convert a constraint satisfiability problem (or energy model) to a probabilistic * `mars` is the list of variables to be marginalized. """ function TensorNetworkModel(problem::ConstraintSatisfactionProblem, β::T; evidence::Dict=Dict{Int,Int}(), - optimizer=GreedyMethod(), openvars=Int[], simplifier=nothing, mars=[[l] for l in variables(problem)]) where T <: Real + optimizer=GreedyMethod(), openvars=Int[], simplifier=nothing, unity_tensors_labels = [[l] for l in variables(problem)]) where T <: Real tensors, ixs = generate_tensors(β, problem) factors = [Factor((ix...,), t) for (ix, t) in zip(ixs, tensors)] - return TensorNetworkModel(variables(problem), fill(num_flavors(problem), num_variables(problem)), factors; openvars, evidence, optimizer, simplifier, mars) + model = UAIModel(num_variables(problem), fill(num_flavors(problem), num_variables(problem)), factors) + return TensorNetworkModel(model; openvars, evidence, optimizer, simplifier, unity_tensors_labels) end """ @@ -47,8 +48,9 @@ The program will regenerate tensors from the problem, without repeated optimizin """ function update_temperature(tnet::TensorNetworkModel, problem::ConstraintSatisfactionProblem, β::Real) tensors, ixs = generate_tensors(β, problem) - alltensors = [tnet.tensors[1:length(tnet.mars)]..., tensors...] - return TensorNetworkModel(tnet.vars, tnet.code, alltensors, tnet.evidence, tnet.mars) + @assert tnet.unity_tensors_idx == collect(1:length(tnet.unity_tensors_idx)) "The target tensor network can not be updated! Got `unity_tensors_idx = $(tnet.unity_tensors_idx)`" + alltensors = [tnet.tensors[tnet.unity_tensors_idx]..., tensors...] + return TensorNetworkModel(tnet.nvars, tnet.code, alltensors, tnet.evidence, tnet.unity_tensors_idx) end function MMAPModel(problem::ConstraintSatisfactionProblem, β::Real; diff --git a/src/map.jl b/src/map.jl index 372133f..a1f5762 100644 --- a/src/map.jl +++ b/src/map.jl @@ -53,13 +53,18 @@ $(TYPEDSIGNATURES) Returns the largest log-probability and the most probable configuration. """ function most_probable_config(tn::TensorNetworkModel; usecuda = false)::Tuple{Real, Vector} - expected_mars = [[l] for l in get_vars(tn)] - @assert tn.mars[1:length(expected_mars)] == expected_mars "To get the the most probable configuration, the leading elements of `tn.vars` must be `$expected_mars`" - vars = get_vars(tn) + tensor_indices = check_queryvars(tn, [[v] for v in 1:tn.nvars]) 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.evidence, vars[k]) ? tn.evidence[vars[k]] : argmax(grads[k]) - 1, 1:length(vars)) + return content(Array(logp)[]), map(k -> haskey(tn.evidence, k) ? tn.evidence[k] : argmax(grads[tensor_indices[k]]) - 1, 1:tn.nvars) +end +# check if the queryvars are included in the unity tensors labels, if yes, return the indices of the unity tensors +function check_queryvars(tn::TensorNetworkModel, queryvars::AbstractVector{Vector{Int}}) + ixs = OMEinsum.getixsv(tn.code) + indices = [findfirst(==(l), ixs[tn.unity_tensors_idx]) for l in queryvars] + @assert !any(isnothing, indices) "To get the the most probable configuration, the unity tensors labels must include all variables. Query variables: $queryvars, Unity tensors labels: $(ixs[tn.unity_tensors_idx])" + return tn.unity_tensors_idx[indices] end """ diff --git a/src/mar.jl b/src/mar.jl index a436d3d..b1b65ae 100644 --- a/src/mar.jl +++ b/src/mar.jl @@ -130,14 +130,16 @@ are their respective marginals. A marginal is a probability distribution over a subset of variables, obtained by integrating or summing over the remaining variables in the model. By default, the function returns the marginals of all individual variables. To specify which marginal variables to query, set the -`mars` field when constructing a [`TensorNetworkModel`](@ref). Note that +`unity_tensors_labels` field when constructing a [`TensorNetworkModel`](@ref). Note that the choice of marginal variables will affect the contraction order of the tensor network. ### Arguments - `tn`: The [`TensorNetworkModel`](@ref) to query. -- `usecuda`: Specifies whether to use CUDA for tensor contraction. -- `rescale`: Specifies whether to rescale the tensors during contraction. + +### Keyword Arguments +- `usecuda::Bool`: Specifies whether to use CUDA for tensor contraction. +- `rescale::Bool`: Specifies whether to rescale the tensors during contraction. ### Example The following example is taken from [`examples/asia-network/main.jl`](https://tensorbfs.github.io/TensorInference.jl/dev/generated/asia-network/main/). @@ -158,7 +160,7 @@ Dict{Vector{Int64}, Vector{Float64}} with 8 entries: [7] => [0.145092, 0.854908] [2] => [0.05, 0.95] -julia> tn2 = TensorNetworkModel(model; evidence=Dict(1=>0), mars=[[2, 3], [3, 4]]); +julia> tn2 = TensorNetworkModel(model; evidence=Dict(1=>0), unity_tensors_labels = [[2, 3], [3, 4]]); julia> marginals(tn2) Dict{Vector{Int64}, Matrix{Float64}} with 2 entries: @@ -186,9 +188,11 @@ function marginals(tn::TensorNetworkModel; usecuda = false, rescale = true)::Dic # sometimes, the cost can overflow, then we need to rescale the tensors during contraction. cost, grads = cost_and_gradient(tn.code, adapt_tensors(tn; usecuda, rescale)) @debug "cost = $cost" + ixs = OMEinsum.getixsv(tn.code) + queryvars = ixs[tn.unity_tensors_idx] if rescale - return Dict(zip(tn.mars, LinearAlgebra.normalize!.(getfield.(grads[1:length(tn.mars)], :normalized_value), 1))) + return Dict(zip(queryvars, LinearAlgebra.normalize!.(getfield.(grads[tn.unity_tensors_idx], :normalized_value), 1))) else - return Dict(zip(tn.mars, LinearAlgebra.normalize!.(grads[1:length(tn.mars)], 1))) + return Dict(zip(queryvars, LinearAlgebra.normalize!.(grads[tn.unity_tensors_idx], 1))) end -end +end \ No newline at end of file diff --git a/src/utils.jl b/src/utils.jl index ce90819..0c67e3b 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -352,11 +352,11 @@ function random_matrix_product_state(::Type{T}, n::Int, chi::Int, d::Int=2) wher push!(ixs_bra, [virtual_indices_bra[n-1], physical_indices[n]]) tensors, ixs = [tensors..., conj.(tensors)...], [ixs_ket..., ixs_bra...] return TensorNetworkModel( - collect(1:3n-2), + 3n-2, optimize_code(DynamicEinCode(ixs, Int[]), OMEinsum.get_size_dict(ixs, tensors), GreedyMethod()), tensors, Dict{Int, Int}(), - Vector{Int}[[i] for i=1:n] + collect(1:n) ) end random_matrix_product_state(n::Int, chi::Int, d::Int=2) = random_matrix_product_state(ComplexF64, n, chi, d) diff --git a/test/cspmodels.jl b/test/cspmodels.jl index c7b559b..cc26457 100644 --- a/test/cspmodels.jl +++ b/test/cspmodels.jl @@ -7,7 +7,7 @@ using GenericTensorNetworks β = 2.0 g = GenericTensorNetworks.Graphs.smallgraph(:petersen) problem = IndependentSet(g) - model = TensorNetworkModel(problem, β; mars=[[2, 3]]) + model = TensorNetworkModel(problem, β; unity_tensors_labels = [[2, 3]]) mars = marginals(model)[[2, 3]] problem2 = IndependentSet(g) mars2 = TensorInference.normalize!(GenericTensorNetworks.solve(GenericTensorNetwork(problem2; openvertices=[2, 3]), PartitionFunction(β)), 1) @@ -28,7 +28,7 @@ using GenericTensorNetworks β = 1.0 problem = SpinGlass(g, -ones(Int, ne(g)), zeros(Int, nv(g))) - model = TensorNetworkModel(problem, β; mars=[[2, 3]]) + model = TensorNetworkModel(problem, β; unity_tensors_labels = [[2, 3]]) samples = sample(model, 100) @test sum(energy.(Ref(problem), samples))/100 <= -14 end \ No newline at end of file diff --git a/test/map.jl b/test/map.jl index abc0c98..2ee148d 100644 --- a/test/map.jl +++ b/test/map.jl @@ -2,16 +2,14 @@ using Test using OMEinsum using TensorInference -@testset "load from code" begin +@testset "load from model" begin model = problem_from_artifact("uai2014", "MAR", "Promedus", 14) tn1 = TensorNetworkModel(read_model(model); evidence=read_evidence(model), optimizer = TreeSA(ntrials = 3, niters = 2, βs = 1:0.1:80)) - tn2 = TensorNetworkModel(read_model(model), tn1.code, evidence=read_evidence(model)) - - @test tn1.code == tn2.code + @test tn1 isa TensorNetworkModel end @testset "gradient-based tensor network solvers" begin diff --git a/test/mar.jl b/test/mar.jl index 7aa10b4..1302ca7 100644 --- a/test/mar.jl +++ b/test/mar.jl @@ -1,6 +1,5 @@ using Test using OMEinsum -using KaHyPar using TensorInference @testset "composite number" begin @@ -116,7 +115,7 @@ end 0.1 0.3 0.2 0.9 """) n = 10000 - tnet = TensorNetworkModel(model; mars=[[2, 3], [3, 4]]) + tnet = TensorNetworkModel(model; unity_tensors_labels = [[2, 3], [3, 4]]) mars = marginals(tnet) tnet23 = TensorNetworkModel(model; openvars=[2,3]) tnet34 = TensorNetworkModel(model; openvars=[3,4]) @@ -124,8 +123,8 @@ end @test mars[[3, 4]] ≈ probability(tnet34) vars = [[2, 4], [3, 5]] - tnet1 = TensorNetworkModel(model; mars=vars, evidence=Dict(3=>1)) - tnet2 = TensorNetworkModel(model; mars=vars, evidence=Dict(3=>0)) + tnet1 = TensorNetworkModel(model; unity_tensors_labels = vars, evidence=Dict(3=>1)) + tnet2 = TensorNetworkModel(model; unity_tensors_labels = vars, evidence=Dict(3=>0)) mars1 = marginals(tnet1) mars2 = marginals(tnet2) update_evidence!(tnet1, Dict(3=>0)) diff --git a/test/pr.jl b/test/pr.jl index 53c1c1e..f6c0c08 100644 --- a/test/pr.jl +++ b/test/pr.jl @@ -1,6 +1,5 @@ using Test using OMEinsum -using KaHyPar using TensorInference @testset "UAI Reference Solution Comparison" begin diff --git a/test/sampling.jl b/test/sampling.jl index 0f06320..005dc87 100644 --- a/test/sampling.jl +++ b/test/sampling.jl @@ -70,10 +70,12 @@ end Random.seed!(140) mps = random_matrix_product_state(n, chi) num_samples = 10000 + ixs = OMEinsum.getixsv(mps.code) + @show ixs samples = map(1:num_samples) do i - sample(mps, 1; queryvars=vcat(mps.mars...)).samples[:,1] + sample(mps, 1; queryvars=collect(1:n)).samples[:,1] end - samples = sample(mps, num_samples; queryvars=vcat(mps.mars...)) + samples = sample(mps, num_samples; queryvars=collect(1:n)) indices = map(samples) do sample sum(i->sample[i] * 2^(i-1), 1:n) + 1 end