More testing

Author: gdalle

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.9.2"
 manifest_format = "2.0"
-project_hash = "d06132272ceec5b60df1cd95ecf68b1ee04ca385"
+project_hash = "76c4ad29b7fdf61caa5b4fe006d178c04b45b443"
 
 [[deps.ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
@@ -142,16 +142,22 @@ deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "
 uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 version = "1.9.2"
 
+[[deps.PrecompileTools]]
+deps = ["Preferences"]
+git-tree-sha1 = "9673d39decc5feece56ef3940e5dafba15ba0f81"
+uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+version = "1.1.2"
+
+[[deps.Preferences]]
+deps = ["TOML"]
+git-tree-sha1 = "7eb1686b4f04b82f96ed7a4ea5890a4f0c7a09f1"
+uuid = "21216c6a-2e73-6563-6e65-726566657250"
+version = "1.4.0"
+
 [[deps.Printf]]
 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

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

src/StochasticBlockModelVariants.jl

@@ -2,7 +2,7 @@ module StochasticBlockModelVariants
 
 using Graphs: AbstractGraph, neighbors
 using LinearAlgebra: dot, mul!
-using ProgressMeter: @showprogress
+using PrecompileTools: @compile_workload
 using Random: AbstractRNG, default_rng
 using SimpleWeightedGraphs: SimpleWeightedGraph
 using Statistics: mean
@@ -16,4 +16,11 @@ 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)
+    (; observations) = rand(rng, csbm)
+    run_amp(rng; observations, csbm, iterations=2)
+end
+
 end

src/csbm_inference.jl

@@ -123,26 +123,27 @@ function update_amp!(
     h̃₊ .= -h₊ .+ log(one(R) / 2) .+ û_no_feat
     h̃₋ .= -h₋ .+ log(one(R) / 2) .- û_no_feat
 
-    # 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)
-            if k != j
-                common = 2λ * sqrt(d) * χ₊eᵗ[k, i]
-                s_ij += log((cₒ + common) / (cᵢ - common))
-            end
-        end
-        χ₊eᵗ⁺¹[i, j] = sigmoid(s_ij)
-    end
-
     # BP update of the marginals
     for i in 1:N
         s_i = h̃₊[i] - h̃₋[i]
         for k in neighbors(g, i)
             common = 2λ * sqrt(d) * χ₊eᵗ[k, i]
             s_i += log((cₒ + common) / (cᵢ - common))
         end
-        χ₊[i] = sigmoid(s_i)
+        χ₊[i] = s_i
+    end
+
+    # BP update of the messages
+    for i in 1:N, j in neighbors(g, i)
+        common = 2λ * sqrt(d) * χ₊eᵗ[j, i]
+        s_ij = log((cₒ + common) / (cᵢ - common))
+        χ₊eᵗ⁺¹[i, j] = χ₊[i] - s_ij
+    end
+
+    # Sigmoidize probabilities
+    χ₊ .= sigmoid.(χ₊)
+    for (key, val) in pairs(χ₊eᵗ⁺¹)
+        χ₊eᵗ⁺¹[key] = sigmoid(val)
     end
 
     # BP estimation of u
@@ -155,17 +156,15 @@ function run_amp(
     rng::AbstractRNG;
     observations::ContextualSBMObservations,
     csbm::ContextualSBM,
-    init_std::Real,
-    iterations::Integer,
+    init_std::Real=1e-3,
+    iterations::Integer=10,
 )
     (; 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
+    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

test/csbm.jl

@@ -1,21 +1,49 @@
+using JET
 using Random: default_rng
 using Statistics
 using StochasticBlockModelVariants
 using Test
 
 rng = default_rng()
 
-function test_recovery(; N, P, d, λ, μ, ρ, init_std=1e-3, iterations=20)
-    csbm = ContextualSBM(; N, P, d, λ, μ, ρ)
+function test_recovery(csbm::ContextualSBM)
     @assert effective_snr(csbm) > 1
     (; latents, observations) = rand(rng, csbm)
-    storage_history = run_amp(rng; observations, csbm, init_std, iterations)
+    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)
     @test last(overlap_history) > 0.5
     return nothing
 end
 
-test_recovery(; N=10^3, P=10^3, d=5, λ=2, μ=2, ρ=0.0)  # AMP-BP
-test_recovery(; N=10^3, P=10^3, d=5, λ=0, μ=2, ρ=0.0)  # AMP
-test_recovery(; N=10^3, P=10^3, d=5, λ=2, μ=0, ρ=0.0)  # BP
+function test_jet(csbm::ContextualSBM)
+    (; observations) = rand(rng, csbm)
+    @test_opt target_modules = (StochasticBlockModelVariants,) run_amp(
+        rng; observations, csbm, iterations=2
+    )
+    @test_call target_modules = (StochasticBlockModelVariants,) run_amp(
+        rng; observations, csbm, iterations=2
+    )
+    return nothing
+end
+
+function test_allocations(csbm::ContextualSBM)
+    (; observations) = rand(rng, csbm)
+    (; storage, next_storage, temp_storage) = init_amp(
+        rng; observations, csbm, init_std=1e-3
+    )
+    alloc = @allocated update_amp!(next_storage, temp_storage; storage, observations, csbm)
+    @test alloc == 0
+    return nothing
+end
+
+@testset "Correct code" begin
+    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
+end
+
+@testset "Good code" begin
+    test_jet(ContextualSBM(; N=10^3, P=10^3, d=5, λ=2, μ=2, ρ=0.0))
+    test_allocations(ContextualSBM(; N=10^3, P=10^3, d=5, λ=2, μ=2, ρ=0.0))
+end