# Module 5 Lab: Networks in Motion (Worked Solutions)
#
# Accompanies the SISMID course "Network Modeling for Epidemics" (NME).
# Source page: mod5-Lab.qmd
#
# The lab is a menu of paths, not a checklist. Each path below changes ONE thing
# about a network from the tutorials; the habit to build is to PREDICT what the
# movie should do before running it, and where a path asks "how much," back the
# prediction with replicated runs rather than a single movie.
#
# Try a path yourself before reading its solution. Run top to bottom, or jump to
# the path you want.


# ---- Setup ----

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

# Shared epidemic (the poker-chip SI process) and animation helper.
param <- param.net(inf.prob = 1)
init <- init.net(i.num = 10)
control <- control.net(type = "SI", nsteps = 25, nsims = 1)

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 <- 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.cex = 0.9, vertex.col = "ndtvcol", edge.col = "darkgrey",
                 vertex.border = "lightgrey", displaylabels = FALSE,
                 filename = file.path(getwd(), file))
}
coef.diss <- dissolution_coefs(dissolution = ~offset(edges), duration = 20)


# =====================================================================
# PATH: If concurrency interests you
# =====================================================================

# The other direction from the tutorial's concurrent = 0. 30 is well ABOVE the
# ~19 that chance (the Poisson benchmark) would give, so we ask for MORE
# concurrency than random mixing. Prediction: faster spread than no-concurrency.
set.seed(1234)
nw <- network_initialize(n = 100)
est.hi <- netest(nw, ~edges + concurrent, c(40, 30), coef.diss)
sim.hi <- netsim(est.hi, param, init, control)
animate(sim.hi, "lab-concurrent30.html")

# Where the model breaks: push the target to 35, 40, 45. At 40 edges there are 80
# degree slots and each concurrent node needs 2, so at most 40 nodes can be
# concurrent. As the target nears and passes that ceiling, netest first MISSES
# concurrent (fitting edges), then fails outright. tryCatch keeps the loop going.
for (ct in c(35, 40, 45)) {
  set.seed(1234)
  nw <- network_initialize(n = 100)
  fit <- tryCatch(
    suppressWarnings(netest(nw, ~edges + concurrent, c(40, ct), coef.diss,
                            verbose = FALSE)),
    error = function(e) NULL)
  cat("\n--- concurrent target =", ct, "---\n")
  if (is.null(fit)) {
    cat("netest could not fit this target (model too far from feasible).\n")
  } else {
    print(netdx(fit, nsims = 10, nsteps = 100, ncores = 5, verbose = FALSE))
  }
}
# Lesson: a mean degree can only support so much of any structure. Ask for too
# much and you get a miss you must catch with netdx; ask for the impossible and
# the fit fails. You do not learn this from the coefficients, so netdx is not
# optional.


# =====================================================================
# PATH: If timing interests you
# =====================================================================

# Fill in the middle of the duration dial. "Long" and "short" are relative to the
# 25-step horizon: at 5 the network churns (each edge turns over several times);
# at 50 it feels nearly frozen over 25 steps. Same mean degree, opposite feel, so
# much of what you are watching is timing, not structure.
for (d in c(5, 50)) {
  set.seed(1234)
  nw <- network_initialize(n = 100)
  est <- netest(nw, ~edges, 40, dissolution_coefs(~offset(edges), duration = d))
  sim <- netsim(est, param, init, control)
  animate(sim, paste0("lab-duration", d, ".html"))
}


# =====================================================================
# PATH: If heterogeneous dissolution interests you
# =====================================================================

# A minimal edges + nodematch("community") model is enough, and it fits robustly
# (dyad-independent, so MPLE, no MCMC). The FIRST duration is the reference
# (non-matched) group, the SECOND is matched.
set.seed(1234)
nw <- network_initialize(n = 100)
nw <- set_vertex_attribute(nw, "community", rbinom(100, 1, 0.5))
form.d <- ~edges + nodematch("community")
targ.d <- c(50, 40)

# Baseline: reference (non-matched) ties last 20, matched ties 10.
est.base.d <- netest(nw, form.d, targ.d,
  dissolution_coefs(~offset(edges) + offset(nodematch("community")),
                    duration = c(20, 10)))
netdx(est.base.d, nsims = 5, nsteps = 500, ncores = 5, verbose = FALSE)

# Swap: matched (within-community) ties become the LONGER ones. The two Duration
# rows exchange values, and the nodematch.community coefficient FLIPS SIGN
# (matched now longer than reference, so the increment turns positive). Nothing
# but the diagnostics tells you a valid-but-backwards model is wrong.
est.swapped <- netest(nw, form.d, targ.d,
  dissolution_coefs(~offset(edges) + offset(nodematch("community")),
                    duration = c(10, 20)))
netdx(est.swapped, nsims = 5, nsteps = 500, ncores = 5, verbose = FALSE)

# Equal durations make the increment exactly 0: the nodematch.community
# dissolution coefficient is 0, and you have fit the HOMOGENEOUS model (one
# duration for every edge) with a redundant second dissolution term.
dissolution_coefs(~offset(edges) + offset(nodematch("community")),
                  duration = c(20, 20))$coef.adj


# =====================================================================
# PATH: If you want to measure, not just watch
# =====================================================================

# Replace one movie with 200 runs. The movie shows the MECHANISM; the replicates
# measure the EFFECT (how much, and whether the direction survives run-to-run
# noise). Here removing concurrency roughly halves the baseline, and adding it
# lifts spread further.
final_prev <- function(formation, target.stats, nsims = 200) {
  set.seed(1234)
  nw <- network_initialize(n = 100)
  est <- netest(nw, formation, target.stats, coef.diss, verbose = FALSE)
  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)
}
data.frame(
  specification = c("No concurrency (0)", "Baseline (edges only)", "High concurrency (30)"),
  mean_infected_of_100 = c(
    final_prev(~edges + concurrent, c(40, 0)),
    final_prev(~edges, 40),
    final_prev(~edges + concurrent, c(40, 30))))