update

Author: GiggleLiu

src/TensorInference.jl

@@ -40,6 +40,9 @@ export MMAPModel
 # for ProblemReductions
 export update_temperature
 
+# belief propagation
+export belief_propagation
+
 # utils
 export random_matrix_product_state
 

src/belief.jl

@@ -1,7 +1,56 @@
 struct BPState{T, VT<:AbstractVector{T}}
     t2v::Vector{Vector{Int}}           # a mapping from tensors to variables
     v2t::Vector{Vector{Int}}           # a mapping from variables to tensors
-    edges_vectors::Vector{Vector{VT}}  # each tensor is associated with a vector of vectors, one for each neighbor
+    tensors::Vector{AbstractArray{T}}                 # the tensors
+    message_in::Vector{Vector{VT}}  # for each variable, we store the incoming messages
+    message_out::Vector{Vector{VT}} # the outgoing messages
+end
+
+# message_in -> message_out
+function process_message!(bp::BPState)
+    for (ov, iv) in zip(bp.message_out, bp.message_in)
+        _process_message!(ov, iv)
+    end
+end
+function _process_message!(ov::Vector, iv::Vector)
+    # process the message, TODO: speed up if needed!
+    for (i, v) in enumerate(ov)
+        fill!(v, one(eltype(v)))  # clear the output vector
+        for (j, u) in enumerate(iv)
+            j != i && (v .*= u)
+        end
+    end
+end
+
+function collect_message!(bp::BPState)
+    for (it, t) in enumerate(bp.t2v)
+        _collect_message!(vectors_on_tensor(bp.message_out, bp, it), t, vectors_on_tensor(bp.message_in, bp, it))
+    end
+end
+# collect the vectors associated with the target tensor
+function vectors_on_tensor(messages, bp::BPState, it::Int)
+    return map(bp.t2v[it]) do v
+        # the message goes to the idx-th tensor from variable v
+        messages[v][findfirst(==(it), bp.v2t[v])]
+    end
+end
+function _collect_message!(vectors_out::Vector, t::AbstractArray, vectors_in::Vector)
+    @assert length(vectors_out) == length(vectors_in) == ndims(t)
+    # TODO: speed up if needed!
+    code = star_code(length(vectors_in))
+    cost, gradient = cost_and_gradient(code, [t, vectors_in...])
+    for (o, g) in zip(vectors_out, gradient[2:end])
+        o .= g
+    end
+    return cost
+end
+function star_code(n::Int)
+    ix1, ixrest = collect(1:n), [[i] for i in 1:n]
+    ne = DynamicNestedEinsum([DynamicNestedEinsum{Int}(1), DynamicNestedEinsum{Int}(2)], DynamicEinCode([ix1, ixrest[1]], collect(2:n)))
+    for i in 2:n
+        ne = DynamicNestedEinsum([ne, DynamicNestedEinsum{Int}(i + 1)], DynamicEinCode([ne.eins.iy, ixrest[i]], collect(i+1:n)))
+    end
+    return ne
 end
 
 function BPState(::Type{T}, n::Int, t2v::Vector{Vector{Int}}, size_dict::Dict{Int, Int}) where T
@@ -39,9 +88,6 @@ function belief_propagation(tn::TensorNetworkModel{T}, bpstate::BPState{T}; max_
     end
 end
 
-function tensor_product()
-end
-
 function belief_propagation(tn::TensorNetworkModel{T}) where T
     return belief_propagation(tn, BPState(T, OMEinsum.get_ixsv(tn.code), tn.size_dict))
 end

test/belief.jl

@@ -0,0 +1,36 @@
+using TensorInference, Test
+
+@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]]
+    mo = similar.(mi)
+    TensorInference._process_message!(mo, mi)
+    @test mo == mo_expected
+end
+
+@testset "star code" begin
+    code = TensorInference.star_code(3)
+    c1, c2, c3, c4 = [DynamicNestedEinsum{Int}(i) for i in 1:4]
+    ne1 = DynamicNestedEinsum([c1, c2], DynamicEinCode([[1, 2, 3], [1]], [2, 3]))
+    ne2 = DynamicNestedEinsum([ne1, c3], DynamicEinCode([[2, 3], [2]], [3]))
+    ne3 = DynamicNestedEinsum([ne2, c4], DynamicEinCode([[3], [3]], Int[]))
+    @test code == ne3
+    t = randn(2, 2, 2)
+    v1 = randn(2)
+    v2 = randn(2)
+    v3 = randn(2)
+    vectors_out = [similar(v1), similar(v2), similar(v3)]
+    TensorInference._collect_message!(vectors_out, t, [v1, v2, v3])
+    @test vectors_out[1] ≈ reshape(t, 2, 4) * kron(v3, v2)  # NOTE: v3 is the little end
+    @test vectors_out[2] ≈ vec(v1' * reshape(reshape(t, 4, 2) * v3, 2, 2))
+    @test vectors_out[3] ≈ vec(kron(v2, v1)' * reshape(t, 4, 2))
+end
+
+@testset "belief propagation" begin
+    n = 5
+    chi = 3
+    Random.seed!(140)
+    mps = random_matrix_product_state(n, chi)
+    model = TensorNetworkModel(mps)
+    state = belief_propagation(model)
+end