Follow updates in TensorTrains: VUMPS truncator

Project.toml

@@ -1,7 +1,7 @@
 name = "MatrixProductBP"
 uuid = "3d39929c-b583-45fa-b331-3f50b693a38a"
 authors = ["stecrotti <stefano.crotti@polito.it>"]
-version = "0.8.0"
+version = "0.9.0"
 
 [deps]
 CavityTools = "217fe2f1-7e3d-419f-9934-cd6900c2759a"
@@ -46,7 +46,7 @@ SparseArrays = "1.8"
 Statistics = "1"
 StatsBase = "0.34"
 TensorCast = "0.4"
-TensorTrains = "0.11"
+TensorTrains = "0.12"
 Tullio = "0.3"
 UnPack = "1"
 Unzip = "0.2"

src/MatrixProductBP.jl

@@ -33,11 +33,10 @@ using TensorTrains:
     bond_dims, flat_tt, rand_tt, flat_periodic_tt, rand_periodic_tt, 
     orthogonalize_right!, orthogonalize_left!, compress!,
     marginals, twovar_marginals, normalization, normalize!,
-    svd, _compose, accumulate_L, accumulate_R
-
-using TensorTrains.UniformTensorTrains:
-    InfiniteUniformTensorTrain, flat_infinite_uniform_tt
-
+    svd, _compose, accumulate_L, accumulate_R,
+    InfiniteUniformTensorTrain, flat_infinite_uniform_tt, TruncVUMPS, dot
+
+
 
 export 
     SVDTrunc, TruncBond, TruncThresh, TruncBondMax, TruncBondThresh,
@@ -67,6 +66,7 @@ export
     RecursiveBPFactor, DampedFactor, RecursiveTraceFactor, GenericFactor,
     RestrictedRecursiveBPFactor,
     mpbp_stationary, mpbp_stationary_infinite_graph, mpbp_stationary_infinite_bipartite_graph,
+    CB_BPVUMPS, TruncVUMPS,
     mean_with_uncertainty
 
 

src/stationary.jl

@@ -206,8 +206,6 @@
     end
 end
 
-default_truncator(::Type{<:InfiniteUniformMPEM2}) = TruncThresh(1e-6)
-
 function mpbp_stationary_infinite_graph(k::Integer, wᵢ::Vector{U}, qi::Int,
     ϕᵢ = fill(ones(qi), length(wᵢ));
     ψₖᵢ = fill(ones(qi, qi), length(wᵢ)),
@@ -256,3 +254,42 @@
     end
     return nothing
 end
+
+default_truncator(::Type{<:InfiniteUniformMPEM2}) = TruncVUMPS(4)
+
+struct CB_BPVUMPS{TP<:ProgressUnknown, F, M2<:InfiniteUniformMPEM2}
+    prog :: TP
+    m    :: Vector{Vector{Vector{Float64}}} 
+    Δs   :: Vector{Float64}     # convergence error on marginals
+    A    :: Vector{Vector{M2}}     
+    εs   :: Vector{Float64}     # convergence error on messages
+    f    :: F
+
+    function CB_BPVUMPS(bp::MPBPStationary{G, T, V, M2}; showprogress::Bool=true, f::F=(x,i)->x, info="") where {G, T, V, M2, F}
+        dt = showprogress ? 0.1 : Inf
+        isempty(info) || (info *= "\n")
+        prog = ProgressUnknown(desc=info*"Running MPBP: iter", dt=dt, showspeed=true)
+        TP = typeof(prog)
+        m = [means(f, bp)]
+        Δs = zeros(0)
+        A = [deepcopy(bp.μ.v)]
+        εs = zeros(0)
+        new{TP,F,M2}(prog, m, Δs, A, εs, f)
+    end
+end
+
+function (cb::CB_BPVUMPS)(bp::MPBPStationary, it::Integer, svd_trunc::SVDTrunc)
+    marg_new = means(cb.f, bp)
+    marg_old = cb.m[end]
+    Δ = isempty(marg_new) ? NaN : maximum(maximum(abs, mn .- mo) for (mn, mo) in zip(marg_new, marg_old))
+    push!(cb.Δs, Δ)
+    push!(cb.m, marg_new)
+    A_new = bp.μ
+    A_old = cb.A[end]
+    ε = isempty(A_new) ? NaN : maximum(abs, 1 - dot(Anew, Aold) for (Anew, Aold) in zip(A_new, A_old))
+    push!(cb.εs, ε)
+    push!(cb.A, deepcopy(bp.μ))
+    next!(cb.prog, showvalues=[(:Δ,Δ), (:trunc, summary_compact(svd_trunc))])
+    flush(stdout)
+    return Δ
+end

test/glauber_small_tree.jl

@@ -320,10 +320,9 @@ end
     @testset "Glauber small tree - stationary" begin
         rng = MersenneTwister(111)
 
-        J = [0 1 0 0
-            1 0 1 1
-            0 1 0 0
-            0 1 0 0] .|> float
+        J = [0 1 0
+             1 0 1
+             0 1 0] .|> float
 
         N = size(J, 1)
         h = randn(rng, N)
@@ -332,20 +331,22 @@ end
         ising = Ising(J, h, β)
         gl = Glauber(ising, 0)
         bp = mpbp_stationary(gl)
+        @test is_periodic(bp)
+        svd_trunc = TruncVUMPS(12)
 
         spin(x, i) = 3-2x; spin(x) = spin(x, 0)
-        iterate!(bp; tol=1e-14, maxiter=10)
+        iterate!(bp; tol=1e-14, maxiter=10, svd_trunc, cb=CB_BPVUMPS(bp))
         m_bp = [only(m) for m in means(spin, bp)]
 
         pb, = pair_beliefs(bp)
 
-        reset!(bp, observations=false)
-
         bp_slow = MPBP(bp.g, [GenericFactor.(w) for w in bp.w], bp.ϕ, bp.ψ, 
             deepcopy(collect(bp.μ)), deepcopy(bp.b), collect(bp.f))
-        iterate!(bp_slow; tol=1e-14, maxiter=10)
+        bp_slow.μ[1] = bp.μ[1]
+        iterate!(bp_slow; tol=1e-14, maxiter=10, svd_trunc)
         m_bp_slow = [only(m) for m in means(spin, bp_slow)]
         @test m_bp_slow ≈ m_bp
+        reset!(bp)
     end
 
 end

test/sis_small_tree.jl

@@ -148,14 +148,15 @@
     @testset "SIS small tree - stationary" begin
         sis = SIS(g, λ, ρ, 0; γ, α)
         bp = mpbp_stationary(sis)
+        svd_trunc = TruncVUMPS(10)
 
-        iterate!(bp; tol=1e-14, maxiter=10)
+        iterate!(bp; tol=1e-14, maxiter=10, svd_trunc)
         local f(x,i) = x-1
         m_bp = [only(m) for m in means(f, bp)]
 
         bp_slow = MPBP(bp.g, [GenericFactor.(w) for w in bp.w], bp.ϕ, bp.ψ, 
             deepcopy(collect(bp.μ)), deepcopy(bp.b), collect(bp.f))
-        iterate!(bp_slow; tol=1e-14, maxiter=10, damp=0.5)
+        iterate!(bp_slow; tol=1e-14, maxiter=10, damp=0.5, svd_trunc)
         m_bp_slow = [only(m) for m in means(f, bp_slow)]
         @test m_bp_slow ≈ m_bp
     end