Finish recoding GLMSBM

Author: gdalle

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0-beta2"
 manifest_format = "2.0"
-project_hash = "5d6b46b606378213b4547f42572ed5ae33cb9d37"
+project_hash = "ba25bdf8dd9a1a3a609260b65801f3325dbdaa20"
 
 [[deps.ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
@@ -22,9 +22,9 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
 [[deps.Compat]]
 deps = ["UUIDs"]
-git-tree-sha1 = "4e88377ae7ebeaf29a047aa1ee40826e0b708a5d"
+git-tree-sha1 = "e460f044ca8b99be31d35fe54fc33a5c33dd8ed7"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "4.7.0"
+version = "4.9.0"
 weakdeps = ["Dates", "LinearAlgebra"]
 
     [deps.Compat.extensions]
@@ -37,9 +37,9 @@ version = "1.0.5+1"
 
 [[deps.DataStructures]]
 deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
-git-tree-sha1 = "cf25ccb972fec4e4817764d01c82386ae94f77b4"
+git-tree-sha1 = "3dbd312d370723b6bb43ba9d02fc36abade4518d"
 uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-version = "0.18.14"
+version = "0.18.15"
 
 [[deps.Dates]]
 deps = ["Printf"]
@@ -75,6 +75,12 @@ 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"
@@ -143,14 +149,20 @@ 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"
 
 [[deps.MacroTools]]
 deps = ["Markdown", "Random"]
-git-tree-sha1 = "42324d08725e200c23d4dfb549e0d5d89dede2d2"
+git-tree-sha1 = "9ee1618cbf5240e6d4e0371d6f24065083f60c48"
 uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
-version = "0.5.10"
+version = "0.5.11"
 
 [[deps.Markdown]]
 deps = ["Base64"]
@@ -189,9 +201,9 @@ uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
 version = "0.5.5+0"
 
 [[deps.OrderedCollections]]
-git-tree-sha1 = "d321bf2de576bf25ec4d3e4360faca399afca282"
+git-tree-sha1 = "2e73fe17cac3c62ad1aebe70d44c963c3cfdc3e3"
 uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-version = "1.6.0"
+version = "1.6.2"
 
 [[deps.Pkg]]
 deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
@@ -200,9 +212,9 @@ version = "1.10.0"
 
 [[deps.PrecompileTools]]
 deps = ["Preferences"]
-git-tree-sha1 = "9673d39decc5feece56ef3940e5dafba15ba0f81"
+git-tree-sha1 = "03b4c25b43cb84cee5c90aa9b5ea0a78fd848d2f"
 uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
-version = "1.1.2"
+version = "1.2.0"
 
 [[deps.Preferences]]
 deps = ["TOML"]
@@ -216,9 +228,9 @@ uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
 [[deps.ProgressMeter]]
 deps = ["Distributed", "Printf"]
-git-tree-sha1 = "d7a7aef8f8f2d537104f170139553b14dfe39fe9"
+git-tree-sha1 = "ae36206463b2395804f2787ffe172f44452b538d"
 uuid = "92933f4c-e287-5a05-a399-4b506db050ca"
-version = "1.7.2"
+version = "1.8.0"
 
 [[deps.REPL]]
 deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
@@ -282,9 +294,9 @@ weakdeps = ["Statistics"]
     StaticArraysStatisticsExt = "Statistics"
 
 [[deps.StaticArraysCore]]
-git-tree-sha1 = "6b7ba252635a5eff6a0b0664a41ee140a1c9e72a"
+git-tree-sha1 = "36b3d696ce6366023a0ea192b4cd442268995a0d"
 uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
-version = "1.4.0"
+version = "1.4.2"
 
 [[deps.Statistics]]
 deps = ["LinearAlgebra", "SparseArrays"]

Project.toml

@@ -5,8 +5,11 @@ version = "0.1.0"
 
 [deps]
 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"

src/StochasticBlockModelVariants.jl

@@ -1,7 +1,18 @@
+"""
+    StochasticBlockModelVariants
+
+A package for inference in SBMs with node features using message-passing algorithms.
+
+# Exports
+
+$(EXPORTS)
+"""
 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!
@@ -18,7 +29,6 @@ export overlaps
 export init_amp, update_amp!, run_amp, evaluate_amp
 
 include("abstract_sbm.jl")
-include("utils.jl")
 include("csbm.jl")
 include("glmsbm.jl")
 include("csbm_inference.jl")

src/abstract_sbm.jl

@@ -1,33 +1,33 @@
 """
-    AbstractSBM
+$(TYPEDEF)
 
 Abstract supertype for Stochastic Block Models with additional node features.
 """
 abstract type AbstractSBM end
 
 """
-    length(sbm)
+    length(sbm::AbstractSBM)
 
 Return the number of nodes `N` in the graph.
 """
 Base.length
 
 """
-    nb_features(sbm)
+    nb_features(sbm::AbstractSBM)
 
-Return the number of features for each node in the graph.
+Return the number of nodes `N` in the graph.
 """
 function nb_features end
 
 """
-    average_degree(sbm)
+    average_degree(sbm::AbstractSBM)
 
 Return the average degre `d` of a node in the graph.
 """
 function average_degree end
 
 """
-    communities_snr(sbm)
+    communities_snr(sbm::AbstractSBM)
 
 Return the signal-to-noise ratio `λ` of the communities in the graph.
 """
@@ -46,19 +46,6 @@ 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)
 
@@ -112,3 +99,19 @@ function sample_mask(rng::AbstractRNG, sbm::AbstractSBM, communities::Vector{<:I
 end
 
 prior(::Type{R}, u, Ξᵢ) where {R} = ismissing(Ξᵢ) ? one(R) / 2 : R(Ξᵢ == u)
+
+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))
+    qᵤ = 2 * (q̂ᵤ - one(R) / 2)
+    return qᵤ
+end
+
+function continuous_overlap(v, v̂)
+    q̂ᵥ = max(abs(dot(v̂, v)), abs(dot(v̂, -v)))
+    qᵥ = q̂ᵥ / (eps() + norm(v̂) * norm(v))
+    return qᵥ
+end

src/csbm.jl

@@ -1,27 +1,28 @@
 ## Model
 
 """
-    CSBM
+$(TYPEDEF)
 
 A generative model for graphs with node features, which combines a Stochastic Block Model with a mixture of Gaussians.
 
 Reference: <https://arxiv.org/abs/2306.07948>
 
 # Fields
 
-- `N`: graph size
-- `P`: feature dimension
-- `d`: average degree
-- `λ`: SNR of the communities
-- `μ`: SNR of the features
-- `ρ`: fraction of node assignments observed
+$(TYPEDFIELDS)
 """
-struct CSBM{R<:Real} <: AbstractSBM
+struct CSBM{R} <: AbstractSBM
+    "graph size"
     N::Int
+    "feature dimension"
     P::Int
+    "average degree"
     d::R
+    "SNR of the communities"
     λ::R
+    "SNR of the features"
     μ::R
+    "fraction of node assignments observed"
     ρ::R
 
     function CSBM(; N::Integer, P::Integer, d::R1, λ::R2, μ::R3, ρ::R4) where {R1,R2,R3,R4}
@@ -50,35 +51,38 @@ end
 ## Latents
 
 """
-    LatentsCSBM
+$(TYPEDEF)
 
 The hidden variables generated by sampling from a [`CSBM`](@ref).
 
 # Fields
 
-- `u::Vector`: community assignments, length `N`
-- `v::Vector`: feature centroids, length `P`
+$(TYPEDFIELDS)
 """
 @kwdef struct LatentsCSBM{R<:Real}
+    "community assignments, length `N`"
     u::Vector{Int}
+    "feature centroids, length `P`"
     v::Vector{R}
 end
 
 ## Observations
 
 """
-    ObservationsCSBM
+$(TYPEDEF)
 
 The observations generated by sampling from a [`CSBM`](@ref).
 
 # Fields
-- `g::AbstractGraph`: undirected unweighted graph with `N` nodes (~ adjacency matrix `A`)
-- `Ξ::Vector`: revealed communities `±1` for a fraction `ρ` of nodes and `missing` for the rest, length `N`
-- `B::Matrix`: feature matrix, size `(P, N)`
+
+$(TYPEDFIELDS)
 """
 @kwdef struct ObservationsCSBM{R<:Real,G<:AbstractGraph{Int}}
+    "undirected unweighted graph with `N` nodes (~ adjacency matrix `A`)"
     g::G
+    "revealed communities `±1` for a fraction `ρ` of nodes and `missing` for the rest, length `N`"
     Ξ::Vector{Union{Int,Missing}}
+    "feature matrix, size `(P, N)`"
     B::Matrix{R}
 end