# Clustering and Heterogeneity
#
# Accompanies the SISMID course "Network Modeling for Epidemics" (NME).
# Source page: mod5-Structure.qmd
#
# From timing to structure: whether contacts cluster into tight local groups
# (triangles), and whether contact patterns and durations depend on who people
# are. This is the most involved network we build by hand in these visualization
# tutorials, so we build it up one term at a time.
#
# Each render writes an HTML movie to your working directory. Run top to bottom.


# ---- Setup ----

library("EpiModel")
library("ndtv")

slice.par <- list(start = 1, end = 25, interval = 1, aggregate.dur = 1, rule = "any")
render.par <- list(tween.frames = 10, show.time = FALSE)
plot.par <- list(mar = c(0, 0, 0, 0))

# animate() wraps the pipeline; extra render.d3movie args (vertex.cex,
# vertex.sides, ...) pass through via ...
animate <- function(sim, file, ...) {
  nw <- get_network(sim)
  nw <- color_tea(nw, verbose = FALSE)
  compute.animation(nw, slice.par = slice.par, verbose = FALSE)
  render.d3movie(
    nw, render.par = render.par, plot.par = plot.par,
    vertex.col = "ndtvcol", edge.col = "darkgrey",
    vertex.border = "lightgrey", displaylabels = FALSE,
    filename = file.path(getwd(), file), ...)
}

# inf.prob = 1: every contact transmits with certainty (a deliberate extreme), so
# differences between the movies and counts come from network structure, not disease.
param <- param.net(inf.prob = 1)
init <- init.net(i.num = 10)
control <- control.net(type = "SI", nsteps = 25, nsims = 1)


# ---- Clustering: triangles ----

# gwesp (geometrically weighted edgewise shared partner) is the well-behaved
# stand-in for a raw, degenerate-prone triangle term. gwesp(0, TRUE) fixes the
# decay at 0, so the statistic is the number of edges with at least one shared
# partner (each closes a triangle). Target 15 of the 40 edges (an EXPECTED value;
# too high a target at fixed mean degree pushes back toward instability).
set.seed(1234)
nw <- network_initialize(n = 100)
coef.diss <- dissolution_coefs(dissolution = ~offset(edges), duration = 20)
est.tri <- netest(nw, ~edges + gwesp(0, TRUE), c(40, 15), coef.diss)
sim.tri <- netsim(est.tri, param, init, control)
animate(sim.tri, "movie-triangles.html", vertex.cex = 0.9)

# In these specifications triangles SLOW spread: they create redundant paths to
# people already reached, so each edge does less work. Boundary: "slower" is the
# direction here, not a law; it depends on the disease, the network, the seeding,
# and the outcome measured. Replicated evidence, 200 runs each:
final_prev <- function(est, nsims = 200) {
  sim <- netsim(est, param, init,
                control.net(type = "SI", nsteps = 25, nsims = nsims,
                            ncores = 5, verbose = FALSE))
  d <- as.data.frame(sim, out = "mean"); round(d$i.num[d$time == 25], 1)
}
set.seed(1234)
nw0 <- network_initialize(n = 100)
est.base <- netest(nw0, ~edges, 40, coef.diss, verbose = FALSE)
data.frame(
  specification = c("Baseline (edges only)", "With clustering (gwesp target 15)"),
  mean_infected_of_100 = c(final_prev(est.base), final_prev(est.tri)))


# ---- Heterogeneous networks: attributes ----

# Two ordinary (non-special) nodal attributes drive contact: binary "community"
# (any social/geographic/demographic split) and integer "age" (18-50).
set.seed(1234)
nw <- network_initialize(n = 100)
nw <- set_vertex_attribute(nw, "community", rbinom(100, 1, 0.5))
nw <- set_vertex_attribute(nw, "age", sample(18:50, 100, TRUE))

