add LSPE

Author: yuehhua

src/GeometricFlux.jl

@@ -82,6 +82,7 @@ include("models.jl")
 
 include("sampling.jl")
 include("embedding/node2vec.jl")
+include("layers/positional.jl")
 
 using .Datasets
 

src/layers/positional.jl

@@ -0,0 +1,91 @@
+"""
+    AbstractPE
+
+Abstract type of positional encoding for GNN.
+"""
+abstract type AbstractPE end
+
+positional_encode(l::AbstractPE, args...) = throw(ErrorException("positional_encode function for $l is not implemented."))
+
+"""
+    LSPE(mp)
+
+Learnable structural positional encoding layer.
+
+# Arguments
+
+- `mp`: message-passing layer.
+"""
+struct LSPE{A<:MessagePassing,F,P} <: AbstractPE
+    layer::A
+    f_p::F
+    pe::P
+end
+
+function LSPE(mp::MessagePassing, pe_dim::Int; init=glorot_uniform, init_pe=random_walk_pe)
+    # f_p = Dense(; init=init)
+    pe = init_pe(A, pe_dim)
+    return LSPE(mp, f_p, pe)
+end
+
+positional_encode(::LSPE, p_i, p_j, e_ij) = p_j
+
+function message(l::LSPE, h_i, h_j, e_ij, p_i, p_j)
+    x_i = isnothing(h_i) ? nothing : vcat(h_i, p_i)
+    x_j = isnothing(h_j) ? nothing : vcat(h_j, p_j)
+    return message(l.layer, x_i, x_j, e_ij)
+end
+
+function Base.show(io::IO, l::LSPE)
+    print(io, "LSPE(", l.layer, ")")
+end
+
+
+"""
+    random_walk_pe(A, k)
+
+Returns positional encoding (PE) of size `(k, N)` where N is node number.
+PE is generated by `k`-step random walk over given graph.
+
+# Arguments
+
+- `A`: Adjacency matrix of a graph.
+- `k::Int`: First dimension of PE.
+"""
+function random_walk_pe(A::AbstractMatrix, k::Int)
+    N = size(A, 1)
+    @assert k ≤ N "k must less or equal to number of nodes"
+    inv_D = GraphSignals.degree_matrix(A, Float32, inverse=true)
+
+    RW = similar(A, size(A)..., k)
+    RW[:, :, 1] .= A * inv_D
+    for i in 2:k
+        RW[:, :, i] .= RW[:, :, i-1] * RW[:, :, 1]
+    end
+
+    pe = similar(RW, k, N)
+    for i in 1:N
+        pe[:, i] .= RW[i, i, :]
+    end
+
+    return pe
+end
+
+"""
+    laplacian_pe(A, k)
+
+Returns positional encoding (PE) of size `(k, N)` where `N` is node number.
+PE is generated from eigenvectors of a graph Laplacian truncated by `k`.
+
+# Arguments
+
+- `A`: Adjacency matrix of a graph.
+- `k::Int`: First dimension of PE.
+"""
+function laplacian_pe(A::AbstractMatrix, k::Int)
+    N = size(A, 1)
+    @assert k ≤ N "k must less or equal to number of nodes"
+    L = GraphSignals.normalized_laplacian(A)
+    U = eigvecs(L)
+    return U[1:k, :]
+end