ERG boilerplate

erg_boilerplate.jl

using Distributed, Statistics, GraphRecipes, BenchmarkTools, LinearAlgebra, LightGraphs, Plots, RCall, ProfileView, FunctionWrappers, LoopVectorization, ThreadPools

using FunctionWrappers: FunctionWrapper
using Profile

n = 100
g_orig = erdos_renyi(n, 0.05; is_directed = true)
g_adj = convert(Array{Bool, 2}, collect(adjacency_matrix(g_orig)));

struct erg{G <: AbstractArray{Bool}} <: AbstractGraph{Int}
    m::G
    fs::Array{FunctionWrapper{Float64,Tuple{erg,Int64,Int64}},1}
    trans::G
    realnodecov::Union{Dict{String,Array{Float64,1}},Nothing}
    catnodecov::Union{Dict{String,Array{Any,1}},Nothing}
    edgecov::Union{Array{Array{Float64,2}},Nothing}
    indegree::Vector{Int64}
    outdegree::Vector{Int64}
    n_funcs::Int
    kstars_precomputed::Array{Float64,2}
end

Signature = FunctionWrapper{
    Float64,
    Tuple{erg, Int, Int}
}

# Basic outer constructor when we just have a graph and set of subgraph functions
function erg(
    g::G,
    fs::Array{FunctionWrapper{Float64,Tuple{erg,Int64,Int64}},1};
    max_star=3,
    realnodecov=nothing,
    catnodecov=nothing,
    edgecov=nothing,
) where {G <: AbstractArray{Bool}}
    p = length(fs)
    n = size(g)[1]

    erg(
        g,
        fs,
        collect(transpose(g)),
        realnodecov,
        catnodecov,
        edgecov,
        vec(sum(g, dims=1)),
        vec(sum(g, dims=2)),
        p,
        kstarpre(n, max_star),
    )
end

# Define some needed helper functions
is_directed(g::E) where {E <: erg} = true
edgetype(g::E) where {E <: erg} = SimpleEdge{Int}
ne(g::erg) = sum(g.m)
nv(g::erg) = @inbounds size(g.m)[1]
vertices(g::E) where {E <: erg} = 1:nv(g)
has_vertex(g::E, v) where {E <: erg} = v <= nv(g) && v > 0
indegree(g::E, v::T) where {E <: erg,T <: Int} = @inbounds g.indegree[v]
outdegree(g::E, v::T) where {E <: erg,T <: Int} = @inbounds g.outdegree[v]
density(x) = ne(x) / (nv(x) * (nv(x) - 1))

function has_edge(g::E, s, r) where {E <: erg}
    g.m[s, r]
end

function outneighbors(g::E, node) where {E <: erg}
    @inbounds return (v for v in 1:nv(g) if g.trans[v, node] && node != v)
end

function inneighbors(g::E, node) where {E <: erg}
    @inbounds return (v for v in 1:nv(g) if g.m[v, node] && v != node)
end


# Functions to turn edge on and off - does no checking
function add_edge!(g::E, s, r) where {E <: erg}
    g.m[s, r] = true
    g.trans[r, s] = true
    g.indegree[r] += 1
    g.outdegree[s] += 1
    return nothing
end

function rem_edge!(g::E, s, r) where {E <: erg}
    g.m[s, r] = false
    g.trans[r, s] = false
    g.indegree[r] -= 1
    g.outdegree[s] -= 1
    return nothing
end


# Toggle a single edge in-place
function edge_toggle!(g::E, s, r) where {E <: erg}
    if has_edge(g, s, r)
        rem_edge!(g, s, r)
        return nothing
    else
        add_edge!(g, s, r)
        return nothing
    end
end

# Toggle and edge and get the changes in total number of edges; really basic, but useful.
function delta_edge(g::E, s, r) where {E <: AbstractGraph}
    if !has_edge(g, s, r)
        return 1.0
    else
        return -1.0
    end
end

# Toggle edge from s -> r, and get changes in count of reciprocal dyads
function delta_mutual(g::E, s, r) where {E <: AbstractGraph}
    if !g.trans[s, r]
        return 0.0
    elseif !g.m[s, r]
        return 1.0
    else
        return -1.0
    end
end

# delta k-stars using a precalculated table of binomials - MUCH faster
# to maintain compatibility with both erg and simplegraphs, explicitly pass the required degree integer
function delta_istar(g::E, s, r, k) where {E <: AbstractGraph}
    if !has_edge(g, s, r)
        return g.kstars_precomputed[indegree(g, r) + 1, k - 1]
    else
        return -g.kstars_precomputed[indegree(g, r) - 1 + 1, k - 1]
    end
end


function delta_ostar(g::E, s, r, k) where {E <: AbstractGraph}
    if !has_edge(g, s, r)
        return g.kstars_precomputed[outdegree(g, s) + 1, k - 1]
    else
        return -g.kstars_precomputed[outdegree(g, s) - 1 + 1, k - 1]
    end
end


# Change in mixed 2-stars
function delta_m2star(g::G, s, r) where {G <: AbstractGraph}
    if !has_edge(g, s, r) && !has_edge(g, r, s)
        return convert(Float64, indegree(g, s) + outdegree(g, r))
    elseif has_edge(g, s, r) && !has_edge(g, r, s)
        return -convert(Float64, indegree(g, s) + outdegree(g, r))
    elseif !has_edge(g, s, r) && has_edge(g, r, s)
        return convert(Float64, indegree(g, s) + outdegree(g, r) - 2)
    else
        return convert(Float64, -indegree(g, s) - outdegree(g, r) + 2)
    end
end

# Change in transitive triads
# Works very well for arrays, horribly for edgelist
# Note: must use + and * operations to get SIMD - using && is much slower
function delta_ttriple(g::E, s, r) where {E <: AbstractGraph}
    x = 0
    for i in vertices(g)
        x +=
            (g.trans[i, r] * g.trans[i, s]) +
            (g.m[i, r] * g.trans[i, s]) +
            (g.m[i, r] * g.m[i, s])
    end
    if !has_edge(g, s, r)
        return convert(Float64, x)
    else
        return -convert(Float64, x)
    end
end

# Effect of covariate on indegrees
function delta_nodeicov(g::G, s, r, cov_name)::Float64 where {G <: AbstractGraph}
    if !has_edge(g, s, r)
        return g.realnodecov[cov_name][r]
    else
        return -g.realnodecov[cov_name][r]
    end
end


# Effect of covariate on outdegrees
function delta_nodeocov(g::G, s, r, cov_name)::Float64 where {G <: AbstractGraph}
    if !has_edge(g, s, r)
        return g.realnodecov[cov_name][s]
    else
        return -g.realnodecov[cov_name][s]
    end
end

# Effect of dyadic distance on covariate p
function delta_nodediff(g::G, s, r, k, cov_name)::Float64 where {G <: AbstractGraph}
    if !has_edge(g, s, r)
        return abs(g.realnodecov[cov_name][s] - g.realnodecov[cov_name][r])^k
    else
        return -abs(g.realnodecov[cov_name][s] - g.realnodecov[cov_name][r])^k
    end
end

function kstarpre(n::Int, k::Int)
    m = zeros(n, k)
    for j = 1:k
        for i = 1:n
            m[i, j] = convert(Float64, binomial(i - 1, j))
        end
    end
    return m
end