BP is working for CSBM

Author: gdalle

Manifest.toml

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

src/contextual_sbm.jl

@@ -79,8 +79,11 @@ function Base.rand(rng::AbstractRNG, csbm::ContextualSBM)
 
     r = rand(rng, N, N)
     Is, Js = Int[], Int[]
-    for i in 1:N, j in 1:i
-        if ((u[i] == u[j]) && (r[i, j] < cᵢ / N)) || ((u[i] != u[j]) && (r[i, j] < cₒ / N))
+    for i in 1:N, j in (i + 1):N
+        if (
+            ((u[i] == u[j]) && (r[i, j] < cᵢ / N)) ||  # same
+            ((u[i] != u[j]) && (r[i, j] < cₒ / N))  # diff
+        )
             push!(Is, i)
             push!(Js, j)
         end
@@ -155,6 +158,16 @@ end
     χ₊::Vector{R}
 end
 
+function Base.copy(temp_storage::AMPTempStorage)
+    return AMPTempStorage(;
+        û_no_feat=copy(temp_storage.û_no_feat),
+        v̂_no_graph=copy(temp_storage.v̂_no_graph),
+        h̃₊=copy(temp_storage.h̃₊),
+        h̃₋=copy(temp_storage.h̃₋),
+        χ₊=copy(temp_storage.χ₊),
+    )
+end
+
 function init_amp(
     rng::AbstractRNG;
     observations::ContextualSBMObservations,
@@ -191,8 +204,8 @@ function update_amp!(
     temp_storage::AMPTempStorage;
     storage::AMPStorage,
     observations::ContextualSBMObservations,
-    csbm::ContextualSBM,
-)
+    csbm::ContextualSBM{R},
+) where {R}
     (; d, λ, μ, N, P) = csbm
     (; B, G) = observations
     (; cᵢ, cₒ) = affinities(csbm)
@@ -223,18 +236,19 @@ function update_amp!(
     h₊ = (1 / N) * sum(cᵢ * (1 + ûᵗ[i]) / 2 + cₒ * (1 - ûᵗ[i]) / 2 for i in 1:N)
     h₋ = (1 / N) * sum(cₒ * (1 + ûᵗ[i]) / 2 + cᵢ * (1 - ûᵗ[i]) / 2 for i in 1:N)
     for i in 1:N
-        h̃₊[i] = -h₊ + û_no_feat[i]
-        h̃₋[i] = -h₋ - û_no_feat[i]
+        h̃₊[i] = -h₊ + log(one(R) / 2) + û_no_feat[i]
+        h̃₋[i] = -h₋ + log(one(R) / 2) - û_no_feat[i]
     end
 
     # BP update of the messages
     for i in 1:N, j in neighbors(G, i)
         s_ij = h̃₊[i] - h̃₋[i]
         for k in neighbors(G, i)
-            k != j || continue
-            num = (cₒ + 2λ * sqrt(d) * χ₊eᵗ[k, i])
-            den = (cᵢ - 2λ * sqrt(d) * χ₊eᵗ[k, i])
-            s_ij += log(num / den)
+            if k != j
+                num = (cₒ + 2λ * sqrt(d) * χ₊eᵗ[k, i])
+                den = (cᵢ - 2λ * sqrt(d) * χ₊eᵗ[k, i])
+                s_ij += log(num / den)
+            end
         end
         χ₊eᵗ⁺¹[i, j] = sigmoid(s_ij)
     end
@@ -267,19 +281,21 @@ function run_amp(
 )
     (; storage, next_storage, temp_storage) = init_amp(rng; observations, csbm, init_std)
     storage_history = [copy(storage)]
+    temp_storage_history = [copy(temp_storage)]
     @showprogress "AMP-BP" for iter in 1:iterations
         update_amp!(next_storage, temp_storage; storage, observations, csbm)
         copy!(storage, next_storage)
         push!(storage_history, copy(storage))
+        push!(temp_storage_history, copy(temp_storage))
     end
     return storage_history
 end
 
 function evaluate_amp(; storage::AMPStorage, latents::ContextualSBMLatents)
+    û = sign.(storage.û)
+    @assert all(abs.(û) .> 0.5)
     u = latents.u
-    N = length(u)
-    û = 2 .* Int.(storage.û .> 0) .- 1
-    q̂ᵤ = (1 / N) * max(count_equalities(û, u), count_equalities(û, -u))
+    q̂ᵤ = (1 / length(û)) * max(count_equalities(û, u), count_equalities(û, -u))
     qᵤ = 2 * (q̂ᵤ - 0.5)
     return qᵤ
 end

src/utils.jl

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

test/contextual_sbm.jl

@@ -1,34 +1,21 @@
 using Random: default_rng
+using Statistics
 using StochasticBlockModelVariants
 using Test
 
 rng = default_rng()
 
-# Parameters
-
-d = 5.0
-λ = 2.0
-μ = 2.0
-N = 3 * 10^3
-P = 3 * 10^2
-
-init_std = 1e-3
-iterations = 10
-
-# Sampling
-
-csbm = ContextualSBM(; d, λ, μ, N, P)
-
-(; latents, observations) = rand(rng, csbm);
-
-# Inference
-
-(; storage, next_storage, temp_storage) = init_amp(rng; observations, csbm, init_std);
-
-storage_history = run_amp(rng; observations, csbm, init_std, iterations);
-
-@test effective_snr(csbm) > 1
-
-overlap_history = [evaluate_amp(; storage, latents) for storage in storage_history]
-
-@test_broken last(overlap_history) > 0.5
+function test_recovery(; d, λ, μ, N, P, init_std=1e-3, iterations=20)
+    csbm = ContextualSBM(; d, λ, μ, N, P)
+    @assert effective_snr(csbm) > 1
+    (; latents, observations) = rand(rng, csbm)
+    storage_history = run_amp(rng; observations, csbm, init_std, iterations)
+    overlap_history = [evaluate_amp(; storage, latents) for storage in storage_history]
+    @test last(overlap_history) > 10 * first(overlap_history)
+    @test last(overlap_history) > 0.5
+    return nothing
+end
+
+test_recovery(; d=5.0, λ=2.0, μ=2.0, N=10^3, P=10^3)  # AMP-BP
+test_recovery(; d=5.0, λ=0.0, μ=2.0, N=10^3, P=10^3)  # AMP
+test_recovery(; d=5.0, λ=2.0, μ=0.0, N=10^3, P=10^3)  # BP