Basic structs

Author: gdalle

.gitignore

@@ -4,4 +4,4 @@
 # /Manifest.toml
 /docs/Manifest.toml
 /docs/build/
-/test/playground.jl
+playground.jl

CITATION.bib

@@ -1,5 +1,5 @@
 @misc{StochasticBlockModelVariants.jl,
-	author  = {Guillaume Dalle <22795598+gdalle@users.noreply.github.com> and contributors},
+	author  = {Guillaume Dalle},
 	title   = {StochasticBlockModelVariants.jl},
 	url     = {https://github.com/gdalle/StochasticBlockModelVariants.jl},
 	version = {v0.1.0},

src/StochasticBlockModelVariants.jl

@@ -14,6 +14,7 @@ export affinities, effective_snr
 export overlaps
 export init_amp, update_amp!, run_amp, evaluate_amp
 
+include("abstract_sbm.jl")
 include("utils.jl")
 include("csbm.jl")
 include("csbm_inference.jl")

src/abstract_sbm.jl

@@ -0,0 +1,33 @@
+"""
+    AbstractSBM
+
+Abstract supertype for Stochastic Block Models with additional node features.
+"""
+abstract type AbstractSBM end
+
+"""
+    average_degree(sbm)
+
+Return the average degre `d` of a node in the graph.
+"""
+function average_degree end
+
+"""
+    communities_snr(sbm)
+
+Return the signal-to-noise ratio `λ` of the communities in the graph.
+"""
+function communities_snr end
+
+"""
+    affinities(sbm)
+
+Return a named tuple `(; cᵢ, cₒ)` containing the affinities inside and outside of a community.
+"""
+function affinities(sbm::AbstractSBM)
+    d = average_degree(sbm)
+    λ = communities_snr(sbm)
+    cᵢ = d + λ * sqrt(d)
+    cₒ = d - λ * sqrt(d)
+    return (; cᵢ, cₒ)
+end

src/csbm.jl

@@ -3,7 +3,7 @@
 """
     ContextualSBM
 
-A generative model for graphs with node features, which combines the Stochastic Block Model with a mixture of Gaussians.
+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>
 
@@ -16,7 +16,7 @@ Reference: <https://arxiv.org/abs/2306.07948>
 - `μ`: SNR of the features
 - `ρ`: fraction of nodes revealed
 """
-struct ContextualSBM{R<:Real}
+struct ContextualSBM{R<:Real} <: AbstractSBM
     N::Int
     P::Int
     d::R
@@ -32,6 +32,9 @@ struct ContextualSBM{R<:Real}
     end
 end
 
+average_degree(csbm::CSBM) = csbm.d
+communities_snr(csbm::CSBM) = csbm.λ
+
 """
     effective_snr(csbm)
 
@@ -42,17 +45,6 @@ function effective_snr(csbm::ContextualSBM)
     return abs2(λ) + abs2(μ) / (N / P)
 end
 
-"""
-    affinities(csbm)
-
-Return a named tuple `(; cᵢ, cₒ)` containing the affinities inside and outside of a community.
-"""
-function affinities(csbm::ContextualSBM)
-    (; d, λ) = csbm
-    cᵢ = d + λ * sqrt(d)
-    cₒ = d - λ * sqrt(d)
-    return (; cᵢ, cₒ)
-end
 
 ## Latents
 

src/glmsbm.jl

@@ -0,0 +1,71 @@
+## Model
+
+"""
+    GLMSBM
+
+A generative model for graphs with node features, which combines a Generalized Linear Model with a Stochastic Block Model.
+
+Reference: <https://arxiv.org/abs/2303.09995>
+
+# Fields
+
+- `N`: graph size
+- `M`: feature dimension
+- `c`: average degree
+- `λ`: SNR of the communities
+- `ρ`: fraction of nodes revealed
+"""
+struct GLMSBM{R<:Real} <: AbstractSBM
+    N::Int
+    M::Int
+    c::R
+    λ::R
+    ρ::R
+
+    function GLMSBM(; N::Integer, M::Integer, c::R1, λ::R2, ρ::R3) where {R1,R2,R3}
+        R = promote_type(R1, R2, R3)
+        return new{R}(N, M, c, λ, ρ)
+    end
+end
+
+average_degree(glmsbm::GLMSBM) = glmsbm.c
+communities_snr(glmsbm::GLMSBM) = glmsbm.λ
+
+## Latents
+
+"""
+    GLMSBMLatents
+
+The hidden variables generated by sampling from a [`GLMSBM`](@ref).
+
+# Fields
+
+- `w::Vector`: feature weights, length `M`
+- `s::Vector`: community assignments, length `N`
+"""
+@kwdef struct GLMSBMLatents{R<:Real}
+    w::Vector{R}
+    s::Vector{Int}
+end
+
+## Observations
+
+"""
+    GLMSBMObservations
+
+The observations generated by sampling from a [`GLMSBM`](@ref).
+
+# Fields
+- `A::AbstractMatrix`: symmetric boolean adjacency matrix, size `(N, N)`
+- `g::AbstractGraph`: undirected unweighted graph generated from `A`
+- `F::Matrix`: feature matrix, size `(M, N)`
+- `Ξ::Vector`: revealed communities `±1` for a fraction of nodes and `0` for the rest, length `N`
+"""
+@kwdef struct ContextualSBMObservations{
+    R<:Real,M<:AbstractMatrix{Bool},G<:AbstractGraph{Int}
+}
+    A::M
+    g::G
+    F::Matrix{R}
+    Ξ::Vector{Int}
+end