update

Author: GiggleLiu

src/TensorInference.jl

@@ -36,7 +36,7 @@ export sample
 export MMAPModel
 
 # utils
-export matrix_product_state
+export random_matrix_product_state
 
 include("Core.jl")
 include("RescaledArray.jl")

src/sampling.jl

@@ -50,6 +50,10 @@ Returns a vector of vector, each element being a configurations defined on `get_
 ### Arguments
 * `tn` is the tensor network model.
 * `n` is the number of samples to be returned.
+
+### Keyword Arguments
+* `usecuda` is a boolean flag to indicate whether to use CUDA for tensor computation.
+* `queryvars` is the variables to be sampled, default is `get_vars(tn)`.
 """
 function sample(tn::TensorNetworkModel, n::Int; usecuda = false, queryvars = get_vars(tn))::Samples
     # generate tropical tensors with its elements being log(p).
@@ -80,7 +84,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) "Complex probability encountered: $x"
+    @assert all(e->abs(imag(e)) < max(100*eps(abs(e)), 1e-10), x) "Complex probability encountered: $x"
     return Weights(real.(x))
 end
 

src/utils.jl

@@ -331,19 +331,22 @@ connected in a chain.
 - `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)]
+random_matrix_product_state(n::Int, chi::Int, d::Int=2) = random_matrix_product_state(ComplexF64, n, chi, d)
+function random_matrix_product_state(::Type{T}, n::Int, chi::Int, d::Int=2) where T
+    # chi ^ (n-1) * (variance^n)^2 == 1/d^n
+    variance = d^(-1/2) * chi^(-1/2+1/2n)
+    tensors = Any[randn(T, d, chi) .* variance]
     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!(tensors, randn(T, chi, d, chi) .* variance)
         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!(tensors, randn(T, chi, d) .* variance)
     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...]

test/sampling.jl

@@ -66,11 +66,9 @@ end
 @testset "sample MPS" begin
     n = 4
     chi = 3
-    mps = matrix_product_state(n, chi)
+    mps = random_matrix_product_state(n, chi)
     num_samples = 1000
-    samples = map(1:num_samples) do i
-        sample(mps, 1; queryvars=vcat(mps.mars...)).samples[:, 1]
-    end
+    sample(mps, num_samples; queryvars=vcat(mps.mars...))
     indices = map(samples) do sample
         sum(i->sample[i] * 2^(i-1), 1:4) + 1
     end