MPBP with infinite messages to describe the stationary state of processes

Project.toml

@@ -1,7 +1,7 @@
 name = "MatrixProductBP"
 uuid = "3d39929c-b583-45fa-b331-3f50b693a38a"
 authors = ["stecrotti <[email protected]>"]
-version = "0.7.0"
+version = "0.8.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.10, 0.11"
+TensorTrains = "0.11"
 Tullio = "0.3"
 UnPack = "1"
 Unzip = "0.2"

src/MatrixProductBP.jl

@@ -35,11 +35,16 @@ using TensorTrains:
     marginals, twovar_marginals, normalization, normalize!,
     svd, _compose, accumulate_L, accumulate_R
 
+using TensorTrains.UniformTensorTrains:
+    InfiniteUniformTensorTrain, flat_infinite_uniform_tt
+
 
 export 
     SVDTrunc, TruncBond, TruncThresh, TruncBondMax, TruncBondThresh,
-    PeriodicMPEM2, PeriodicMPEM3, PeriodicMPEM1,
-    MPEM1, MPEM2, MPEM3, mpem2, rand_mpem1, rand_mpem2, normalization, normalize!, marginalize,
+    PeriodicMPEM2, PeriodicMPEM3, PeriodicMPEM1, getT,
+    MPEM1, MPEM2, MPEM3, mpem2, rand_mpem1, rand_mpem2, 
+    InfiniteUniformMPEM1, InfiniteUniformMPEM2, InfiniteUniformMPEM3,
+    normalization, normalize!, marginalize,
     orthogonalize_right!, orthogonalize_left!, compress!, twovar_marginals, evaluate,
     BPFactor, nstates, MPBP, mpbp, reset_messages!, reset_beliefs!, reset_observations!,
     reset!, is_free_dynamics, onebpiter!, CB_BP, iterate!, 
@@ -61,6 +66,7 @@ export
     draw_node_observations!, AtomicVector,
     RecursiveBPFactor, DampedFactor, RecursiveTraceFactor, GenericFactor,
     RestrictedRecursiveBPFactor,
+    mpbp_stationary, mpbp_stationary_infinite_graph, mpbp_stationary_infinite_bipartite_graph,
     mean_with_uncertainty
 
 
@@ -74,6 +80,7 @@ include("test_factors.jl")
 include("infinite_graph.jl")
 include("exact.jl")
 include("sampling.jl")
+include("stationary.jl")
 
 include("Models/Models.jl")
 

src/Models/Models.jl

@@ -5,7 +5,7 @@ import MatrixProductBP: exact_prob, getT, nstates, mpbp, compress!,
     beliefs, beliefs_tu, marginals, pair_belief, pair_beliefs,
     marginalize, cavity, onebpiter!, check_ψs, _compose,
     RecursiveBPFactor, nstates, prob_y, prob_xy, prob_yy, prob_y0, prob_y_partial,
-    prob_y_dummy, periodic_mpbp
+    prob_y_dummy, periodic_mpbp, mpbp_stationary
 using MatrixProductBP
 
 import IndexedGraphs: IndexedGraph, IndexedBiDiGraph, AbstractIndexedDiGraph, ne, nv, 
@@ -26,7 +26,7 @@ import Lazy: @forward
 
 export 
     Ising, Glauber, energy,
-    HomogeneousGlauberFactor, GenericGlauberFactor, PMJGlauberFactor, mpbp,
+    HomogeneousGlauberFactor, GenericGlauberFactor, PMJGlauberFactor, mpbp, mpbp_stationary,
     equilibrium_magnetization, equilibrium_observables, RandomRegular, ErdosRenyi, CB_Pop,
     SIS, SISFactor, SIRS, SIRSFactor, SIS_heterogeneous, SIS_heterogeneousFactor, SUSCEPTIBLE, INFECTIOUS, RECOVERED,
     kl_marginals, l1_marginals, roc, auc

src/Models/epidemics/sis_bp.jl

@@ -51,6 +51,12 @@ function periodic_mpbp(sis::SIS{T,N,F}; kw...) where {T,N,F}
     return periodic_mpbp(g, w, fill(2, nv(g)), T, ϕ=sis.ϕ, ψ=sis.ψ; kw...)
 end
 
