update

Author: GiggleLiu

src/belief.jl

@@ -43,24 +43,29 @@ struct BPState{T, VT<:AbstractVector{T}}
 end
 
 # message_in -> message_out
-function process_message!(bp::BPState)
+function process_message!(bp::BPState; normalize, damping)
     for (ov, iv) in zip(bp.message_out, bp.message_in)
-        _process_message!(ov, iv)
+        _process_message!(ov, iv, normalize, damping)
     end
 end
-function _process_message!(ov::Vector, iv::Vector)
+function _process_message!(ov::Vector, iv::Vector, normalize::Bool, damping)
     # process the message, TODO: speed up if needed!
     for (i, v) in enumerate(ov)
-        fill!(v, one(eltype(v)))  # clear the output vector
+        w = similar(v)
+        fill!(w, one(eltype(v)))  # clear the output vector
         for (j, u) in enumerate(iv)
-            j != i && (v .*= u)
+            j != i && (w .*= u)
         end
+        normalize && normalize!(w, 1)
+        v .= v .* damping + (1 - damping) * w
     end
 end
 
-function collect_message!(bp::BeliefPropgation, state::BPState)
+function collect_message!(bp::BeliefPropgation, state::BPState; normalize::Bool)
     for it in 1:num_tensors(bp)
-        _collect_message!(vectors_on_tensor(state.message_in, bp, it), bp.tensors[it], vectors_on_tensor(state.message_out, bp, it))
+        out = vectors_on_tensor(state.message_in, bp, it)
+        _collect_message!(out, bp.tensors[it], vectors_on_tensor(state.message_out, bp, it))
+        normalize && normalize!.(out, 1)
     end
 end
 # collect the vectors associated with the target tensor
