Convergence criterion

Author: gdalle

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.9.2"
 manifest_format = "2.0"
-project_hash = "eabebeb24d3db0211e95e73a9e2e965bf16ba4e1"
+project_hash = "dd3777d035d7f906717e2a66a2e6f5c3854a6b1a"
 
 [[deps.ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

docs/Project.toml

@@ -1,3 +1,8 @@
 [deps]
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StochasticBlockModelVariants = "37adabcb-1964-4c56-850d-39a4302c0c39"

docs/plots.jl

@@ -0,0 +1,8 @@
+using Base.Threads
+using CairoMakie
+using LinearAlgebra
+using Random: default_rng
+using StochasticBlockModelVariants
+using ProgressMeter
+
+BLAS.set_num_threads(1)

src/StochasticBlockModelVariants.jl

@@ -1,28 +1,27 @@
 module StochasticBlockModelVariants
 
 using Graphs: AbstractGraph, neighbors
-using LinearAlgebra: dot, mul!
+using LinearAlgebra: dot, mul!, norm
 using PrecompileTools: @compile_workload
 using ProgressMeter: Progress, next!
 using Random: AbstractRNG, default_rng
 using SimpleWeightedGraphs: SimpleWeightedGraph
-using Statistics: mean
+using Statistics: mean, std
 using SparseArrays: SparseMatrixCSC, sparse, findnz
 
 export ContextualSBM, ContextualSBMLatents, ContextualSBMObservations
 export affinities, effective_snr
-export init_amp, update_amp!, run_amp
-export overlap
+export init_amp, update_amp!, run_amp, overlaps
 
 include("utils.jl")
 include("csbm.jl")
 include("csbm_inference.jl")
 
 @compile_workload begin
     rng = default_rng()
-    csbm = ContextualSBM(; N=10^2, P=10^2, d=5, λ=2, μ=2, ρ=0.0)
+    csbm = ContextualSBM(; N=10^2, P=10^2, d=5, λ=2, μ=2, ρ=0.1)
     (; observations) = rand(rng, csbm)
-    run_amp(rng; observations, csbm, iterations=2)
+    run_amp(rng; observations, csbm, max_iterations=2)
 end
 
 end

src/csbm_inference.jl

@@ -1,33 +1,50 @@
-## Storage
+## Marginals
 
 """
-    AMPStorage
+    AMPMarginals
 
 # Fields
 
 - `û::Vector`: posterior mean of `u`, length `N`
 - `v̂::Vector`: posterior mean of `v`, length `P`
 - `χ₊e::Dict`: messages about the marginal distribution of `u`, size `(N, N)`
 """
-@kwdef struct AMPStorage{R<:Real}
+@kwdef struct AMPMarginals{R<:Real}
     û::Vector{R}
     v̂::Vector{R}
     χ₊e::Dict{Tuple{Int,Int},R}
 end
 
-function Base.copy(storage::AMPStorage)
-    return AMPStorage(; û=copy(storage.û), v̂=copy(storage.v̂), χ₊e=copy(storage.χ₊e))
+function Base.copy(marginals::AMPMarginals)
+    return AMPMarginals(;
+        û=copy(marginals.û), v̂=copy(marginals.v̂), χ₊e=copy(marginals.χ₊e)
+    )
+end
+
+function Base.copy!(marginals_dest::AMPMarginals, marginals_source::AMPMarginals)
+    marginals_dest.û .= marginals_source.û
+    marginals_dest.v̂ .= marginals_source.v̂
+    copy!(marginals_dest.χ₊e, marginals_source.χ₊e)
+    return marginals_dest
 end
 
-function Base.copy!(storage_dest::AMPStorage, storage_source::AMPStorage)
-    storage_dest.û .= storage_source.û
-    storage_dest.v̂ .= storage_source.v̂
-    copy!(storage_dest.χ₊e, storage_source.χ₊e)
-    return storage_dest
+function overlaps(;
+    u::Vector{<:Integer}, v::Vector{R}, û::Vector{R}, v̂::Vector{R}
+) where {R}
+    û .= sign.(û)
+    û[abs.(û) .< eps(R)] .= one(R)
+
+    q̂ᵤ = max(freq_equalities(û, u), freq_equalities(û, -u))
+    qᵤ = 2 * (q̂ᵤ - one(R) / 2)
+
+    q̂ᵥ = max(abs(dot(v̂, v)), abs(dot(v̂, -v)))
+    qᵥ = q̂ᵥ / (eps(R) + norm(v̂) * norm(v))
+
+    return (; qᵤ, qᵥ)
 end
 
 """
-    AMPTempStorage
+    AMPStorage
 
 # Fields
 
@@ -37,21 +54,21 @@ end
 - `h̃₋::Vector`: individual external field for `u=-1`, length `N`
 - `χ₊::Vector`: marginal probability of `u=1`, length `N`
 """
-@kwdef struct AMPTempStorage{R<:Real}
+@kwdef struct AMPStorage{R<:Real}
     û_no_feat::Vector{R}
     v̂_no_comm::Vector{R}
     h̃₊::Vector{R}
     h̃₋::Vector{R}
     χ₊::Vector{R}
 end
 
-function Base.copy(temp_storage::AMPTempStorage)
-    return AMPTempStorage(;
-        û_no_feat=copy(temp_storage.û_no_feat),
-        v̂_no_comm=copy(temp_storage.v̂_no_comm),
-        h̃₊=copy(temp_storage.h̃₊),
-        h̃₋=copy(temp_storage.h̃₋),
-        χ₊=copy(temp_storage.χ₊),
+function Base.copy(storage::AMPStorage)
+    return AMPStorage(;
+        û_no_feat=copy(storage.û_no_feat),
+        v̂_no_comm=copy(storage.v̂_no_comm),
+        h̃₊=copy(storage.h̃₊),
+        h̃₋=copy(storage.h̃₋),
+        χ₊=copy(storage.χ₊),
     )
 end
 
@@ -66,42 +83,42 @@ function init_amp(
     (; N, P) = csbm
     (; g, Ξ) = observations
 
-    û = prior₊.(R, Ξ) + init_std .* randn(rng, R, N)
-    v̂ = 2 .* prior₊.(R, Ξ) .- one(R) .+ init_std .* randn(rng, R, P)
+    û = 2 .* prior₊.(R, Ξ) .- one(R) .+ init_std .* randn(rng, R, N)
+    v̂ = init_std .* randn(rng, R, P)
     χ₊e = Dict{Tuple{Int,Int},R}()
     for i in 1:N, j in neighbors(g, i)
-        χ₊e[i, j] = (one(R) / 2) + init_std * randn(rng, R)
+        χ₊e[i, j] = prior₊(R, Ξ[i]) + init_std * randn(rng, R)
     end
-    storage = AMPStorage(; û, v̂, χ₊e)
-    next_storage = copy(storage)
+    marginals = AMPMarginals(; û, v̂, χ₊e)
+    next_marginals = copy(marginals)
 
     û_no_feat = zeros(R, N)
     v̂_no_comm = zeros(R, P)
     h̃₊ = zeros(R, N)
     h̃₋ = zeros(R, N)
     χ₊ = zeros(R, N)
-    temp_storage = AMPTempStorage(; û_no_feat, v̂_no_comm, h̃₊, h̃₋, χ₊)
+    storage = AMPStorage(; û_no_feat, v̂_no_comm, h̃₊, h̃₋, χ₊)
 
-    return (; storage, next_storage, temp_storage)
+    return (; marginals, next_marginals, storage)
 end
 
 prior₊(::Type{R}, Ξᵢ) where {R} = Ξᵢ == 0 ? one(R) / 2 : R(Ξᵢ == 1)
 prior₋(::Type{R}, Ξᵢ) where {R} = Ξᵢ == 0 ? one(R) / 2 : R(Ξᵢ == -1)
 
 function update_amp!(
-    next_storage::AMPStorage{R},
-    temp_storage::AMPTempStorage{R};
-    storage::AMPStorage{R},
+    next_marginals::AMPMarginals{R},
+    storage::AMPStorage{R};
+    marginals::AMPMarginals{R},
     observations::ContextualSBMObservations{R},
     csbm::ContextualSBM{R},
 ) where {R}
     (; d, λ, μ, N, P) = csbm
     (; g, B, Ξ) = observations
     (; cᵢ, cₒ) = affinities(csbm)
 
-    ûᵗ, v̂ᵗ, χ₊eᵗ = storage.û, storage.v̂, storage.χ₊e
-    ûᵗ⁺¹, v̂ᵗ⁺¹, χ₊eᵗ⁺¹ = next_storage.û, next_storage.v̂, next_storage.χ₊e
-    (; û_no_feat, v̂_no_comm, h̃₊, h̃₋, χ₊) = temp_storage
+    ûᵗ, v̂ᵗ, χ₊eᵗ = marginals.û, marginals.v̂, marginals.χ₊e
+    ûᵗ⁺¹, v̂ᵗ⁺¹, χ₊eᵗ⁺¹ = next_marginals.û, next_marginals.v̂, next_marginals.χ₊e
+    (; û_no_feat, v̂_no_comm, h̃₊, h̃₋, χ₊) = storage
 
     ûₜ_sum = sum(ûᵗ)
     ûₜ_sum2 = sum(abs2, ûᵗ)
@@ -156,29 +173,52 @@ end
 
 function run_amp(
     rng::AbstractRNG;
-    observations::ContextualSBMObservations,
-    csbm::ContextualSBM,
+    observations::ContextualSBMObservations{R},
+    csbm::ContextualSBM{R},
     init_std::Real=1e-3,
-    iterations::Integer=10,
-    show_progress=false,
-)
-    (; storage, next_storage, temp_storage) = init_amp(rng; observations, csbm, init_std)
-    storage_history = [copy(storage)]
-    prog = Progress(iterations; desc="AMP-BP", enabled=show_progress)
-    for _ in 1:iterations
-        update_amp!(next_storage, temp_storage; storage, observations, csbm)
-        copy!(storage, next_storage)
-        push!(storage_history, copy(storage))
-        next!(prog)
+    max_iterations::Integer=200,
+    convergence_threshold=1e-3,
+    recent_past=10,
+    show_progress::Bool=false,
+) where {R}
+    (; N, P) = csbm
+    (; marginals, next_marginals, storage) = init_amp(rng; observations, csbm, init_std)
+
+    û_history = Matrix{R}(undef, N, max_iterations)
+    v̂_history = Matrix{R}(undef, P, max_iterations)
+    converged = false
+    prog = Progress(max_iterations; desc="AMP-BP", enabled=show_progress)
+
+    for t in 1:max_iterations
+        update_amp!(next_marginals, storage; marginals, observations, csbm)
+        copy!(marginals, next_marginals)
+
+        û_history[:, t] .= marginals.û
+        v̂_history[:, t] .= marginals.v̂
+
+        if t <= recent_past
+            û_recent_std = typemax(R)
+            v̂_recent_std = typemax(R)
+        else
+            û_recent_std = maximum(std(view(û_history, :, (t - recent_past):t); dims=2))
+            v̂_recent_std = maximum(std(view(v̂_history, :, (t - recent_past):t); dims=2))
+        end
+        converged = (
+            û_recent_std < convergence_threshold && v̂_recent_std < convergence_threshold
+        )
+        if converged
+            û_history = û_history[:, 1:t]
+            v̂_history = v̂_history[:, 1:t]
+            break
+        else
+            showvalues = [
+                (:û_recent_std, û_recent_std),
+                (:v̂_recent_std, v̂_recent_std),
+                (:convergence_threshold, convergence_threshold),
+            ]
+            next!(prog; showvalues)
+        end
     end
-    return storage_history
-end
 
-function overlap(; storage::AMPStorage, latents::ContextualSBMLatents)
-    û = sign.(storage.û)
-    @assert all(abs.(û) .> 0.5)
-    u = latents.u
-    q̂ᵤ = (1 / length(û)) * max(count_equalities(û, u), count_equalities(û, -u))
-    qᵤ = 2 * (q̂ᵤ - 0.5)
-    return qᵤ
+    return (; û_history, v̂_history, converged)
 end

src/utils.jl

@@ -1,3 +1,3 @@
 sigmoid(x) = 1 / (1 + exp(-x))
 
-count_equalities(x, y) = sum(x[i] ≈ y[i] for i in eachindex(x, y))
+freq_equalities(x, y) = mean(x[i] ≈ y[i] for i in eachindex(x, y))

test/csbm.jl

@@ -4,44 +4,54 @@ using Statistics
 using StochasticBlockModelVariants
 using Test
 
-rng = default_rng()
-
-function test_recovery(csbm::ContextualSBM)
+function test_recovery(csbm::ContextualSBM; test_u=true, test_v=true)
+    rng = default_rng()
     @assert effective_snr(csbm) > 1
     (; latents, observations) = rand(rng, csbm)
-    storage_history = run_amp(rng; observations, csbm)
-    overlap_history = [overlap(; storage, latents) for storage in storage_history]
-    @test last(overlap_history) > first(overlap_history)
-    @test last(overlap_history) > 0.5
+    (; u, v) = latents
+    (; û_history, v̂_history, converged) = run_amp(rng; observations, csbm)
+    @test converged
+    first_q = overlaps(; u, v, û=û_history[:, begin], v̂=v̂_history[:, begin])
+    last_q = overlaps(; u, v, û=û_history[:, end], v̂=v̂_history[:, end])
+    if test_u
+        @test last_q.qᵤ > first_q.qᵤ
+        @test last_q.qᵤ > 0.5
+    end
+    if test_v
+        @test last_q.qᵥ > first_q.qᵥ
+        @test last_q.qᵥ > 0.5
+    end
     return nothing
 end
 
 function test_jet(csbm::ContextualSBM)
+    rng = default_rng()
     (; observations) = rand(rng, csbm)
     @test_opt target_modules = (StochasticBlockModelVariants,) run_amp(
-        rng; observations, csbm, iterations=2
+        rng; observations, csbm, max_iterations=2
     )
     @test_call target_modules = (StochasticBlockModelVariants,) run_amp(
-        rng; observations, csbm, iterations=2
+        rng; observations, csbm, max_iterations=2
     )
     return nothing
 end
 
 function test_allocations(csbm::ContextualSBM)
+    rng = default_rng()
     (; observations) = rand(rng, csbm)
-    (; storage, next_storage, temp_storage) = init_amp(
+    (; marginals, next_marginals, storage) = init_amp(
         rng; observations, csbm, init_std=1e-3
     )
-    alloc = @allocated update_amp!(next_storage, temp_storage; storage, observations, csbm)
+    alloc = @allocated update_amp!(next_marginals, storage; marginals, observations, csbm)
     @test alloc == 0
     return nothing
 end
 
 @testset "Correct code" begin
     test_recovery(ContextualSBM(; N=10^3, P=10^3, d=5, λ=2, μ=2, ρ=0.0))  # AMP-BP
     test_recovery(ContextualSBM(; N=10^3, P=10^3, d=5, λ=0, μ=2, ρ=0.0))  # AMP
-    test_recovery(ContextualSBM(; N=10^3, P=10^3, d=5, λ=2, μ=0, ρ=0.0))  # BP
-    test_recovery(ContextualSBM(; N=10^2, P=10^2, d=5, λ=2, μ=0, ρ=0.5))  # semi-supervised
+    test_recovery(ContextualSBM(; N=10^3, P=10^3, d=5, λ=2, μ=0, ρ=0.0); test_v=false)  # BP
+    test_recovery(ContextualSBM(; N=10^2, P=10^2, d=5, λ=2, μ=2, ρ=0.5))  # semi-supervised
 end
 
 @testset "Good code" begin