Added Personalized PageRank Diffusion [ppr_diffusion function]

src/GNNGraphs/GNNGraphs.jl

@@ -78,6 +78,7 @@ export add_nodes,
        to_unidirected,
        random_walk_pe,
        remove_nodes,
+       ppr_diffusion,
 # from Flux
        batch,
        unbatch,

src/GNNGraphs/transform.jl

@@ -1078,3 +1078,53 @@ ci2t(ci::AbstractVector{<:CartesianIndex}, dims) = ntuple(i -> map(x -> x[i], ci
 @non_differentiable remove_self_loops(x...)  # TODO this is wrong, since g carries feature arrays, needs rrule
 @non_differentiable dense_zeros_like(x...)
 
+"""
+    ppr_diffusion(g::GNNGraph{<:COO_T}; alpha_f32::Float32=0.85f0) -> GNNGraph
+
+Calculates the Personalized PageRank (PPR) diffusion based on the edge weight matrix of a GNNGraph and updates the graph with new edge weights derived from the PPR matrix.
+
+The function performs the following steps:
+1. Constructs a modified adjacency matrix `A` using the graph's edge weights, where `A` is adjusted by `(α - 1) * A + I`, with `α` being the damping factor (`alpha_f32`) and `I` the identity matrix.
+2. Normalizes `A` to ensure each column sums to 1, representing transition probabilities.
+3. Applies the PPR formula `α * (I + (α - 1) * A)^-1` to compute the diffusion matrix.
+4. Updates the original edge weights of the graph based on the PPR diffusion matrix, assigning new weights for each edge from the PPR matrix.
+
+# Arguments
+- `g::GNNGraph`: The input graph for which PPR diffusion is to be calculated. It should have edge weights available.
+- `alpha_f32::Float32`: The damping factor used in PPR calculation, controlling the teleport probability in the random walk. Defaults to `0.85f0`.
+
+# Returns
+- A new `GNNGraph` instance with the same structure as `g` but with updated edge weights according to the PPR diffusion calculation.
+"""
+function ppr_diffusion(g::GNNGraph{<:COO_T}; alpha_f32::Float32=0.85f0)
+    s, t = edge_index(g)
+    w = get_edge_weight(g)  
+
+    N = g.num_nodes
+
+    A = zeros(Float32, N, N)
+
+    for (idx, src, dst) in zip(1:length(s), s, t)
+        A[dst, src] += w[idx]  
+    end
+
+    A .*= (alpha_f32 - 1)
+
+    for i in 1:N
+        A[i, i] += 1
+    end
+
+    # α * (I + (α - 1) * A)^-1
+    PPR = alpha_f32 * inv(A)
+    println(PPR)
+
+    new_w = Float32[]
+    for (src, dst) in zip(s, t)
+        push!(new_w, PPR[dst, src]) 
+    end
+
+    return GNNGraph((s, t, new_w),
+                    g.num_nodes, length(s), g.num_graphs,
+                    g.graph_indicator,
+                    g.ndata, g.edata, g.gdata)
+end

test/GNNGraphs/transform.jl

@@ -519,4 +519,25 @@ end
 
         @test g.graph[(:A, :to1, :A)][3] == vcat([2, 2, 2], fill(1, n))
     end
+end
+
+@testset "ppr_diffusion" begin
+    if GRAPH_T == :coo
+        s = [1, 1, 2, 3]
+        t = [2, 3, 4, 5]
+        eweights = [0.1, 0.2, 0.3, 0.4]
+
+        g = GNNGraph(s,t,eweights)
+        get_edge_weight(g)
+
+        g_new = ppr_diffusion(g)
+        w_new = get_edge_weight(g_new)
+
+        check_ew = Float32[0.012749999
+        0.025499998
+        0.038249996
+        0.050999995]
+
+        isequal(w_new, check_ew)
+    end
 end