# Build the formation up one term at a time (each target is an expected value):
#   edges = 50                   overall mean degree 1 (a bit denser, so added
#                                structure is visible)
#   nodematch("community") = 40  40 of 50 edges within-community (structure)
#   nodefactor("community") = 70 total degree of community = 1 nodes (activity
#                                heterogeneity; default reference is the lowest
#                                value, 0)
#   absdiff("age") = 100         summed |age gap| over the 50 edges, so partners
#                                average 2 years apart (age assortativity)
#   concurrent = 30              slightly above the ~26 chance value at degree 1
formation <- ~edges + nodematch("community") + nodefactor("community") +
              absdiff("age") + concurrent
target.stats <- c(50, 40, 70, 100, 30)

# Homogeneous duration for now; heterogeneous dissolution comes next.
est.attr <- netest(nw, formation, target.stats,
                   dissolution_coefs(~offset(edges), duration = 20))
sim.attr <- netsim(est.attr, param, init, control)

# Encode community as node shape (square vs round) and age as node size.
nw.a <- get_network(sim.attr)
community <- get_vertex_attribute(nw.a, "community")
age <- get_vertex_attribute(nw.a, "age")
animate(sim.attr, "movie-attributes.html",
        vertex.sides = ifelse(community == 1, 4, 50),
        vertex.cex = age / 25)


# ---- Heterogeneous dissolution ----

# Let duration depend on partnership type: within-community ties average 10
# steps, cross-community 20. The dissolution formula must be dyad-independent.
# ORDER IS THE TRAP: the first duration is the REFERENCE group (edges NOT matched
# on community), the second is the matched group.
coef.diss.het <- dissolution_coefs(
  dissolution = ~offset(edges) + offset(nodematch("community")),
  duration = c(20, 10))

# The dissolution model must be a LEADING SUBSET of the formation model: its
# terms (edges, nodematch("community")) are the FIRST TWO formation terms, in the
# same order. dissolution_coefs computes one correction per dissolution term
# (coef.form.corr), and netest subtracts each from the positionally corresponding
# formation coefficient. (Current EpiModel convention; confirm ?dissolution_coefs
# on an unfamiliar version.)
#
# Reading it back is hard: the stored coefficients are on the logit scale of
# PERSISTENCE (p = 1 - 1/duration). The first, logit(1 - 1/20) = 2.944, is the
# reference and matches the homogeneous duration = 20 model. The second is NOT
# the 10-step coefficient; it is the INCREMENT: logit(1 - 1/10) - 2.944 =
# 2.197 - 2.944 = -0.747. So the sign is direction (negative = shorter than
# reference), a coefficient is meaningless without its reference, and equal
# durations give an increment of exactly 0 (the homogeneous model as a special
# case). Dissolution heterogeneity must be a property of the EDGE (nodematch,
# nodemix), not the node (nodefactor is unsupported here: an edge has two nodes).

est.het <- netest(nw, formation, target.stats, coef.diss.het)

# Check the durations came out where we asked: the ONLY way to see the
# heterogeneity rather than assert it. The edges row is the reference duration,
# nodematch.community the matched one; they land near 20 and 10.
dx <- netdx(est.het, nsims = 5, nsteps = 500, ncores = 5, verbose = FALSE)
dx

# NOTE: the edges dissolution approximation corrects the formation coefficients
# for EACH dissolution term, here BOTH edges AND nodematch("community"), not just
# edges. Compare est.het$coef.form.crude (cross-sectional) with est.het$coef.form
# (corrected): the first two entries differ, the rest are untouched.
round(est.het$coef.form.crude - est.het$coef.form, 4)
# That coefficient correction is exact for these dyad-independent terms. But in
# dynamic simulation every formation statistic still lands a few percent under
# target (see dx above: all 5 to 8% low) because the edges approximation is only
# first-order; the dyad-DEPENDENT concurrent term drifts the most (~8% here), and
# that bias grows as durations shorten.

plot(dx, type = "duration")

# The movie, with community and age encoded as before.
sim.het <- netsim(est.het, param, init, control)
nw.h <- get_network(sim.het)
community <- get_vertex_attribute(nw.h, "community")
age <- get_vertex_attribute(nw.h, "age")
animate(sim.het, "movie-heterodissolution.html",
        vertex.sides = ifelse(community == 1, 4, 50),
        vertex.cex = age / 25)