Return Samples type instead of a matrix (#60)

* sample type change

* update_evidence! and update_temperature

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"

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], ", ")

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

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."

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}

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)

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

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