@@ -78,7 +83,7 @@ function _collect_message!(vectors_out::Vector, t::AbstractArray, vectors_in::Ve
     for (o, g) in zip(vectors_out, gradient[2:end])
         o .= g
     end
-    return cost
+    return cost[]
 end
 
 # star code: contract a tensor with multiple vectors, one for each dimension
@@ -112,20 +117,20 @@ Run the belief propagation algorithm, and return the final state and the informa
 - `max_iter::Int=100`: the maximum number of iterations
 - `tol::Float64=1e-6`: the tolerance for the convergence
 """
-function belief_propagate(bp::BeliefPropgation; max_iter::Int=100, tol::Float64=1e-6)
+function belief_propagate(bp::BeliefPropgation; kwargs...)
     state = initial_state(bp)
-    info = belief_propagate!(bp, state; max_iter=max_iter, tol=tol)
+    info = belief_propagate!(bp, state; kwargs...)
     return state, info
 end
 struct BPInfo
     converged::Bool
     iterations::Int
 end
-function belief_propagate!(bp::BeliefPropgation, state::BPState{T}; max_iter::Int=100, tol::Float64=1e-6) where T
+function belief_propagate!(bp::BeliefPropgation, state::BPState{T}; max_iter::Int=100, tol=1e-6, damping=0.2) where T
     pre_message_in = deepcopy(state.message_in)
     for i in 1:max_iter
-        collect_message!(bp, state)
-        process_message!(state)
+        collect_message!(bp, state; normalize=true)
+        process_message!(state; normalize=true, damping=damping)
         # check convergence
         if all(iv -> all(it -> isapprox(state.message_in[iv][it], pre_message_in[iv][it], atol=tol), 1:length(bp.v2t[iv])), 1:num_variables(bp))
             return BPInfo(true, i)

src/utils.jl

@@ -380,21 +380,21 @@ Tensor train (TT) is a tensor network model that is widely used in quantum
 many-body physics. This model is different from the matrix product state (MPS)
 in that it does not have an extra copy for representing the bra state.
 """
-function random_tensor_train_uai(::Type{T}, n::Int, chi::Int, d::Int=2) where T
+function random_tensor_train_uai(::Type{T}, n::Int, chi::Int, d::Int=2; periodic=false) 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 = collect(n+1:2n-1)
-    ixs = [[physical_indices[1], virtual_indices[1]]]
+    virtual_indices = collect(n+1:2n)
+    tensors = Any[(periodic ? rand(T, chi, d, chi) : rand(T, d, chi)) .* variance]
+    ixs = [periodic ? [virtual_indices[n], physical_indices[1], virtual_indices[1]] : [physical_indices[1], virtual_indices[1]]]
     for i = 2:n-1
-        push!(tensors, randn(T, chi, d, chi) .* variance)
+        push!(tensors, rand(T, chi, d, chi) .* variance)
         push!(ixs, [virtual_indices[i-1], physical_indices[i], virtual_indices[i]])
     end
-    push!(tensors, randn(T, chi, d) .* variance)
-    push!(ixs, [virtual_indices[n-1], physical_indices[n]])
+    push!(tensors, (periodic ? rand(T, chi, d, chi) : rand(T, chi, d)) .* variance)
+    push!(ixs, periodic ? [virtual_indices[n-1], physical_indices[n], virtual_indices[n]] : [virtual_indices[n-1], physical_indices[n]])
     size_dict = OMEinsum.get_size_dict(ixs, tensors)
-    nvars = 2n-1
+    nvars = periodic ? 2n : 2n-1
     return UAIModel(
         nvars,
         [size_dict[i] for i=1:nvars],

test/belief.jl

@@ -1,12 +1,15 @@
 using TensorInference, Test
-using OMEinsum
+using OMEinsum, LinearAlgebra
 
 @testset "process message" begin
-    mi = [[1, 2, 3], [2, 3, 4], [3, 4, 5]]
-    mo_expected = [[6, 12, 20], [3, 8, 15], [2, 6, 12]]
+    mi = [[1.0, 2, 3], [2.0, 3, 4], [3.0, 4, 5]]
+    mo_expected = [[6.0, 12, 20], [3.0, 8, 15], [2.0, 6, 12]]
     mo = similar.(mi)
-    TensorInference._process_message!(mo, mi)
-    @test mo == mo_expected
+    TensorInference._process_message!(mo, mi, false, 0)
+    @test all(mo .≈ mo_expected)
+
+    TensorInference._process_message!(mo, mi, true, 0)
+    @test all(mo .≈ normalize!.(mo_expected, 1))
 end
 
 @testset "star code" begin
@@ -44,14 +47,29 @@ end
     @test TensorInference.initial_state(bp) isa TensorInference.BPState
     state, info = belief_propagate(bp)
     @test info.converged
-    @test info.iterations < 10
+    @test info.iterations < 20
+    mars = marginals(state)
+    tnet = TensorNetworkModel(mps_uai)
+    mars_tnet = marginals(tnet)
+    for v in 1:TensorInference.num_variables(bp)
+        @test mars[[v]] ≈ mars_tnet[[v]] atol=1e-6
+    end
+end
+
+@testset "belief propagation on circle" begin
+    n = 10
+    chi = 3
+    mps_uai = TensorInference.random_tensor_train_uai(Float64, n, chi; periodic=true)
+    bp = BeliefPropgation(mps_uai)
+    @test TensorInference.initial_state(bp) isa TensorInference.BPState
+    state, info = belief_propagate(bp; max_iter=100, tol=1e-6)
+    @test info.converged
+    @test info.iterations < 100
     contraction_res = TensorInference.contraction_results(state)
     tnet = TensorNetworkModel(mps_uai)
-    expected_result = probability(tnet)[]
-    @test all(r -> isapprox(r, expected_result), contraction_res)
     mars = marginals(state)
     mars_tnet = marginals(tnet)
     for v in 1:TensorInference.num_variables(bp)
-        @test mars[[v]] ≈ mars_tnet[[v]]
+        @test mars[[v]] ≈ mars_tnet[[v]] atol=1e-4
     end
 end

test/utils.jl

@@ -3,6 +3,9 @@ using TensorInference, Test
 @testset "tensor train" begin
     tt = random_tensor_train_uai(Float64, 5, 3)
     @test tt.nvars == length(unique(vcat([collect(Int, f.vars) for f in tt.factors]...)))
+
+    tt = random_tensor_train_uai(Float64, 5, 3; periodic=true)
+    @test tt.nvars == length(unique(vcat([collect(Int, f.vars) for f in tt.factors]...)))
 end
 
 @testset "mps" begin