update matrix product state

Author: GiggleLiu

src/TensorInference.jl

@@ -35,6 +35,9 @@ export sample
 # MMAP
 export MMAPModel
 
+# utils
+export matrix_product_state
+
 include("Core.jl")
 include("RescaledArray.jl")
 include("utils.jl")

src/sampling.jl

@@ -70,7 +70,7 @@ function sample(tn::TensorNetworkModel, n::Int; usecuda = false, queryvars = get
     # back-propagate
     env = copy(cache.content)
     fill!(env, one(eltype(env)))
-    generate_samples!(tn.code, cache, env, samples, samples.labels, size_dict)
+    generate_samples!(tn.code, cache, env, samples, copy(samples.labels), size_dict)  # note: `copy` is necessary
     # set evidence variables
     for (k, v) in tn.evidence
         idx = findfirst(==(k), samples.labels)
@@ -80,7 +80,7 @@ function sample(tn::TensorNetworkModel, n::Int; usecuda = false, queryvars = get
 end
 _Weights(x::AbstractVector{<:Real}) = Weights(x)
 function _Weights(x::AbstractArray{<:Complex})
-    @assert all(e->abs(imag(e)) < 100*eps(abs(e)), x)
+    @assert all(e->abs(imag(e)) < 100*eps(abs(e)), x) "Complex probability encountered: $x"
     return Weights(real.(x))
 end
 
@@ -118,7 +118,7 @@ function generate_samples!(code::NestedEinsum, cache::CacheTree{T}, env::Abstrac
             end
 
             # recurse
-            generate_samples!(subcode, child, subenv, samples, setdiff(pool, sample_vars), size_dict)
+            generate_samples!(subcode, child, subenv, samples, pool, size_dict)
         end
     end
 end

src/utils.jl

@@ -305,3 +305,53 @@ function get_artifact_path(artifact_name::String)
     artifact_hash = Pkg.Artifacts.artifact_hash(artifact_name, artifact_toml)
     return Pkg.Artifacts.artifact_path(artifact_hash)
 end
+
+"""
+$TYPEDSIGNATURES
+
+Matrix product state (MPS) is a tensor network model that is widely used in
+quantum many-body physics. It is a special case of tensor network model where
+the tensors are rank-3 tensors and the physical indices are connected in a
+chain. The MPS is defined as:
+
+```math
+\\begin{align*}
+\\left| \\psi \\right\\rangle &= \\sum_{x_1, x_2, \\ldots, x_n} \\text{Tr}(A_1^{x_1} A_2^{x_2} \\cdots A_n^{x_n}) \\left| x_1, x_2, \\ldots, x_n \\right\\rangle \\\\
+\\left\\langle \\psi \\right| &= \\sum_{x_1, x_2, \\ldots, x_n} \\text{Tr}(A_n^{x_n} \\cdots A_2^{x_2} A_1^{x_1}) \\left\\langle x_1, x_2, \\ldots, x_n \\right|
+\\end{align*}
+```
+
+where \$A_i^{x_i}\$ is a rank-3 tensor with physical index \$x_i\$ and two virtual
+indices connecting to the next tensor. The MPS is a special case of the tensor
+network model where the tensors are rank-3 tensors and the physical indices are
+connected in a chain.
+
+### Arguments
+- `n` is the number of physical indices.
+- `chi` is the bond dimension of the virtual indices.
+- `d` is the dimension of the physical indices.
+"""
+function matrix_product_state(n::Int, chi::Int, d::Int=2)
+    tensors = Any[randn(ComplexF64, d, chi)]
+    physical_indices = collect(1:n)
+    virtual_indices_ket = collect(n+1:2n-1)
+    virtual_indices_bra = collect(2n:3n-2)
+    ixs_ket = [[physical_indices[1], virtual_indices_ket[1]]]
+    ixs_bra = [[physical_indices[1], virtual_indices_bra[1]]]
+    for i = 2:n-1
+        push!(tensors, randn(ComplexF64, chi, d, chi))
+        push!(ixs_ket, [virtual_indices_ket[i-1], physical_indices[i], virtual_indices_ket[i]])
+        push!(ixs_bra, [virtual_indices_bra[i-1], physical_indices[i], virtual_indices_bra[i]])
+    end
+    push!(tensors, randn(ComplexF64, chi, d))
+    push!(ixs_ket, [virtual_indices_ket[n-1], physical_indices[n]])
+    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),
+        optimize_code(DynamicEinCode(ixs, Int[]), OMEinsum.get_size_dict(ixs, tensors), GreedyMethod()),
+        tensors,
+        Dict{Int, Int}(),
+        Vector{Int}[[i] for i=1:n]
+    )
+end

test/sampling.jl

@@ -64,32 +64,20 @@ using OMEinsum
 end
 
 @testset "sample MPS" begin
-    tensors = [
-        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),
-        optimize_code(DynamicEinCode(ixs, Int[]), OMEinsum.get_size_dict(ixs, tensors), GreedyMethod()),
-        tensors,
-        Dict{Int, Int}(),
-        Vector{Int}[]
-    )
+    n = 4
+    chi = 3
+    mps = matrix_product_state(n, chi)
     num_samples = 1000
     samples = map(1:num_samples) do i
-        sample(mps, 1; queryvars=[1, 2, 3, 4]).samples[:, 1]
+        sample(mps, 1; queryvars=vcat(mps.mars...)).samples[:, 1]
     end
     indices = map(samples) do sample
         sum(i->sample[i] * 2^(i-1), 1:4) + 1
     end
     distribution = map(1:16) do i
         count(j->j==i, indices) / num_samples
     end
-    probs = normalize!(real.(vec(DynamicEinCode(ixs, collect(1:4))(tensors...))), 1)
+    probs = normalize!(real.(vec(DynamicEinCode(ixs, collect(1:4))(mps.tensors...))), 1)
     negative_loglikelyhood(probs, samples) = -sum(log.(probs[samples]))/length(samples)
     entropy(probs) = -sum(probs .* log.(probs))
     @test negative_loglikelyhood(probs, indices) ≈ entropy(probs) atol=1e-1