First plot

Author: gdalle

docs/plots.jl

@@ -2,7 +2,40 @@ using Base.Threads
 using CairoMakie
 using LinearAlgebra
 using Random: default_rng
+using Statistics
 using StochasticBlockModelVariants
 using ProgressMeter
 
 BLAS.set_num_threads(1)
+
+function compute_fig1(; N=3 * 10^3, P=N ÷ 10, d=5, μ=2, ρ=0.1, λ_values=0:0.1:2, trials=10)
+    rng = default_rng()
+    qᵤ_values = zeros(trials, length(λ_values))
+    qᵥ_values = zeros(trials, length(λ_values))
+    prog = Progress(trials * length(λ_values); desc="Fig 1 - trials")
+    @threads for i in 1:trials
+        for j in eachindex(λ_values)
+            λ = λ_values[j]
+            csbm = ContextualSBM(; N, P, d, μ, λ, ρ)
+            (; qᵤ, qᵥ) = evaluate_amp(rng; csbm)
+            qᵤ_values[i, j] = qᵤ
+            qᵥ_values[i, j] = qᵥ
+            next!(prog)
+        end
+    end
+    return (; λ_values, qᵤ_values, qᵥ_values)
+end
+
+function plot_fig1(res)
+    f = Figure()
+    ax = Axis(f[1, 1]; title="Fig 1", xlabel="λ", ylabel="qᵤ", limits=(0, 2, 0, 1))
+    qᵤ_means = dropdims(mean(res.qᵤ_values; dims=1); dims=1)
+    qᵤ_stds = dropdims(std(res.qᵤ_values; dims=1); dims=1)
+    lines!(ax, λ_values, qᵤ_means)
+    scatter!(ax, λ_values, qᵤ_means)
+    errorbars!(ax, λ_values, qᵤ_means, qᵤ_stds)
+    return f
+end
+
+res1 = compute_fig1()
+plot_fig1(res1)

src/StochasticBlockModelVariants.jl

@@ -11,7 +11,8 @@ using SparseArrays: SparseMatrixCSC, sparse, findnz
 
 export ContextualSBM, ContextualSBMLatents, ContextualSBMObservations
 export affinities, effective_snr
-export init_amp, update_amp!, run_amp, overlaps
+export overlaps
+export init_amp, update_amp!, run_amp, evaluate_amp
 
 include("utils.jl")
 include("csbm.jl")

src/csbm_inference.jl

@@ -175,11 +175,11 @@ function run_amp(
     rng::AbstractRNG;
     observations::ContextualSBMObservations{R},
     csbm::ContextualSBM{R},
-    init_std::Real=1e-3,
-    max_iterations::Integer=200,
+    init_std=1e-3,
+    max_iterations=200,
     convergence_threshold=1e-3,
     recent_past=10,
-    show_progress::Bool=false,
+    show_progress=false,
 ) where {R}
     (; N, P) = csbm
     (; marginals, next_marginals, storage) = init_amp(rng; observations, csbm, init_std)
@@ -222,3 +222,12 @@ function run_amp(
 
     return (; û_history, v̂_history, converged)
 end
+
+function evaluate_amp(rng::AbstractRNG; csbm::ContextualSBM, kwargs...)
+    (; latents, observations) = rand(rng, csbm)
+    (; û_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]
+    )
+    return (; qᵤ, qᵥ)
+end

test/csbm.jl

@@ -14,11 +14,11 @@ function test_recovery(csbm::ContextualSBM; test_u=true, test_v=true)
     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ᵤ >= first_q.qᵤ
         @test last_q.qᵤ > 0.5
     end
     if test_v
-        @test last_q.qᵥ > first_q.qᵥ
+        @test last_q.qᵥ >= first_q.qᵥ
         @test last_q.qᵥ > 0.5
     end
     return nothing