GLMSBM runs but fails

Author: gdalle

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0-beta2"
 manifest_format = "2.0"
-project_hash = "3ba5719890857eaf711cde76af0aade7729abd07"
+project_hash = "5d6b46b606378213b4547f42572ed5ae33cb9d37"
 
 [[deps.ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

src/StochasticBlockModelVariants.jl

@@ -2,7 +2,7 @@ module StochasticBlockModelVariants
 
 using DensityInterface: logdensityof, densityof
 using Graphs: AbstractGraph, neighbors
-using LinearAlgebra: dot, mul!, norm
+using LinearAlgebra: LinearAlgebra, dot, mul!, norm, normalize!
 using PrecompileTools: @compile_workload
 using ProgressMeter: Progress, next!
 using Random: AbstractRNG, default_rng
@@ -22,6 +22,7 @@ include("utils.jl")
 include("csbm.jl")
 include("glmsbm.jl")
 include("csbm_inference.jl")
+include("glmsbm_inference.jl")
 
 # @compile_workload begin
 #     rng = default_rng()

src/abstract_sbm.jl

@@ -46,6 +46,19 @@ function affinities(sbm::AbstractSBM)
     return (; cᵢ, cₒ)
 end
 
+struct AffinityMatrix{R}
+    cᵢ::R
+    cₒ::R
+end
+
+function Base.getindex(C::AffinityMatrix, i, j)
+    if i == j
+        return C.cᵢ
+    else
+        return C.cₒ
+    end
+end
+
 """
     fraction_observed(sbm)
 
@@ -88,7 +101,7 @@ Sample a vector `Ξ` (Xi) whose components are equal to the community assignment
 function sample_mask(rng::AbstractRNG, sbm::AbstractSBM, communities::Vector{<:Integer})
     N = length(sbm)
     ρ = fraction_observed(sbm)
-    Ξ = Vector{Union{Missing, Int}}(undef, N)
+    Ξ = Vector{Union{Missing,Int}}(undef, N)
     Ξ .= missing
     for i in 1:N
         if rand(rng) < ρ
@@ -98,5 +111,4 @@ function sample_mask(rng::AbstractRNG, sbm::AbstractSBM, communities::Vector{<:I
     return Ξ
 end
 
-prior₊(::Type{R}, Ξᵢ) where {R} = ismissing(Ξᵢ) ? one(R) / 2 : R(Ξᵢ == 1)
-prior₋(::Type{R}, Ξᵢ) where {R} = ismissing(Ξᵢ) ? one(R) / 2 : R(Ξᵢ == -1)
+prior(::Type{R}, u, Ξᵢ) where {R} = ismissing(Ξᵢ) ? one(R) / 2 : R(Ξᵢ == u)

src/csbm_inference.jl

@@ -53,12 +53,12 @@ end
 ## Message-passing
 
 function init_amp(
-    rng::AbstractRNG; observations::ObservationsCSBM{R}, csbm::CSBM{R}, init_std
+    rng::AbstractRNG, observations::ObservationsCSBM{R}, csbm::CSBM{R}; init_std
 ) where {R}
     (; N, P) = csbm
     (; g, Ξ) = observations
 
-    û = 2 .* prior₊.(R, Ξ) .- one(R) .+ init_std .* randn(rng, R, N)
+    û = 2 .* prior.(R, 1, Ξ) .- one(R) .+ init_std .* randn(rng, R, N)
     v̂ = init_std .* randn(rng, R, P)
 
     û_no_feat = zeros(R, N)
@@ -69,7 +69,7 @@ function init_amp(
 
     χe₊ = Dict{Tuple{Int,Int},R}()
     for i in 1:N, j in neighbors(g, i)
-        χe₊[i, j] = prior₊(R, Ξ[i]) + init_std * randn(rng, R)
+        χe₊[i, j] = prior(R, 1, Ξ[i]) + init_std * randn(rng, R)
     end
     χ₊ = zeros(R, N)
 
@@ -79,7 +79,7 @@ function init_amp(
 end
 
 function update_amp!(
-    next_marginals::MarginalsCSBM{R};
+    next_marginals::MarginalsCSBM{R},
     marginals::MarginalsCSBM{R},
     observations::ObservationsCSBM{R},
     csbm::CSBM{R},
@@ -111,8 +111,8 @@ function update_amp!(
     # Estimation of the field h
     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
+    h̃₊ .= -h₊ .+ log.(prior.(R, +1, Ξ)) .+ û_no_feat
+    h̃₋ .= -h₋ .+ log.(prior.(R, -1, Ξ)) .- û_no_feat
 
     # BP update of the marginals
     for i in 1:N
@@ -144,25 +144,25 @@ function update_amp!(
 end
 
 function run_amp(
-    rng::AbstractRNG;
+    rng::AbstractRNG,
     observations::ObservationsCSBM{R},
-    csbm::CSBM{R},
+    csbm::CSBM{R};
     init_std=1e-3,
     max_iterations=200,
     convergence_threshold=1e-3,
     recent_past=10,
     show_progress=false,
 ) where {R}
     (; N, P) = csbm
-    (; marginals, next_marginals) = init_amp(rng; observations, csbm, init_std)
+    (; marginals, next_marginals) = 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 for CSBM", enabled=show_progress)
 
     for t in 1:max_iterations
-        update_amp!(next_marginals; marginals, observations, csbm)
+        update_amp!(next_marginals, marginals, observations, csbm)
         copy!(marginals, next_marginals)
 
         û_history[:, t] .= marginals.û
@@ -197,7 +197,7 @@ end
 
 function evaluate_amp(rng::AbstractRNG; csbm::CSBM, kwargs...)
     (; latents, observations) = rand(rng, csbm)
-    (; û_history, v̂_history, converged) = run_amp(rng; observations, csbm, kwargs...)
+    (; û_history, v̂_history, converged) = run_amp(rng, observations, csbm; kwargs...)
     (; qᵤ, qᵥ) = overlaps(;
         u=latents.u, v=latents.v, û=û_history[:, end], v̂=v̂_history[:, end]
     )

src/glmsbm.jl

@@ -110,3 +110,8 @@ fᵥ(::RademacherWeightPrior, Λ, Γ) = inv(abs2(cosh(Γ)))
 
 fₐ(::GaussianWeightPrior, Λ, Γ) = Γ / (Λ + 1)
 fᵥ(::GaussianWeightPrior, Λ, Γ) = 1 / (Λ + 1)
+
+function gₒ(ω, χ, V)
+    Znn = (1 + (2χ[1] - 1) * erf(ω / sqrt(2V))) / 2
+    return inv(sqrt(2π * V)) * (2χ[1] - 1) * exp(-ω^2 / (2V)) / Znn
+end

src/glmsbm_inference.jl

@@ -5,15 +5,15 @@
 
 # Fields
 
-- `a::Vector{R}`: length `M`
-- `v::Vector{R}`: length `M`
-- `Γ::Vector{R}`: length `M`
-- `ω::Vector{R}`: length `N`
-- `gₒ::Vector{R}`: length `N`
-- `ψl₊::Vector{R}`: length `N`
-- `χl₊::Vector{R}`: length `N`
-- `χe₊::Dict{Tuple{Int,Int},R}`: size `(N, N)`
-- `χ₊::Vector{R}`: length `N`
+- `a::Vector`: length `M`
+- `v::Vector`: length `M`
+- `Γ::Vector`: length `M`
+- `ω::Vector`: length `N`
+- `gₒ::Vector`: length `N`
+- `ψl::Vector{PlusMinusMeasure}`: length `N`
+- `χe::Dict{Tuple{Int,Int},PlusMinusMeasure}`: size `(N, N)`
+- `χl::Vector{PlusMinusMeasure}`: length `N`
+- `χ::Vector{PlusMinusMeasure}`: length `N`
 
 """
 @kwdef struct MarginalsGLMSBM{R<:Real}
@@ -22,31 +22,23 @@
     Γ::Vector{R}
     ω::Vector{R}
     gₒ::Vector{R}
-    ψl₊::Vector{R}
-    ψl₋::Vector{R}
-    χl₊::Vector{R}
-    χl₋::Vector{R}
-    χe₊::Dict{Tuple{Int,Int},R}
-    χe₋::Dict{Tuple{Int,Int},R}
-    χ₊::Vector{R}
-    χ₋::Vector{R}
+    ψl::Vector{PlusMinusMeasure{R}}
+    χe::Dict{Tuple{Int,Int},PlusMinusMeasure{R}}
+    χl::Vector{PlusMinusMeasure{R}}
+    χ::Vector{PlusMinusMeasure{R}}
 end
 
 function Base.copy(marginals::MarginalsGLMSBM)
-    return MarginalsCSBM(;
+    return MarginalsGLMSBM(;
         a=copy(marginals.a),
         v=copy(marginals.v),
         Γ=copy(marginals.Γ),
         ω=copy(marginals.ω),
         gₒ=copy(marginals.gₒ),
-        ψl₊=copy(marginals.ψl₊),
-        ψl₋=copy(marginals.ψl₋),
-        χl₊=copy(marginals.χl₊),
-        χl₋=copy(marginals.χl₋),
-        χe₊=copy(marginals.χe₊),
-        χe₋=copy(marginals.χe₋),
-        χ₊=copy(marginals.χ₊),
-        χ₋=copy(marginals.χ₋),
+        ψl=copy(marginals.ψl),
+        χe=copy(marginals.χe),
+        χl=copy(marginals.χl),
+        χ=copy(marginals.χ),
     )
 end
 
@@ -56,14 +48,10 @@ function Base.copy!(marginals_dest::MarginalsGLMSBM, marginals_source::Marginals
     copy!(marginals_dest.Γ, marginals_source.Γ)
     copy!(marginals_dest.ω, marginals_source.ω)
     copy!(marginals_dest.gₒ, marginals_source.gₒ)
-    copy!(marginals_dest.ψl₊, marginals_source.ψl₊)
-    copy!(marginals_dest.ψl₋, marginals_source.ψl₋)
-    copy!(marginals_dest.χl₊, marginals_source.χl₊)
-    copy!(marginals_dest.χl₋, marginals_source.χl₋)
-    copy!(marginals_dest.χe₊, marginals_source.χe₊)
-    copy!(marginals_dest.χe₋, marginals_source.χe₋)
-    copy!(marginals_dest.χ₊, marginals_source.χ₊)
-    copy!(marginals_dest.χ₋, marginals_source.χ₋)
+    copy!(marginals_dest.ψl, marginals_source.ψl)
+    copy!(marginals_dest.χe, marginals_source.χe)
+    copy!(marginals_dest.χl, marginals_source.χl)
+    copy!(marginals_dest.χ, marginals_source.χ)
     return marginals_dest
 end
 
@@ -82,14 +70,16 @@ function init_amp(
     ω = zeros(R, N)
     gₒ = zeros(R, N)
 
-    χl₊ = ones(R, N) / 2
-    χe₊ = Dict{Tuple{Int,Int},R}()
+    ψl = [PlusMinusMeasure(one(R) / 2) for μ in 1:N]
+    χe = Dict{Tuple{Int,Int},PlusMinusMeasure{R}}()
     for μ in 1:N, ν in neighbors(g, μ)
-        χe₊[μ, ν] = one(R) / 2 + init_std * randn(rng, R)
+        p = one(R) / 2 + init_std * randn(rng, R)
+        χe[μ, ν] = PlusMinusMeasure(clamp(p, zero(R), one(R)))
     end
-    χ₊ = zeros(R, N)
+    χl = [PlusMinusMeasure(one(R) / 2) for μ in 1:N]
+    χ = [PlusMinusMeasure(one(R) / 2) for μ in 1:N]
 
-    marginals = MarginalsGLMSBM(; a, v, Γ, ω, gₒ, ψl₊, ψl₋, χl₊, χl₋, χe₊, χe₋, χ₊, χ₋)
+    marginals = MarginalsGLMSBM(; a, v, Γ, ω, gₒ, ψl, χl, χe, χ)
     next_marginals = copy(marginals)
     return (; marginals, next_marginals)
 end
@@ -103,36 +93,30 @@ function update_amp!(
     (; N, M, c, λ, Pʷ) = glmsbm
     (; g, Ξ, F) = observations
     (; cᵢ, cₒ) = affinities(glmsbm)
+    C = AffinityMatrix(cᵢ, cₒ)
 
-    (aᵗ, vᵗ, Γᵗ, ωᵗ, gₒᵗ, ψl₊ᵗ, ψl₋ᵗ, χl₊ᵗ, χl₋ᵗ, χe₊ᵗ, χe₋ᵗ, χ₊ᵗ, χ₋ᵗ) = (
+    (aᵗ, vᵗ, Γᵗ, ωᵗ, gₒᵗ, ψlᵗ, χeᵗ, χlᵗ, χᵗ) = (
         marginals.a,
         marginals.v,
         marginals.Γ,
         marginals.ω,
         marginals.gₒ,
-        marginals.ψl₊,
-        marginals.ψl₋,
-        marginals.χl₊,
-        marginals.χl₋,
-        marginals.χe₊,
-        marginals.χe₋,
-        marginals.χ₊,
-        marginals.χ₋,
+        marginals.ψl,
+        marginals.χe,
+        marginals.χl,
+        marginals.χ,
     )
-    (aᵗ⁺¹, vᵗ⁺¹, Γᵗ⁺¹, ωᵗ⁺¹, gₒᵗ⁺¹, ψl₊ᵗ⁺¹, ψl₋ᵗ⁺¹, χl₊ᵗ⁺¹, χl₋ᵗ⁺¹, χe₊ᵗ⁺¹, χe₋ᵗ⁺¹, χ₊ᵗ⁺¹, χ₋ᵗ⁺¹) = (
+
+    (aᵗ⁺¹, vᵗ⁺¹, Γᵗ⁺¹, ωᵗ⁺¹, gₒᵗ⁺¹, ψlᵗ⁺¹, χeᵗ⁺¹, χlᵗ⁺¹, χᵗ⁺¹) = (
         next_marginals.a,
         next_marginals.v,
         next_marginals.Γ,
         next_marginals.ω,
         next_marginals.gₒ,
-        next_marginals.ψl₊,
-        next_marginals.ψl₋,
-        next_marginals.χl₊,
-        next_marginals.χl₋,
-        next_marginals.χe₊,
-        next_marginals.χe₋,
-        next_marginals.χ₊,
-        next_marginals.χ₋,
+        next_marginals.ψl,
+        next_marginals.χe,
+        next_marginals.χl,
+        next_marginals.χ,
     )
 
     # AMP update of ω, V
@@ -141,29 +125,77 @@ function update_amp!(
     ωᵗ⁺¹ .-= Vᵗ⁺¹ .* gₒᵗ
 
     # AMP update of ψl, gₒ, μ, Λ, Γ
-    ψl₊ .= missing  # TODO: fix
-    gₒᵗ⁺¹ = gₒ.(ωᵗ⁺¹, χl₊ᵗ, χl₋ᵗ, Ref(Vᵗ⁺¹))
+    for μ in 1:N, sμ in (-1, 1)
+        ψlᵗ⁺¹[μ][sμ] = (one(R) + sμ * erf(ωᵗ⁺¹[μ] / sqrt(2Vᵗ⁺¹))) / 2
+    end
+    gₒᵗ⁺¹ = gₒ.(ωᵗ⁺¹, χlᵗ, Ref(Vᵗ⁺¹))
     Λᵗ⁺¹ = sum(abs2, gₒᵗ⁺¹) / M
     mul!(Γᵗ⁺¹, F, gₒᵗ⁺¹)
     Γᵗ⁺¹ .+= Λᵗ⁺¹ .* aᵗ
 
     # AMP update of the estimated marginals a, v
-    aᵗ⁺¹ .= fₐ.(Ref(Λᵗ⁺¹), Γᵗ⁺¹)
-    vᵗ⁺¹ .= fᵥ.(Ref(Λᵗ⁺¹), Γᵗ⁺¹)
+    aᵗ⁺¹ .= fₐ.(Ref(Pʷ), Ref(Λᵗ⁺¹), Γᵗ⁺¹)
+    vᵗ⁺¹ .= fᵥ.(Ref(Pʷ), Ref(Λᵗ⁺¹), Γᵗ⁺¹)
 
     # BP update of the field h
-    h₊ᵗ⁺¹
-    h₋ᵗ⁺¹
+    hᵗ⁺¹ = PlusMinusMeasure(zero(R), zero(R))
+    for s in (-1, 1)
+        hᵗ⁺¹[s] = sum(C[s, sμ] * χᵗ[μ][sμ] for μ in 1:N for sμ in (-1, 1))
+    end
 
     # BP update of the messages χe and of the marginals χ
-    χe₊ᵗ⁺¹
-    χe₋ᵗ⁺¹
-    χ₊ᵗ⁺¹
-    χ₋ᵗ⁺¹
+    for μ in 1:N, ν in neighbors(g, μ)
+        for sμ in (-1, 1)
+            χeᵗ⁺¹[μ, ν][sμ] = prior(R, sμ, Ξ[μ]) * exp(-hᵗ⁺¹[sμ]) * ψlᵗ⁺¹[μ][sμ]
+            for η in neighbors(g, μ)
+                if η != ν
+                    χeᵗ⁺¹[μ, ν][sμ] *= sum(C[sη, sμ] * χeᵗ[η, μ][sη] for sη in (-1, 1))
+                end
+            end
+        end
+        normalize!(χeᵗ⁺¹[μ, ν])
+    end
+    for μ in 1:N
+        for sμ in (-1, 1)
+            χᵗ⁺¹[μ][sμ] = prior(R, sμ, Ξ[μ]) * exp(-hᵗ⁺¹[sμ]) * ψlᵗ⁺¹[μ][sμ]
+            for η in neighbors(g, μ)
+                χᵗ⁺¹[μ][sμ] *= sum(C[sη, sμ] * χeᵗ[η, μ][sη] for sη in (-1, 1))
+            end
+        end
+        normalize!(χᵗ⁺¹[μ])
+    end
 
     # BP update of the SBM-to-GLM messages χl
-    χl₊ᵗ⁺¹
-    χl₋ᵗ⁺¹
+    for μ in 1:N
+        for sμ in (-1, 1)
+            χlᵗ⁺¹[μ][sμ] = prior(R, sμ, Ξ[μ]) * exp(-hᵗ⁺¹[sμ])
+            for η in neighbors(g, μ)
+                χlᵗ⁺¹[μ][sμ] *= sum(C[sη, sμ] * χeᵗ[η, μ][sη] for sη in (-1, 1))
+            end
+        end
+        normalize!(χlᵗ⁺¹[μ])
+    end
 
     return nothing
 end
+
+function run_amp(
+    rng::AbstractRNG,
+    observations::ObservationsGLMSBM{R},
+    glmsbm::GLMSBM{R};
+    init_std=1e-3,
+    max_iterations=200,
+    convergence_threshold=1e-3,
+    recent_past=10,
+    show_progress=false,
+) where {R}
+    (; marginals, next_marginals) = init_amp(rng; observations, glmsbm, init_std)
+    converged = false
+    prog = Progress(max_iterations; desc="AMP-BP for GLM-SBM", enabled=show_progress)
+    for t in 1:max_iterations
+        update_amp!(next_marginals; marginals, observations, glmsbm)
+        copy!(marginals, next_marginals)
+        next!(prog)
+    end
+    return marginals
+end