From 3b6f51d8f30c026809bf4fc42c14e16681a23572 Mon Sep 17 00:00:00 2001 From: Jinguo Liu Date: Thu, 17 Aug 2023 23:18:19 +0800 Subject: [PATCH] Return `Samples` type instead of a matrix (#60) * sample type change * update_evidence! and update_temperature --- Project.toml | 2 +- src/Core.jl | 17 +++++++++++++++++ src/TensorInference.jl | 2 +- src/generictensornetworks.jl | 25 +++++++++++++++++++++++++ src/sampling.jl | 10 ++++++---- test/generictensornetworks.jl | 8 ++++++++ test/mar.jl | 9 +++++++++ test/sampling.jl | 4 ++-- 8 files changed, 69 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index 264c5b8..a1a760b 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "TensorInference" uuid = "c2297e78-99bd-40ad-871d-f50e56b81012" authors = ["Jin-Guo Liu", "Martin Roa Villescas"] -version = "0.2.1" +version = "0.3.0" [deps] Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" diff --git a/src/Core.jl b/src/Core.jl index ac021f0..0127ca2 100644 --- a/src/Core.jl +++ b/src/Core.jl @@ -59,6 +59,23 @@ struct TensorNetworkModel{LT, ET, MT <: AbstractArray} mars::Vector{Vector{LT}} end +""" +$TYPEDSIGNATURES + +Update the evidence of a tensor network model, without changing the set of observed variables! + +### Arguments +- `tnet` is the [`TensorNetworkModel`](@ref) instance. +- `evidence` is the new evidence, the keys must be a subset of existing evidence. +""" +function update_evidence!(tnet::TensorNetworkModel, evidence::Dict) + for (k, v) in evidence + haskey(tnet.evidence, k) || error("`update_evidence!` can only update observed variables!") + tnet.evidence[k] = v + end + return tnet +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], ", ") diff --git a/src/TensorInference.jl b/src/TensorInference.jl index 6fddcd8..bf2bfd5 100644 --- a/src/TensorInference.jl +++ b/src/TensorInference.jl @@ -23,7 +23,7 @@ export problem_from_artifact, ArtifactProblemSpec export read_model, UAIModel, read_evidence, read_solution, read_queryvars, dataset_from_artifact # marginals -export TensorNetworkModel, get_vars, get_cards, log_probability, probability, marginals +export TensorNetworkModel, get_vars, get_cards, log_probability, probability, marginals, update_evidence! # MAP export most_probable_config, maximum_logp diff --git a/src/generictensornetworks.jl b/src/generictensornetworks.jl index 77141f1..e8a9f6f 100644 --- a/src/generictensornetworks.jl +++ b/src/generictensornetworks.jl @@ -1,5 +1,8 @@ using .GenericTensorNetworks: generate_tensors, GraphProblem, flavors, labels +# update models +export update_temperature + """ $TYPEDSIGNATURES @@ -20,6 +23,24 @@ function TensorInference.TensorNetworkModel(problem::GraphProblem, β::Real; evi factors = [Factor((ix...,), t) for (ix, t) in zip(ixs, tensors)] return TensorNetworkModel(lbs, fill(nflavors, length(lbs)), factors; openvars=iy, evidence, optimizer, simplifier, mars) end + +""" +$TYPEDSIGNATURES + +Update the temperature of a tensor network model. +The program will regenerate tensors from the problem, without repeated optimizing the contraction order. + +### Arguments +- `tnet` is the [`TensorNetworkModel`](@ref) instance. +- `problem` is the target constraint satisfiability problem. +- `β` is the inverse temperature. +""" +function update_temperature(tnet::TensorNetworkModel, problem::GraphProblem, β::Real) + tensors = generate_tensors(exp(β), problem) + alltensors = [tnet.tensors[1:end-length(tensors)]..., tensors...] + return TensorNetworkModel(tnet.vars, tnet.code, alltensors, tnet.evidence, tnet.mars) +end + function TensorInference.MMAPModel(problem::GraphProblem, β::Real; queryvars, evidence = Dict{labeltype(problem.code), Int}(), @@ -37,6 +58,10 @@ function TensorInference.MMAPModel(problem::GraphProblem, β::Real; optimizer, simplifier, marginalize_optimizer, marginalize_simplifier) end +function update_temperature(tnet::MMAPModel, problem::GraphProblem, β::Real) + error("We haven't got time to implement setting temperatures for `MMAPModel`. +It is about one or two hours of works. If you need it, please file an issue to let us know: https://github.com/TensorBFS/TensorInference.jl/issues") +end @info "`TensorInference` loaded `GenericTensorNetworks` extension successfully, `TensorNetworkModel` and `MMAPModel` can be used for converting a `GraphProblem` to a probabilistic model now." \ No newline at end of file diff --git a/src/sampling.jl b/src/sampling.jl index 16f8fca..954f4a3 100644 --- a/src/sampling.jl +++ b/src/sampling.jl @@ -9,7 +9,7 @@ The sampled configurations are stored in `samples`, which is a vector of vector. `labels` is a vector of variable names for labeling configurations. The `setmask` is an boolean indicator to denote whether the sampling process of a variable is complete. """ -struct Samples{L} +struct Samples{L} <: AbstractVector{SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}} samples::Matrix{Int} # size is nvars × nsample labels::Vector{L} setmask::BitVector @@ -22,7 +22,9 @@ function setmask!(samples::Samples, eliminated_variables) end return samples end - +Base.getindex(s::Samples, i::Int) = view(s.samples, :, i) +Base.length(s::Samples) = size(s.samples, 2) +Base.size(s::Samples) = (size(s.samples, 2),) idx4labels(totalset, labels)::Vector{Int} = map(v->findfirst(==(v), totalset), labels) """ @@ -99,7 +101,7 @@ Returns a vector of vector, each element being a configurations defined on `get_ * `tn` is the tensor network model. * `n` is the number of samples to be returned. """ -function sample(tn::TensorNetworkModel, n::Int; usecuda = false)::AbstractMatrix{Int} +function sample(tn::TensorNetworkModel, n::Int; usecuda = false)::Samples # generate tropical tensors with its elements being log(p). xs = adapt_tensors(tn; usecuda, rescale = false) # infer size from the contraction code and the input tensors `xs`, returns a label-size dictionary. @@ -125,7 +127,7 @@ function sample(tn::TensorNetworkModel, n::Int; usecuda = false)::AbstractMatrix idx = findfirst(==(k), labels) samples.samples[idx, :] .= v end - return samples.samples + return samples end function generate_samples(se::SlicedEinsum, cache::CacheTree{T}, samples, size_dict::Dict) where {T} diff --git a/test/generictensornetworks.jl b/test/generictensornetworks.jl index 4031624..56454f2 100644 --- a/test/generictensornetworks.jl +++ b/test/generictensornetworks.jl @@ -12,6 +12,14 @@ using GenericTensorNetworks, TensorInference mars2 = TensorInference.normalize!(GenericTensorNetworks.solve(problem2, PartitionFunction(β)), 1) @test mars ≈ mars2 + # update temperature + β2 = 3.0 + model = update_temperature(model, problem, β2) + pa = probability(model)[] + model2 = TensorNetworkModel(problem, β2) + pb = probability(model2)[] + @test pa ≈ pb + # mmap model = MMAPModel(problem, β; queryvars=[1,4]) logp, config = most_probable_config(model) diff --git a/test/mar.jl b/test/mar.jl index bd1ee20..352c080 100644 --- a/test/mar.jl +++ b/test/mar.jl @@ -122,4 +122,13 @@ end tnet34 = TensorNetworkModel(model; openvars=[3,4]) @test mars[1] ≈ probability(tnet23) @test mars[2] ≈ probability(tnet34) + + tnet1 = TensorNetworkModel(model; mars=[[2, 3], [3, 4]], evidence=Dict(3=>1)) + tnet2 = TensorNetworkModel(model; mars=[[2, 3], [3, 4]], evidence=Dict(3=>0)) + mars1 = marginals(tnet1) + mars2 = marginals(tnet2) + update_evidence!(tnet1, Dict(3=>0)) + mars1b = marginals(tnet1) + @test !(mars1 ≈ mars2) + @test mars1b ≈ mars2 end \ No newline at end of file diff --git a/test/sampling.jl b/test/sampling.jl index 12549f2..eca855c 100644 --- a/test/sampling.jl +++ b/test/sampling.jl @@ -50,13 +50,13 @@ using TensorInference, Test tnet = TensorNetworkModel(model) samples = sample(tnet, n) mars = getindex.(marginals(tnet), 2) - mars_sample = [count(i->samples[k, i]==(1), axes(samples, 2)) for k=1:8] ./ n + mars_sample = [count(s->s[k]==(1), samples) for k=1:8] ./ n @test isapprox(mars, mars_sample, atol=0.05) # fix the evidence tnet = TensorNetworkModel(model, optimizer=TreeSA(), evidence=Dict(7=>1)) samples = sample(tnet, n) mars = getindex.(marginals(tnet), 1) - mars_sample = [count(i->samples[k, i]==(0), axes(samples, 2)) for k=1:8] ./ n + mars_sample = [count(s->s[k]==(0), samples) for k=1:8] ./ n @test isapprox([mars[1:6]..., mars[8]], [mars_sample[1:6]..., mars_sample[8]], atol=0.05) end \ No newline at end of file