update

Author: GiggleLiu

Project.toml

@@ -12,6 +12,7 @@ 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"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TropicalNumbers = "b3a74e9c-7526-4576-a4eb-79c0d4c32334"
@@ -25,6 +26,7 @@ LinearAlgebra = "1"
 OMEinsum = "0.8"
 Pkg = "1"
 PrecompileTools = "1"
+PrettyTables = "2"
 Requires = "1"
 StatsBase = "0.34"
 TropicalNumbers = "0.5.4, 0.6"

src/RescaledArray.jl

@@ -23,12 +23,12 @@ $(TYPEDSIGNATURES)
 Returns a rescaled array that equivalent to the input tensor.
 """
 function rescale_array(tensor::AbstractArray{T})::RescaledArray where {T}
-    maxf = maximum(tensor)
+    maxf = maximum(abs, tensor)
     if iszero(maxf)
         @warn("The maximum value of the array to rescale is 0!")
         return RescaledArray(zero(T), tensor)
     end
-    return RescaledArray(log(maxf), OMEinsum.asarray(tensor ./ maxf, tensor))
+    return RescaledArray(T(log(maxf)), OMEinsum.asarray(tensor ./ maxf, tensor))
 end
 
 for CT in [:DynamicEinCode, :StaticEinCode]

src/TensorInference.jl

@@ -11,6 +11,7 @@ using OMEinsum, LinearAlgebra
 using DocStringExtensions, TropicalNumbers
 # The Tropical GEMM support
 using StatsBase
+using PrettyTables
 import Pkg
 
 # reexport OMEinsum functions

src/mar.jl

@@ -16,10 +16,10 @@ end
 # `CacheTree` stores intermediate `NestedEinsum` contraction results.
 # It is a tree structure that isomorphic to the contraction tree,
 # `content` is the cached intermediate contraction result.
-# `siblings` are the siblings of current node.
+# `children` are the children of current node, e.g. tensors that are contracted to get `content`.
 struct CacheTree{T}
     content::AbstractArray{T}
-    siblings::Vector{CacheTree{T}}
+    children::Vector{CacheTree{T}}
 end
 
 function cached_einsum(se::SlicedEinsum, @nospecialize(xs), size_dict)
@@ -62,7 +62,7 @@ function generate_gradient_tree(code::NestedEinsum, cache::CacheTree{T}, dy::Abs
     if OMEinsum.isleaf(code)
         return CacheTree(dy, CacheTree{T}[])
     else
-        xs = ntuple(i -> cache.siblings[i].content, length(cache.siblings))
+        xs = ntuple(i -> cache.children[i].content, length(cache.children))
         # `einsum_grad` is the back-propagation rule for einsum function.
         # If the forward pass is `y = einsum(EinCode(inputs_labels, output_labels), (A, B, ...), size_dict)`
         # Then the back-propagation pass is
@@ -73,7 +73,7 @@ function generate_gradient_tree(code::NestedEinsum, cache::CacheTree{T}, dy::Abs
         # ```
         # Let `L` be the loss, we will have `y̅ := ∂L/∂y`, `A̅ := ∂L/∂A`...
         dxs = einsum_backward_rule(code.eins, xs, cache.content, size_dict, dy)
-        return CacheTree(dy, generate_gradient_tree.(code.args, cache.siblings, dxs, Ref(size_dict)))
+        return CacheTree(dy, generate_gradient_tree.(code.args, cache.children, dxs, Ref(size_dict)))
     end
 end
 
@@ -116,7 +116,7 @@ function extract_leaves!(code, cache, res)
         res[code.tensorindex] = cache.content
     else
         # resurse deeper
-        extract_leaves!.(code.args, cache.siblings, Ref(res))
+        extract_leaves!.(code.args, cache.children, Ref(res))
     end
     return res
 end

src/sampling.jl

@@ -25,8 +25,22 @@ 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),)
+function Base.show(io::IO, s::Samples)  # display with PrettyTables
+    println(io, typeof(s))
+    PrettyTables.pretty_table(io, s.samples', header=s.labels)
+end
+num_samples(samples::Samples) = size(samples.samples, 2)
 eliminated_variables(samples::Samples) = samples.labels[samples.setmask]
-idx4labels(totalset, labels)::Vector{Int} = map(v->findfirst(==(v), totalset), labels)
+is_eliminated(samples::Samples{L}, var::L) where L = samples.setmask[findfirst(==(var), samples.labels)]
+function idx4labels(totalset::AbstractVector{L}, labels::AbstractVector{L})::Vector{Int} where L
+    map(v->findfirst(==(v), totalset), labels)
+end
+idx4labels(samples::Samples{L}, lb::L) where L = findfirst(==(lb), samples.labels)
+function subset(samples::Samples{L}, labels::AbstractVector{L}) where L
+    idx = idx4labels(samples.labels, labels)
+    @assert all(i->samples.setmask[i], idx)
+    return samples.samples[idx, :]
+end
 
 """
 $(TYPEDSIGNATURES)
