Working GLM-SBM ?

Author: gdalle

.vscode/settings.json

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0-beta2"
 manifest_format = "2.0"
-project_hash = "ba25bdf8dd9a1a3a609260b65801f3325dbdaa20"
+project_hash = "06fe3f2e3a2d6f17958684cbeb77c96e1e718454"
 
 [[deps.ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
@@ -75,12 +75,6 @@ git-tree-sha1 = "1cf1d7dcb4bc32d7b4a5add4232db3750c27ecb4"
 uuid = "86223c79-3864-5bf0-83f7-82e725a168b6"
 version = "1.8.0"
 
-[[deps.Infiltrator]]
-deps = ["InteractiveUtils", "Markdown", "REPL", "UUIDs"]
-git-tree-sha1 = "04de041e4590428cccbb026d86e5d670513be2e3"
-uuid = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
-version = "1.6.4"
-
 [[deps.Inflate]]
 git-tree-sha1 = "5cd07aab533df5170988219191dfad0519391428"
 uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9"
@@ -149,12 +143,6 @@ version = "0.3.26"
     ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
     InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 
-[[deps.LogarithmicNumbers]]
-deps = ["Random"]
-git-tree-sha1 = "8522befb54ff3b4bcf17d57b14b884d536a22015"
-uuid = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
-version = "1.2.1"
-
 [[deps.Logging]]
 uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
 

Manifest.toml

@@ -7,9 +7,7 @@ version = "0.1.0"
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
-Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -20,6 +18,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 DensityInterface = "0.4"
+DocStringExtensions = "0.9"
 Graphs = "1.8"
 PrecompileTools = "1.1"
 ProgressMeter = "1.7"

Project.toml

@@ -12,7 +12,6 @@ module StochasticBlockModelVariants
 using DensityInterface: logdensityof, densityof
 using DocStringExtensions
 using Graphs: AbstractGraph, neighbors
-using Infiltrator
 using LinearAlgebra: LinearAlgebra, dot, mul!, norm, normalize!
 using PrecompileTools: @compile_workload
 using ProgressMeter: Progress, next!
@@ -22,10 +21,11 @@ using SpecialFunctions: erf
 using Statistics: mean, std
 using SparseArrays: SparseMatrixCSC, sparse, findnz
 
+export AbstractSBM
 export nb_features, average_degree, communities_snr, affinities, features_snr, effective_snr
 export CSBM, LatentsCSBM, ObservationsCSBM
 export GLMSBM, LatentsGLMSBM, ObservationsGLMSBM
-export overlaps
+export GaussianWeightPrior, RademacherWeightPrior
 export init_amp, update_amp!, run_amp, evaluate_amp
 
 include("abstract_sbm.jl")
@@ -34,11 +34,4 @@ include("glmsbm.jl")
 include("csbm_inference.jl")
 include("glmsbm_inference.jl")
 
-# @compile_workload begin
-#     rng = default_rng()
-#     csbm = CSBM(; N=10^2, P=10^2, d=5, λ=2, μ=2, ρ=0.1)
-#     (; observations) = rand(rng, csbm)
-#     run_amp(rng; observations, csbm, max_iterations=2)
-# end
-
 end

src/StochasticBlockModelVariants.jl

@@ -105,13 +105,17 @@ sigmoid(x) = 1 / (1 + exp(-x))
 freq_equalities(x, y) = mean(x[i] ≈ y[i] for i in eachindex(x, y))
 
 function discrete_overlap(u, û)
-    q̂ᵤ = max(freq_equalities(û, u), freq_equalities(û, -u))
+    R = eltype(û)
+    û_sign = sign.(û)
+    û_sign[abs.(û) .< eps(R)] .= 1
+    q̂ᵤ = max(freq_equalities(u, û_sign), freq_equalities(u, -û_sign))
     qᵤ = 2 * (q̂ᵤ - one(R) / 2)
     return qᵤ
 end
 
 function continuous_overlap(v, v̂)
+    R = eltype(v)
     q̂ᵥ = max(abs(dot(v̂, v)), abs(dot(v̂, -v)))
-    qᵥ = q̂ᵥ / (eps() + norm(v̂) * norm(v))
+    qᵥ = q̂ᵥ / (eps(R) + norm(v̂) * norm(v))
     return qᵥ
 end

src/abstract_sbm.jl

@@ -152,7 +152,7 @@ function run_amp(
     observations::ObservationsCSBM,
     csbm::CSBM;
     init_std=1e-3,
-    max_iterations=200,
+    max_iterations=100,
     convergence_threshold=1e-3,
     recent_past=10,
     show_progress=false,
@@ -177,8 +177,8 @@ function run_amp(
             û_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))
+            û_recent_std = mean(std(view(û_history, :, (t - recent_past):t); dims=2))
+            v̂_recent_std = mean(std(view(v̂_history, :, (t - recent_past):t); dims=2))
         end
         converged = (
             û_recent_std < convergence_threshold && v̂_recent_std < convergence_threshold
@@ -203,7 +203,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...)
-    qᵤ = discrete_overlap(latents.u, û_history[:, end])
-    qᵥ = continuous_overlap(latents.v, v̂_history[:, end])
-    return (; qᵤ, qᵥ)
+    q_dis = discrete_overlap(latents.u, û_history[:, end])
+    q_cont = continuous_overlap(latents.v, v̂_history[:, end])
+    return (; q_dis, q_cont, converged)
 end

src/csbm_inference.jl

@@ -71,7 +71,7 @@ $(TYPEDFIELDS)
     g::G
     "revealed communities `±1` for a fraction `ρ` of nodes and `missing` for the rest, ngth `N`"
     Ξ::Vector{Union{Int,Missing}}
-    "feature matrix, size `(M, N)`"
+    "feature matrix, size `(N, M)`"
     F::Matrix{R}
 end
 
@@ -85,7 +85,7 @@ Sample from a [`GLMSBM`](@ref) and return a named tuple `(; latents, observation
 function Base.rand(rng::AbstractRNG, glmsbm::GLMSBM)
     (; N, M, ρ, Pʷ) = glmsbm
 
-    F = randn(rng, M, N) ./ sqrt(M)
+    F = randn(rng, N, M) ./ sqrt(M)
 
     w = [rand(Pʷ) for l in 1:M]
     s = round.(Int, sign.(F * w))
@@ -103,6 +103,9 @@ end
 struct RademacherWeightPrior{R} end
 struct GaussianWeightPrior{R} end
 
+RademacherWeightPrior() = RademacherWeightPrior{Float64}()
+GaussianWeightPrior() = GaussianWeightPrior{Float64}()
+
 function Base.rand(rng::AbstractRNG, ::RademacherWeightPrior{R}) where {R}
     return rand(rng, (-one(R), +one(R)))
 end

src/glmsbm.jl

@@ -66,7 +66,7 @@ end
 ind(s) = mod(s, 3)  # sends 1 to 1 and -1 to 2
 
 function init_amp(
-    rng::AbstractRNG; observations::ObservationsGLMSBM{R1}, glmsbm::GLMSBM{R2}, init_std::R3
+    rng::AbstractRNG, observations::ObservationsGLMSBM{R1}, glmsbm::GLMSBM{R2}; init_std::R3
 ) where {R1,R2,R3}
     R = promote_type(R1, R2, R3)
     (; N, M) = glmsbm
@@ -96,7 +96,7 @@ function init_amp(
 end
 
 function update_amp!(
-    next_marginals::MarginalsGLMSBM{R};
+    next_marginals::MarginalsGLMSBM{R},
     marginals::MarginalsGLMSBM{R},
     observations::ObservationsGLMSBM,
     glmsbm::GLMSBM,
@@ -143,38 +143,21 @@ function update_amp!(
     end
     @views gₒᵗ⁺¹ = gₒ.(ωᵗ⁺¹, χlᵗ[1, :], Ref(Vᵗ⁺¹))
     Λᵗ⁺¹ = sum(abs2, gₒᵗ⁺¹) / M
-    mul!(Γᵗ⁺¹, F, gₒᵗ⁺¹)
+    mul!(Γᵗ⁺¹, F', gₒᵗ⁺¹)
     Γᵗ⁺¹ .+= Λᵗ⁺¹ .* ŵᵗ
 
     # AMP update of the estimated marginals a, v
-    ŵᵗ⁺¹ .= fₐ.(Ref(Pʷ), Ref(Λᵗ⁺¹), Γᵗ⁺¹)
-    vᵗ⁺¹ .= fᵥ.(Ref(Pʷ), Ref(Λᵗ⁺¹), Γᵗ⁺¹)
+    ŵᵗ⁺¹ .= fₐ.(Ref(Pʷ), Λᵗ⁺¹, Γᵗ⁺¹)
+    vᵗ⁺¹ .= fᵥ.(Ref(Pʷ), Λᵗ⁺¹, Γᵗ⁺¹)
 
     # BP update of the field h
     hᵗ⁺¹ = Vector{R}(undef, 2)
     for s in (-1, 1)
-        hᵗ⁺¹[ind(s)] = sum(
-            C[ind(s), ind(sμ)] * χᵗ[ind(sμ), μ] for μ in 1:N for sμ in (-1, 1)
-        )
+        hᵗ⁺¹[ind(s)] =
+            sum(C[ind(s), ind(sμ)] * χᵗ[ind(sμ), μ] for μ in 1:N for sμ in (-1, 1)) / N
     end
 
     # BP update of the messages χe and of the marginals χ
-    for μ in 1:N, ν in neighbors(g, μ)
-        for sμ in (-1, 1)
-            χeᵗ⁺¹[ind(sμ), μ, ν] =
-                prior(R, sμ, Ξ[μ]) * exp(-hᵗ⁺¹[ind(sμ)]) * ψlᵗ⁺¹[ind(sμ), μ]
-            for η in neighbors(g, μ)
-                if η != ν
-                    χeᵗ⁺¹[ind(sμ), μ, ν] *= sum(
-                        C[ind(sη), ind(sμ)] * χeᵗ[ind(sη), η, μ] for sη in (-1, 1)
-                    )
-                end
-            end
-        end
-        normalization = χeᵗ⁺¹[1, μ, ν] + χeᵗ⁺¹[2, μ, ν]
-        χeᵗ⁺¹[1, μ, ν] /= normalization
-        χeᵗ⁺¹[2, μ, ν] /= normalization
-    end
 
     for μ in 1:N
         for sμ in (-1, 1)
@@ -188,6 +171,16 @@ function update_amp!(
         @views χᵗ⁺¹[:, μ] ./= sum(χᵗ⁺¹[:, μ])
     end
 
+    for μ in 1:N, ν in neighbors(g, μ)
+        for sμ in (-1, 1)
+            extra_factor = sum(C[ind(sν), ind(sμ)] * χeᵗ[ind(sν), ν, μ] for sν in (-1, 1))
+            χeᵗ⁺¹[ind(sμ), μ, ν] = χᵗ⁺¹[ind(sμ), μ] / extra_factor
+        end
+        normalization = χeᵗ⁺¹[1, μ, ν] + χeᵗ⁺¹[2, μ, ν]
+        χeᵗ⁺¹[1, μ, ν] /= normalization
+        χeᵗ⁺¹[2, μ, ν] /= normalization
+    end
+
     # BP update of the SBM-to-GLM messages χl
     for μ in 1:N
         for sμ in (-1, 1)
@@ -201,7 +194,7 @@ function update_amp!(
         @views χlᵗ⁺¹[:, μ] ./= sum(χlᵗ⁺¹[:, μ])
     end
 
-    @views ᵗ⁺¹ .= 2 .* χ[1, :] .- one(R)
+    @views ŝᵗ⁺¹ .= 2 .* χᵗ⁺¹[1, :] .- one(R)
 
     return nothing
 end
@@ -211,13 +204,13 @@ function run_amp(
     observations::ObservationsGLMSBM,
     glmsbm::GLMSBM;
     init_std=1e-3,
-    max_iterations=200,
+    max_iterations=100,
     convergence_threshold=1e-3,
     recent_past=10,
     show_progress=false,
 )
-    (; N, M) = csbm
-    (; marginals, next_marginals) = init_amp(rng; observations, glmsbm, init_std)
+    (; N, M) = glmsbm
+    (; marginals, next_marginals) = init_amp(rng, observations, glmsbm; init_std)
 
     R = eltype(marginals)
     ŝ_history = Matrix{R}(undef, N, max_iterations)
@@ -226,7 +219,7 @@ function run_amp(
     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)
+        update_amp!(next_marginals, marginals, observations, glmsbm)
         copy!(marginals, next_marginals)
 
         ŝ_history[:, t] .= marginals.ŝ
@@ -236,8 +229,8 @@ function run_amp(
             ŝ_recent_std = typemax(R)
             ŵ_recent_std = typemax(R)
         else
-            ŝ_recent_std = maximum(std(view(ŝ_history, :, (t - recent_past):t); dims=2))
-            ŵ_recent_std = maximum(std(view(ŵ_history, :, (t - recent_past):t); dims=2))
+            ŝ_recent_std = mean(std(view(ŝ_history, :, (t - recent_past):t); dims=2))
+            ŵ_recent_std = mean(std(view(ŵ_history, :, (t - recent_past):t); dims=2))
         end
         converged = (
             ŝ_recent_std < convergence_threshold && ŵ_recent_std < convergence_threshold
@@ -261,8 +254,8 @@ end
 
 function evaluate_amp(rng::AbstractRNG, glmsbm::GLMSBM; kwargs...)
     (; latents, observations) = rand(rng, glmsbm)
-    (; ŝ_history, ŵ_history, converged) = run_amp(rng, observations, csbm; kwargs...)
-    qᵤ = discrete_overlap(latents.s, ŝ_history[:, end])
-    qᵥ = continuous_overlap(latents.w, ŵ_history[:, end])
-    return (; qᵤ, qᵥ)
+    (; ŝ_history, ŵ_history, converged) = run_amp(rng, observations, glmsbm; kwargs...)
+    q_dis = discrete_overlap(latents.s, ŝ_history[:, end])
+    q_cont = continuous_overlap(latents.w, ŵ_history[:, end])
+    return (; q_dis, q_cont, converged)
 end

src/glmsbm_inference.jl

@@ -0,0 +1,19 @@
+using JET
+using Random: default_rng
+using Statistics
+using StochasticBlockModelVariants
+using Test
+
+function test_allocations(sbm::AbstractSBM)
+    rng = default_rng()
+    (; observations) = rand(rng, sbm)
+    (; marginals, next_marginals) = init_amp(rng, observations, sbm; init_std=1e-3)
+    alloc = @allocated update_amp!(next_marginals, marginals, observations, sbm)
+    @test alloc == 0
+    return nothing
+end
+
+test_allocations(CSBM(; N=10^3, P=10^3, d=5, λ=2, μ=2, ρ=0.0))
+@test_skip test_allocations(
+    GLMSBM(; N=10^3, M=10^3, c=5, λ=2, ρ=0.0, Pʷ=GaussianWeightPrior())
+)