+function mpbp_stationary(sis::SIS{T,N,F}; kw...) where {T,N,F<:AbstractFloat}
+    g = IndexedBiDiGraph(sis.g.A)
+    w = sis_factors(sis)
+    return mpbp_stationary(g, w, fill(2, nv(g)); ϕ=sis.ϕ, ψ=sis.ψ, kw...)
+end
+
 # neighbor j is susceptible -> does nothing
 function prob_y(wᵢ::SISFactor, xᵢᵗ⁺¹, xᵢᵗ, yᵗ, d)
     @unpack λ, ρ, α = wᵢ

src/Models/glauber/glauber_bp.jl

@@ -107,6 +107,14 @@ function periodic_mpbp(gl::Glauber{T,N,F}; kw...) where {T,N,F<:AbstractFloat}
     return periodic_mpbp(g, w, fill(2, nv(g)), T; ϕ, ψ, kw...)
 end
 
+function mpbp_stationary(gl::Glauber{T,N,F}; kw...) where {T,N,F<:AbstractFloat}
+    g = IndexedBiDiGraph(gl.ising.g.A)
+    w = glauber_factors(gl.ising, T)
+    ϕ = gl.ϕ
+    ψ = pair_obs_undirected_to_directed(gl.ψ, gl.ising.g)
+    return mpbp_stationary(g, w, fill(2, nv(g)); ϕ, ψ, kw...)
+end
+
 
 # construct an array of GlauberFactors corresponding to gl
 # seems to be type stable

src/bp_core.jl

@@ -2,7 +2,7 @@
 Factor for the factor graph of a model solvable with MPBP.
 
 Any `BPFactor` subtype must implement:
-- A functor that computes the Boltzmann contribution to the joint probability
+- A functor that computes the factor contribution to the joint probability
 
 That's it!
 
@@ -17,7 +17,7 @@ Base.broadcastable(b::BPFactor) = Ref(b)
 # ψᵢⱼ are the ones living on the outedges of node i
 function f_bp(A::Vector{M2}, wᵢ::Vector{U}, ϕᵢ::Vector{Vector{F}}, 
         ψₙᵢ::Vector{Vector{Matrix{F}}}, j_index::Integer; showprogress=false, 
-        svd_trunc::SVDTrunc=TruncThresh(0.0), periodic=false) where {F,U<:BPFactor,M2<:AbstractMPEM2}
+        svd_trunc::SVDTrunc=default_truncator(M2), periodic=false) where {F,U<:BPFactor,M2<:AbstractMPEM2}
     T = length(A[1]) - 1
     @assert all(length(a) == T + 1 for a in A)
     @assert length(wᵢ) == T + 1