@@ -39,34 +53,49 @@ The backward process for sampling configurations.
 * `samples` is the samples generated for eliminated variables,
 * `size_dict` is a key-value map from tensor label to dimension size.
 """
-function backward_sampling!(code::EinCode, @nospecialize(xs::Tuple), @nospecialize(y), @nospecialize(env), samples::Samples, size_dict)
+function backward_sampling!(code::EinCode, @nospecialize(xs::Tuple), @nospecialize(env), samples::Samples, size_dict)
     ixs, iy = getixsv(code), getiyv(code)
-    el = setdiff(vcat(ixs...), iy)
+    el = setdiff(vcat(ixs...), iy) ∩ samples.labels
+
     # get probability
     prob_code = optimize_code(EinCode([ixs..., iy], el), size_dict, GreedyMethod(; nrepeat=1))
+    el_prev = eliminated_variables(samples)
+    xs = [eliminate_dimensions(x, ix, el_prev=>subset(samples, el_prev)[:,1]) for (ix, x) in zip(ixs, xs)]
     probabilities = einsum(prob_code, (xs..., env), size_dict)
 
     # sample from the probability tensor
     totalset = CartesianIndices((map(x->size_dict[x], el)...,))
     eliminated_locs = idx4labels(samples.labels, el)
-    for i=axes(samples.samples, 2)
-        config = StatsBase.sample(totalset, Weights(vec(selectdim(probabilities, ndims(probabilities), i))))
-        samples.samples[eliminated_locs, i] .= config.I .- 1
-    end
+    config = StatsBase.sample(totalset, _Weights(vec(probabilities)))
+    samples.samples[eliminated_locs, 1] .= config.I .- 1
 
     # eliminate the sampled variables
     set_eliminated!(samples, el)
     for l in el
         size_dict[l] = 1
     end
-    for sample in sampels
-        map(x->eliminate_dimensions!(x, el=>sample), xs)
-    end
+    sub = subset(samples, el)[:, 1]
+    xs = [eliminate_dimensions(x, ix, el=>sub) for (ix, x) in zip(ixs, xs)]
+    env = eliminate_dimensions(env, iy, el=>sub)
 
     # update environment
-    for (i, ix) in enumerate(ixs)
+    return map(1:length(ixs)) do i
+        rest = setdiff(1:length(ixs), i)
+        code = optimize_code(EinCode([ixs[rest]..., iy], ixs[i]), size_dict, GreedyMethod(; nrepeat=1))
+        einsum(code, (xs[rest]..., env), size_dict)
     end
-    return envs
+end
+
+function eliminate_dimensions(x::AbstractArray{T, N}, ix::AbstractVector{L}, el::Pair{<:AbstractVector{L}}) where {T, N, L}
+    idx = ntuple(N) do i
+        if ix[i] ∈ el.first
+            k = el.second[findfirst(==(ix[i]), el.first)] + 1
+            k:k
+        else
+            1:size(x, i)
+        end
+    end
+    return asarray(x[idx...], x)
 end
 
 function addbatch(samples::Samples, eliminated_variables)
@@ -113,48 +142,54 @@ 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)::Samples
+function sample(tn::TensorNetworkModel, n::Int; usecuda = false, queryvars = get_vars(tn))::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.
     size_dict = OMEinsum.get_size_dict!(getixsv(tn.code), xs, Dict{Int, Int}())
     # forward compute and cache intermediate results.
     cache = cached_einsum(tn.code, xs, size_dict)
     # initialize `y̅` as the initial batch of samples.
-    labels = get_vars(tn)
     iy = getiyv(tn.code)
-    setmask = falses(length(labels))
-    idx = map(l->findfirst(==(l), labels), iy)
+    setmask = falses(length(queryvars))
+    idx = map(l->findfirst(==(l), queryvars), iy ∩ queryvars)
     setmask[idx] .= true
-    indices = StatsBase.sample(CartesianIndices(size(cache.content)), Weights(normalize!(vec(LinearAlgebra.normalize!(cache.content)))), n)
-    configs = zeros(Int, length(labels), n)
+    indices = StatsBase.sample(CartesianIndices(size(cache.content)), _Weights(vec(cache.content)), n)
+    configs = zeros(Int, length(queryvars), n)
     for i=1:n
         configs[idx, i] .= indices[i].I .- 1
     end
-    samples = Samples(configs, labels, setmask)
+    samples = Samples(configs, queryvars, setmask)
     # back-propagate
-    generate_samples(tn.code, cache, samples, size_dict)
+    env = copy(cache.content)
+    fill!(env, one(eltype(env)))
+    generate_samples!(tn.code, cache, env, samples, size_dict)
     # set evidence variables
     for (k, v) in tn.evidence
-        idx = findfirst(==(k), labels)
+        idx = findfirst(==(k), samples.labels)
         samples.samples[idx, :] .= v
     end
     return samples
 end
+_Weights(x::AbstractVector{<:Real}) = Weights(x)
+function _Weights(x::AbstractArray{<:Complex})
+    @assert all(e->abs(imag(e)) < 100*eps(abs(e)), x)
+    return Weights(real.(x))
+end
 
-function generate_samples(se::SlicedEinsum, cache::CacheTree{T}, samples, size_dict::Dict) where {T}
+function generate_samples!(se::SlicedEinsum, cache::CacheTree{T}, env::AbstractArray{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)
+    return generate_samples!(se.eins, cache, env, samples, size_dict)
 end
-function generate_samples(code::NestedEinsum, cache::CacheTree{T}, env::AbstractArray{T}, samples, size_dict::Dict) where {T}
+function generate_samples!(code::NestedEinsum, cache::CacheTree{T}, env::AbstractArray{T}, samples::Samples, size_dict::Dict) where {T}
     if !OMEinsum.isleaf(code)
-        xs = ntuple(i -> cache.siblings[i].content, length(cache.siblings))
-        envs = backward_sampling!(code.eins, xs, cache.content, env, samples, copy(size_dict))
-        for (arg, sib, env) in zip(code.args, cache.siblings, envs)
-            generate_samples(arg, sib, env, samples, size_dict)
+        xs = ntuple(i -> cache.children[i].content, length(cache.children))
+        envs = backward_sampling!(code.eins, xs, env, samples, size_dict)
+        for (arg, sib, env) in zip(code.args, cache.children, envs)
+            generate_samples!(arg, sib, env, samples, size_dict)
         end
     end
 end

test/sampling.jl

@@ -1,5 +1,6 @@
 using TensorInference, Test
 using StatsBase: kldivergence
+using OMEinsum
 
 @testset "sampling" begin
     model = TensorInference.read_model_from_string("""MARKOV
