Partially revealed communities

Author: gdalle

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.9.2"
 manifest_format = "2.0"
-project_hash = "76c4ad29b7fdf61caa5b4fe006d178c04b45b443"
+project_hash = "eabebeb24d3db0211e95e73a9e2e965bf16ba4e1"
 
 [[deps.ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
@@ -158,6 +158,12 @@ version = "1.4.0"
 deps = ["Unicode"]
 uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
+[[deps.ProgressMeter]]
+deps = ["Distributed", "Printf"]
+git-tree-sha1 = "d7a7aef8f8f2d537104f170139553b14dfe39fe9"
+uuid = "92933f4c-e287-5a05-a399-4b506db050ca"
+version = "1.7.2"
+
 [[deps.REPL]]
 deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
 uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

Project.toml

@@ -7,6 +7,7 @@ version = "0.1.0"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -15,7 +16,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 Graphs = "1.8"
 PrecompileTools = "1.1"
-ProgressLogging = "0.1"
+ProgressMeter = "1.7"
 SimpleWeightedGraphs = "1.4"
 julia = "1.9"
 

src/StochasticBlockModelVariants.jl

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

src/csbm.jl

@@ -79,13 +79,18 @@ end
 The observations generated by sampling from a [`ContextualSBM`](@ref).
 
 # Fields
-
-- `g::AbstractGraph`: undirected, unweighted graph generated from the adjacency matrix
+- `A::AbstractMatrix`: symmetric boolean adjacency matrix, size `(N, N)`
+- `g::AbstractGraph`: undirected unweighted graph generated from `A`
 - `B::Matrix`: feature matrix, size `(P, N)`
+- `Ξ::Vector`: revealed communities `±1` for a fraction of nodes and `0` for the rest, length `N`
 """
-@kwdef struct ContextualSBMObservations{R<:Real,G<:AbstractGraph{Int}}
-    B::Matrix{R}
+@kwdef struct ContextualSBMObservations{
+    R<:Real,M<:AbstractMatrix{Bool},G<:AbstractGraph{Int}
+}
+    A::M
     g::G
+    B::Matrix{R}
+    Ξ::Vector{Int}
 end
 
 ## Simulation
@@ -96,7 +101,7 @@ end
 Sample from a [`ContextualSBM`](@ref) and return a named tuple `(; latents, observations)`.
 """
 function Base.rand(rng::AbstractRNG, csbm::ContextualSBM)
-    (; μ, N, P) = csbm
+    (; N, P, μ, ρ) = csbm
     (; cᵢ, cₒ) = affinities(csbm)
 
     u = rand(rng, (-1, +1), N)
@@ -121,10 +126,15 @@ function Base.rand(rng::AbstractRNG, csbm::ContextualSBM)
     g = SimpleWeightedGraph(A)
 
     B = randn(rng, P, N)
-    for i in 1:N, α in 1:P
-        B[α, i] += sqrt(μ / N) * v[α] * u[i]
+    B .+= sqrt(μ / N) .* v .* u'
+
+    Ξ = zeros(Int, N)
+    for i in 1:N
+        if rand(rng) < ρ
+            Ξ[i] = u[i]
+        end
     end
 
-    observations = ContextualSBMObservations(; g, B)
+    observations = ContextualSBMObservations(; A, g, B, Ξ)
     return (; latents, observations)
 end

src/csbm_inference.jl

@@ -59,16 +59,15 @@ end
 
 function init_amp(
     rng::AbstractRNG;
-    observations::ContextualSBMObservations{R1},
-    csbm::ContextualSBM{R2},
-    init_std::R3,
-) where {R1,R2,R3}
-    R = promote_type(R1, R2, R3)
+    observations::ContextualSBMObservations{R},
+    csbm::ContextualSBM{R},
+    init_std,
+) where {R}
     (; N, P) = csbm
-    (; g) = observations
+    (; g, Ξ) = observations
 
-    û = init_std .* randn(rng, R, N)
-    v̂ = init_std .* randn(rng, R, P)
+    û = prior₊.(R, Ξ) + init_std .* randn(rng, R, N)
+    v̂ = 2 .* prior₊.(R, Ξ) .- one(R) .+ 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)
@@ -86,15 +85,18 @@ function init_amp(
     return (; storage, next_storage, temp_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,
-    temp_storage::AMPTempStorage;
-    storage::AMPStorage,
-    observations::ContextualSBMObservations,
+    next_storage::AMPStorage{R},
+    temp_storage::AMPTempStorage{R};
+    storage::AMPStorage{R},
+    observations::ContextualSBMObservations{R},
     csbm::ContextualSBM{R},
 ) where {R}
     (; d, λ, μ, N, P) = csbm
-    (; g, B) = observations
+    (; g, B, Ξ) = observations
     (; cᵢ, cₒ) = affinities(csbm)
 
     ûᵗ, v̂ᵗ, χ₊eᵗ = storage.û, storage.v̂, storage.χ₊e
@@ -109,19 +111,19 @@ function update_amp!(
     mul!(v̂_no_comm, B, ûᵗ)
     v̂_no_comm .*= sqrt(μ / N)
     v̂_no_comm .-= (μ / N) .* v̂ᵗ .* (N - ûₜ_sum2)
-    v̂ᵗ⁺¹ .= v̂_no_comm ./ (1 + σᵥ_no_comm)
-    σᵥ = 1 / (1 + σᵥ_no_comm)
+    v̂ᵗ⁺¹ .= v̂_no_comm ./ (one(R) + σᵥ_no_comm)
+    σᵥ = one(R) / (one(R) + σᵥ_no_comm)
 
     # BP estimation of u
     mul!(û_no_feat, B', v̂ᵗ⁺¹)
     û_no_feat .*= sqrt(μ / N)
     û_no_feat .-= (μ / (N / P)) .* σᵥ .* ûᵗ
 
     # Estimation of the field h
-    h₊ = (1 / 2N) * (cᵢ * (N + ûₜ_sum) + cₒ * (N - ûₜ_sum))
-    h₋ = (1 / 2N) * (cₒ * (N + ûₜ_sum) + cᵢ * (N - ûₜ_sum))
-    h̃₊ .= -h₊ .+ log(one(R) / 2) .+ û_no_feat
-    h̃₋ .= -h₋ .+ log(one(R) / 2) .- û_no_feat
+    h₊ = (one(R) / 2N) * (cᵢ * (N + ûₜ_sum) + cₒ * (N - ûₜ_sum))
+    h₋ = (one(R) / 2N) * (cₒ * (N + ûₜ_sum) + cᵢ * (N - ûₜ_sum))
+    h̃₊ .= -h₊ .+ log.(prior₊.(R, Ξ)) .+ û_no_feat
+    h̃₋ .= -h₋ .+ log.(prior₋.(R, Ξ)) .- û_no_feat
 
     # BP update of the marginals
     for i in 1:N
@@ -147,7 +149,7 @@ function update_amp!(
     end
 
     # BP estimation of u
-    ûᵗ⁺¹ .= 2 .* χ₊ .- 1
+    ûᵗ⁺¹ .= 2 .* χ₊ .- one(R)
 
     return nothing
 end
@@ -158,18 +160,21 @@ function run_amp(
     csbm::ContextualSBM,
     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)]
-    for iter in 1:iterations
+    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)
     end
     return storage_history
 end
 
-function evaluate_amp(; storage::AMPStorage, latents::ContextualSBMLatents)
+function overlap(; storage::AMPStorage, latents::ContextualSBMLatents)
     û = sign.(storage.û)
     @assert all(abs.(û) .> 0.5)
     u = latents.u

test/csbm.jl

@@ -10,8 +10,8 @@ function test_recovery(csbm::ContextualSBM)
     @assert effective_snr(csbm) > 1
     (; latents, observations) = rand(rng, csbm)
     storage_history = run_amp(rng; observations, csbm)
-    overlap_history = [evaluate_amp(; storage, latents) for storage in storage_history]
-    @test last(overlap_history) > 10 * first(overlap_history)
+    overlap_history = [overlap(; storage, latents) for storage in storage_history]
+    @test last(overlap_history) > first(overlap_history)
     @test last(overlap_history) > 0.5
     return nothing
 end
@@ -41,6 +41,7 @@ end
     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
 end
 
 @testset "Good code" begin