@@ -57,9 +57,9 @@ function f_bp(A::Vector{M2}, wᵢ::Vector{U}, ϕᵢ::Vector{Vector{F}},
 end
 
 # compute outgoing message to dummy neighbor to get the belief
-function f_bp_dummy_neighbor(A::Vector{<:AbstractMPEM2}, 
+function f_bp_dummy_neighbor(A::Vector{M2}, 
         wᵢ::Vector{U}, ϕᵢ::Vector{Vector{F}}, ψₙᵢ::Vector{Vector{Matrix{F}}};
-        showprogress=false, svd_trunc::SVDTrunc=TruncThresh(0.0), periodic=false) where {F,U<:BPFactor}
+        showprogress=false, svd_trunc::SVDTrunc=default_truncator(M2), periodic=false) where {F,U<:BPFactor,M2<:AbstractMPEM2}
     q = length(ϕᵢ[1])
     T = length(ϕᵢ) - 1
     @assert all(length(a) == T + 1 for a in A)

src/mpbp.jl

@@ -114,43 +114,43 @@
 is_periodic(bp::MPBP{G,F,V,<:PeriodicMPEM2,<:PeriodicMPEM1}) where {G,F,V} = true
 
 # compute outgoing messages from node `i`
-function onebpiter!(bp::MPBP, i::Integer, ::Type{U}; 
-        svd_trunc::SVDTrunc=TruncThresh(1e-6), damp=0.0) where {U<:BPFactor}
+function onebpiter!(bp::MPBP{G,F,V,MsgType}, i::Integer, ::Type{U}; 
+        svd_trunc::SVDTrunc=default_truncator(MsgType), damp=0.0) where {U<:BPFactor,G,F,V,MsgType}
     @unpack g, w, ϕ, ψ, μ = bp
     ein = inedges(g,i)
     eout = outedges(g, i)
     A = μ[ein.|>idx]
-    @assert all(float(normalization(a)) ≈ 1 for a in A)
+    @debug @assert all(float(normalization(a)) ≈ 1 for a in A)
     sumlogzᵢ₂ⱼ = 0.0
     for (j_ind, e_out) in enumerate(eout)
         B, logzᵢ₂ⱼ = f_bp(A, w[i], ϕ[i], ψ[eout.|>idx], j_ind; svd_trunc, periodic=is_periodic(bp))
         sumlogzᵢ₂ⱼ += logzᵢ₂ⱼ
         C = mpem2(B)
-        μj = orthogonalize_right!(C; svd_trunc)
+        μj = compress!(C; svd_trunc, is_orthogonal=:left)
         sumlogzᵢ₂ⱼ += normalize!(μj)
         μ[idx(e_out)] = μj
     end
     dᵢ = length(ein)
     bp.b[i] = onebpiter_dummy_neighbor(bp, i; svd_trunc) |> marginalize
-    logzᵢ = log(normalization(bp.b[i]))
+    logzᵢ = real(log(normalization(bp.b[i])))
     bp.f[i] = (dᵢ/2-1)*logzᵢ - (1/2)*sumlogzᵢ₂ⱼ
     nothing
 end
 
-function onebpiter!(bp::MPBP, i::Integer; svd_trunc::SVDTrunc=TruncThresh(1e-6))
+function onebpiter!(bp::MPBP{G,F,V,MsgType}, i::Integer; svd_trunc::SVDTrunc=default_truncator(MsgType)) where {G,F,V,MsgType}
     onebpiter!(bp, i, eltype(bp.w[i]); svd_trunc, damp=0.0)
 end
 
 
-function onebpiter_dummy_neighbor(bp::MPBP, i::Integer; 
-        svd_trunc::SVDTrunc=TruncThresh(1e-6))
+function onebpiter_dummy_neighbor(bp::MPBP{G,F,V,MsgType}, i::Integer; 
+        svd_trunc::SVDTrunc=default_truncator(MsgType)) where {G,F,V,MsgType}
     @unpack g, w, ϕ, ψ, μ = bp
     ein = inedges(g,i)
     eout = outedges(g, i)
     A = μ[ein.|>idx]
     B, _ = f_bp_dummy_neighbor(A, w[i], ϕ[i], ψ[eout.|>idx]; svd_trunc, periodic=is_periodic(bp))
     C = mpem2(B)
-    return orthogonalize_right!(C; svd_trunc)
+    return compress!(C; svd_trunc, is_orthogonal=:left)
 end
 
 # A callback to print info and save stuff during the iterations 
@@ -177,15 +177,15 @@
     Δ = 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)
-    next!(cb.prog, showvalues=[(:Δ,Δ), summary_compact(svd_trunc)])
+    next!(cb.prog, showvalues=[(:Δ,Δ), (:trunc, summary_compact(svd_trunc))])
     flush(stdout)
     return Δ
 end
 
-function iterate!(bp::MPBP; maxiter::Integer=5, 
-        svd_trunc::SVDTrunc=TruncThresh(1e-6),
+function iterate!(bp::MPBP{G,F,V,MsgType}; maxiter::Integer=5, 
+        svd_trunc::SVDTrunc=default_truncator(MsgType),
         showprogress=true, cb=CB_BP(bp; showprogress), tol=1e-10, 
-        nodes = collect(vertices(bp.g)), shuffle_nodes::Bool=true, damp=0.0)
+        nodes = collect(vertices(bp.g)), shuffle_nodes::Bool=true, damp=0.0) where {G,F,V,MsgType}
     for it in 1:maxiter
         Threads.@threads for i in nodes
             onebpiter!(bp, i, eltype(bp.w[i]); svd_trunc, damp)
@@ -254,7 +254,7 @@
 
 autocorrelations(bp::MPBP; kw...) = autocorrelations((x,i)->x, bp; kw...)
 
-function means(f, bp::MPBP; sites=vertices(bp.g))
+function means(f, bp::MPBP{G,F,V,M2}; sites=vertices(bp.g)) where {G,F,V,M2}
     map(sites) do i
         expectation.(x->f(x, i), marginals(bp.b[i]))
     end

src/mpems.jl

@@ -6,7 +6,7 @@ const PeriodicMPEM1{F} = PeriodicTensorTrain{F, 3}
 flat_mpem1(q::Int, T::Int; d::Int=2, bondsizes=[1; fill(d, T); 1]) = flat_tt(bondsizes, q)
 flat_periodic_mpem1(q::Int, T::Int; d::Int=2, bondsizes=fill(d, T+1)) = flat_periodic_tt(bondsizes, q)
 