@@ -64,21 +65,26 @@ end
 
 @testset "sample MPS" begin
     tensors = [
-        [rand(2, 2) for i=1:2],
-        [rand(2, 2, 2) for i=1:2],
-        [rand(2, 2, 2) for i=1:2],
-        [rand(2, 2) for i=1:2],
+        randn(ComplexF64, 2, 3),
+        randn(ComplexF64, 3, 2, 3),
+        randn(ComplexF64, 3, 2, 3),
+        randn(ComplexF64, 3, 2),
     ]
+    tensors = [tensors..., conj.(tensors)...]
     ixs = [[1, 5], [5, 2, 6], [6, 3, 7], [7, 4], [1, 8], [8, 2, 9], [9, 3, 10], [10, 4]]
     mps = TensorNetworkModel(
         collect(1:10),
-        DynamicEinCode(ixs, Int[]),
-        [tensors..., conj.(tensors)...],
+        optimize_code(DynamicEinCode(ixs, Int[]), OMEinsum.get_size_dict(ixs, tensors), GreedyMethod()),
+        tensors,
         Dict{Int, Int}(),
-        collect(5:10)
+        [[i] for i=5:10]
     )
-    samples = sample(mps, 1000)
-    indices = samples.samples
+    num_samples = 1
+    samples = sample(mps, num_samples; queryvars=[1, 2, 3, 4])
+    indices = map(samples) do sample
+        sum(i->sample[i] * 2^(i-1), 1:4) + 1
+    end
+    @show samples
     @show indices
     probs = vec(DynamicEinCode(ixs, collect(1:4))(tensors...))
     negative_loglikelyhood(samples, probs) = -sum(log.(probs[indices]))