-# construct a flat mpem with given bond dimensions
+# construct a random mpem with given bond dimensions
 rand_mpem1(q::Int, T::Int; d::Int=2, bondsizes=[1; fill(d, T); 1]) = rand_tt(bondsizes, q)
 rand_periodic_mpem1(q::Int, T::Int; d::Int=2, bondsizes=fill(d, T+1)) = rand_periodic_tt(bondsizes, q)
 
@@ -41,7 +41,9 @@ struct MPEM3{F<:Real}
             throw(ArgumentError("First matrix must have 1 row, last matrix must have 1 column"))
         check_bond_dims(tensors) ||
             throw(ArgumentError("Matrix indices for matrix product non compatible"))
-        new{F}(tensors, z)
+        all(size(At, 3) == size(At, 5) for At in tensors) ||
+            throw(ArgumentError("First and third variable indices should have matching ranges because they both represent `xᵢ`"))
+        return new{F}(tensors, z)
     end
 end
 
@@ -99,7 +101,9 @@ struct PeriodicMPEM3{F<:Real}
             throw(ArgumentError("Number of rows of the first matrix should coincide with the number of columns of the last matrix"))
         check_bond_dims(tensors) ||
             throw(ArgumentError("Matrix indices for matrix product non compatible"))
-        new{F}(tensors, z)
+        all(size(At, 3) == size(At, 5) for At in tensors) ||
+            throw(ArgumentError("First and third variable indices should have matching ranges because they both represent `xᵢ`"))
+        return new{F}(tensors, z)
     end
 end
 
@@ -152,4 +156,6 @@ end
 
 
 mpem3from2(::Type{MPEM2{F}}) where F = MPEM3
-mpem3from2(::Type{PeriodicMPEM2{F}}) where F = PeriodicMPEM3
+mpem3from2(::Type{PeriodicMPEM2{F}}) where F = PeriodicMPEM3
+
+default_truncator(::Type{<:AbstractMPEM2}) = TruncThresh(1e-6)

src/recursive_bp_factor.jl

@@ -143,8 +143,8 @@ function compute_prob_ys(wᵢ::Vector{U}, qi::Int, μin::Vector{M2}, ψout, T, s
 end
 
 # compute outgoing messages from node `i`
-function onebpiter!(bp::MPBP{G,F}, i::Integer, ::Type{U}; 
-        svd_trunc::SVDTrunc=TruncThresh(1e-6), damp::Real=0.0) where {G<:AbstractIndexedDiGraph,F<:Real,U<:RecursiveBPFactor}
+function onebpiter!(bp::MPBP{G,F,V,MsgType}, i::Integer, ::Type{U}; 
+    svd_trunc::SVDTrunc=default_truncator(MsgType), damp::Real=0.0) where {G<:AbstractIndexedDiGraph,F<:Real,U<:RecursiveBPFactor,V,MsgType}
     @unpack g, w, ϕ, ψ, μ = bp
     ein, eout = inedges(g,i), outedges(g, i)
     wᵢ, ϕᵢ, dᵢ  = w[i], ϕ[i], length(ein)
@@ -153,7 +153,8 @@ function onebpiter!(bp::MPBP{G,F}, i::Integer, ::Type{U};
     sumlogzᵢ₂ⱼ = zero(F)
     for (j,e) = enumerate(eout)
         B = f_bp_partial_ij(C[j], wᵢ, ϕᵢ, dᵢ - 1, nstates(bp, dst(e)), j)
-        μj = orthogonalize_right!(mpem2(B); svd_trunc)
+        μj = compress!(mpem2(B); svd_trunc, is_orthogonal=:left)
+        normalize_eachmatrix!(μj)
         sumlogzᵢ₂ⱼ += set_msg!(bp, μj, idx(e), damp, svd_trunc)
     end
     B = f_bp_partial_i(full, wᵢ, ϕᵢ, dᵢ)

src/sampling.jl

@@ -177,7 +177,8 @@ function autocorrelations(f, sms::SoftMarginSampler; showprogress::Bool=true,
 end
 
 function autocovariances(f, sms::SoftMarginSampler; showprogress::Bool=true,
-        sites=vertices(sms.bp.g), r = autocorrelations(f, sms; showprogress, sites), 
+        sites=vertices(sms.bp.g), maxdist = getT(sms.bp), 
+        r = autocorrelations(f, sms; showprogress, sites, maxdist), 
         μ = means(f, sms; sites))
     @assert length(μ) == size(r, 1)
     return covariance.(r, μ)