# Lab 2: Mixing & Concurrency in Two-Group Network Models (Worked Solution)
#
# Accompanies the SISMID course "Network Modeling for Epidemics" (NME).
# Source page: mod4-Lab2.qmd
#
# Builds on the two-group model from the Nodal Attributes tutorial
# (mod4-NodalAttributes), where Group 1 is women and Group 2 is men.
#
#   Part 1 (everyone) varies the group-specific acquisition probability on
#   the tutorial's own network model.
#   Part 2 (if you have time) relaxes the tutorial's all-cross-sex mixing
#   assumption and works out what that costs and how to pay for it.
#
# No set.seed is used, so results vary slightly run to run. Run top to bottom.


# ---- Setup ----

# (The tutorial clears the workspace here; we leave it commented so sourcing
# this script will not erase your other objects. Uncomment for a clean slate.)
# rm(list = ls())
library(EpiModel)


# ---- Rebuild the Tutorial's Network Model ----

# 250 women (Group 1) and 250 men (Group 2). The `group` attribute must take
# the values 1 and 2 to unlock EpiModel's two-group functionality.
num.g1 <- num.g2 <- 250
n <- num.g1 + num.g2
nw <- network_initialize(n)
nw <- set_vertex_attribute(nw, "group", rep(1:2, times = c(num.g1, num.g2)))

# Same targets as the tutorial: mean degree 0.66, 5% of women and 15% of men
# with concurrent partners.
edges.target <- n * 0.66 / 2        # 165
conc.g1      <- num.g1 * 0.05       # 12.5
conc.g2      <- num.g2 * 0.15       # 37.5

coef.diss <- dissolution_coefs(dissolution = ~offset(edges), duration = 20)

# nodematch("group") = 0 forces every tie to be cross-sex.
form.ref    <- ~edges + nodematch("group") + concurrent(by = "group", levels = NULL)
targets.ref <- c(edges.target, 0, conc.g1, conc.g2)

est.ref <- netest(nw, form.ref, targets.ref, coef.diss)

# Shared epidemic settings. Only the parameters change below. ncores = 5 runs
# the replicates in parallel; EpiModel caps this at the cores you actually have.
init    <- init.net(i.num = 10, i.num.g2 = 10, r.num = 0, r.num.g2 = 0)
control <- control.net(type = "SIR", nsims = 5, nsteps = 500, ncores = 5)

# Cumulative number of infections by group: total new infections over the whole
# run, from the across-simulation means. na.rm drops the undefined first step.
burden <- function(sim, label) {
  d <- as.data.frame(sim, out = "mean")
  inf.F <- sum(d$si.flow, na.rm = TRUE)
  inf.M <- sum(d$si.flow.g2, na.rm = TRUE)
  data.frame(model = label, cuml.infs.F = inf.F, cuml.infs.M = inf.M,
             ratio.F.over.M = inf.F / inf.M)
}


# =====================================================================
# PART 1: GROUP-SPECIFIC ACQUISITION PROBABILITY
# =====================================================================

# EpiModel's two-group naming: inf.prob is the per-act acquisition
# probability for Group 1 (women), inf.prob.g2 for Group 2 (men). Both are
# keyed to the SUSCEPTIBLE partner's group, never the infected partner's.

# Baseline: the tutorial's equal-risk run.
sim0 <- netsim(est.ref, param.net(inf.prob = 0.3, inf.prob.g2 = 0.3,
                                  rec.rate = 0.02, rec.rate.g2 = 0.02),
               init, control)

# Scenario A: double women's per-act acquisition probability to 0.6.
simA <- netsim(est.ref, param.net(inf.prob = 0.6, inf.prob.g2 = 0.3,
                                  rec.rate = 0.02, rec.rate.g2 = 0.02),
               init, control)

# Scenario B: women stay at 0.6, men drop to 0.1.
simB <- netsim(est.ref, param.net(inf.prob = 0.6, inf.prob.g2 = 0.1,
                                  rec.rate = 0.02, rec.rate.g2 = 0.02),
               init, control)

# Prevalence (left) and incidence (right) by group, for Scenario A.
par(mfrow = c(1, 2))
plot(simA, y = c("i.num", "i.num.g2"), popfrac = TRUE, qnts = FALSE,
     ylim = c(0, 0.6), legend = TRUE, main = "Scenario A: Prevalence")
plot(simA, y = c("si.flow", "si.flow.g2"), qnts = FALSE, legend = TRUE,
     main = "Scenario A: Incidence")

rbind(burden(sim0, "Baseline: 0.3 / 0.3"),
      burden(simA, "A: women 0.6 / men 0.3"),
      burden(simB, "B: women 0.6 / men 0.1"))


# ---- Part 1 Answers ----

# Q1. Doubling women's per-act acquisition (0.3 to 0.6) does NOT double
# women's cumulative number of infections. Two things limit it.
#
# First, and most directly, transmission across a single partnership is already
# near certain at 0.3. Partnerships persist (mean 20 steps) and the infectious
# period is long (mean 50 steps = 1 / rec.rate), so a susceptible woman paired
# with an infected man is exposed step after step, not just once. With a per-act
# probability of 0.3 and on the order of 20 exposures, her chance of eventually
# acquiring in that one partnership is already very high: 1 - 0.7^20 is
# essentially 1, and even a dozen exposures gives about 0.97. Raising the
# per-act probability to 0.6 has almost no room left to act, since an already
# near-certain per-partnership transmission cannot get much more certain. The
# per-act number and the per-partnership outcome live on very different scales;
# this per-act-to-per-partnership gap is the network-specific piece.
#
# Second, this is a closed SIR population. The infection count cannot exceed
# group size, and as susceptibles deplete the epidemic runs out of people to
# infect. So the doubled risk shows up mostly in the SPEED and PEAK of the
# epidemic (it climbs sooner and higher) rather than in its final total. Near
# saturation, the raw infection count is a blunt summary.
#
# Q2. Lowering men's acquisition to 0.1 also lowers WOMEN's incidence, even
# though women's own acquisition is still 0.6. Every tie is cross-sex, so
# women can only be infected by men; fewer infected men means fewer infectious
# partners for women. Reducing risk in one group protects the other through
# the shared network, which is the core rationale for partner-based and
# group-targeted interventions.
#
# Q3. Even at equal risk, women carry the higher burden (the baseline row
# above). Mean degree is equal by construction, but men's concurrency is
# higher, and concurrency raises a person's onward TRANSMISSION without
# raising their own ACQUISITION. So women absorb the cost of men's
# concurrency. Part 2 shows this conclusion is more fragile than it looks.


# =====================================================================
# PART 2: RELAXING THE ALL-CROSS-SEX ASSUMPTION
# =====================================================================

# The tutorial's nodematch("group") = 0 is an ABSOLUTE mixing rule, a
# prohibition. Real mixing is usually STATISTICAL, a tendency. What happens
# if we just allow some same-sex ties?

# The diagnostic. nodefactor("group") counts TIE-ENDS per group, so
# nodefactor.group.k / 250 is group k's mean degree. levels = NULL is
# required: by default nodefactor drops the first level as a reference, which
# would leave you no group-1 statistic to compare against.
mon <- ~edges + nodematch("group", diff = TRUE) +
  nodefactor("group", levels = NULL) + concurrent(by = "group")

mixing_check <- function(dx, label) {
  s <- get_nwstats(dx)
  out <- data.frame(model = label,
                    FF = mean(s$nodematch.group.1),
                    MM = mean(s$nodematch.group.2),
                    meandeg.F = mean(s$nodefactor.group.1) / num.g1,
                    meandeg.M = mean(s$nodefactor.group.2) / num.g2)
  out$ratio.M.over.F <- out$meandeg.M / out$meandeg.F
  out
}

dx.ref <- netdx(est.ref, nsims = 10, nsteps = 500, ncores = 5,
                nwstats.formula = mon)


# ---- Step 1: The Simple Change, and What Breaks ----

# In a 250/250 population, random mixing would make about half of all ties
# same-sex. Target that, changing nothing else.
match.target <- 0.5 * edges.target        # 82.5

est.naive <- netest(nw, ~edges + nodematch("group") +
                      concurrent(by = "group", levels = NULL),
                    c(edges.target, match.target, conc.g1, conc.g2),
                    coef.diss)

dx.naive <- netdx(est.naive, nsims = 10, nsteps = 500, ncores = 5,
                  nwstats.formula = mon)
dx.naive

rbind(mixing_check(dx.ref,   "Reference: nodematch = 0"),
      mixing_check(dx.naive, "Naive: nodematch = 82.5"))

# Two things are true at once, and this is the whole point of Step 1.
#
# 1. Every target we set is MET. Check the netdx table: edges, nodematch and
#    both concurrent statistics all sit on target. The model fits. Nothing is
#    degenerate or over-constrained.
# 2. Mean degree has come apart anyway. Men now run about 30% above women,
#    and FF and MM are nowhere near each other. We never asked for that, and
#    we never asked against it either.


# ---- Step 2: Why, and the Fix ----

# The accounting. With FF = within-women ties, MM = within-men, FM = cross:
#
#   edges        = FF + MM + FM
#   women's ends = 2*FF + FM        men's ends = 2*MM + FM
#   difference   = 2 * (FF - MM)
#
# Groups are the same size, so mean degree is equal IF AND ONLY IF FF = MM.
#
# nodematch("group") supplies ONE statistic: FF + MM. It pins the SUM and
# says nothing about the SPLIT. But mean degree depends on the DIFFERENCE.
# The split is a free direction, and the ERGM uses it to satisfy the
# constraints we DID impose: men need 15% concurrency against women's 5%, and
# the cheapest way to buy it is to give men more ties. So MM drifts above FF.
#
# The problem is not too few degrees of freedom. It is too few CONSTRAINTS.
#
# diff = TRUE reports the within-group counts separately, one statistic per
# group, which pins the difference as well as the sum. Set FF = MM.
ff.target <- mm.target <- match.target / 2      # 41.25 each

est.fix <- netest(nw, ~edges + nodematch("group", diff = TRUE) +
                    concurrent(by = "group", levels = NULL),
                  c(edges.target, ff.target, mm.target, conc.g1, conc.g2),
                  coef.diss)

dx.fix <- netdx(est.fix, nsims = 10, nsteps = 500, ncores = 5,
                nwstats.formula = mon)

rbind(mixing_check(dx.ref,   "Reference: nodematch = 0"),
      mixing_check(dx.naive, "Naive: nodematch = 82.5"),
      mixing_check(dx.fix,   "Fixed: diff = TRUE, 41.25 each"))

# The ratio will be near 1.00 but not exactly 1.00, and that is correct: we
# pinned the EXPECTED FF and MM, not their value in any one draw. Check that
# the residual is only this: 2 * (MM - FF) / 250 predicts the observed
# mean-degree gap to the decimal. In the reference model the ratio is exactly
# 1.000 in every draw, because there equality is structural, not estimated.
#
# Do NOT generalize this to "always use diff = TRUE". At a target of 0 it
# returns -Inf coefficients, which is correct but unhelpful.


# ---- Step 3: The Epidemic, Both Ways ----

param <- param.net(inf.prob = 0.3, inf.prob.g2 = 0.3,
                   rec.rate = 0.02, rec.rate.g2 = 0.02)

sim.ref <- netsim(est.ref, param, init, control)
sim.fix <- netsim(est.fix, param, init, control)

rbind(burden(sim.ref, "Absolute mixing (tutorial)"),
      burden(sim.fix, "Statistical mixing (50% same-sex)"))


# ---- Part 2 Answers ----

# Q1. Why did mean degree diverge in Step 1, when every target was met?
# Because nodematch("group") pins FF + MM and mean degree depends on FF - MM.
# The split was unconstrained, so the model spent it buying men the higher
# concurrency we asked for. See Step 2.
#
# Q2. What does the epidemic comparison show?
# Under the tutorial's absolute mixing women carry roughly 9% more infections
# than men. Under statistical mixing, at the same mean degree, the same
# concurrency targets and the same acquisition probability, that gap largely
# disappears. Do not over-read one run: at nsims = 5 the ratio lands on either
# side of 1.00 (roughly 0.98 to 1.03 across runs), so if yours comes out below
# 1 that is noise, not a reversal. Push the power up (dozens of sims per arm)
# and the ratio settles right on 1.0, with a confidence interval that
# comfortably straddles it: a coin flip.
#
# Q3. Why should the gap close?
# Part 1's Q3 answer said women bear the burden of men's concurrency. That
# argument needs women's partners to BE men. With FF = MM, women's tie-ends
# split half to women and half to men, and men's split the same way. Both
# groups now face an identical mix of partners, so both face an identical
# force of infection, and the asymmetry has nowhere to land.
#
# The lesson is not that concurrency stopped mattering. It is that the
# tutorial's burden gap was never about concurrency alone; it was about
# concurrency being SEGREGATED by sex. Mixing structure decides who bears the
# cost of someone else's risk behavior.
#
# Note what same-sex ties bought epidemiologically: transmission between two
# women is now possible, and it happens. In the tutorial's model it was
# structurally impossible.


# ---- Optional: Things to Try ----

# - Change the same-sex target from 50% to 20% or 80%; re-run Steps 1 and 2.
#   How does the size of the divergence track the target?
# - Set FF and MM to different values on purpose (say 60 and 22.5). Predict
#   the mean degrees from the Step 2 accounting BEFORE you fit, then check.
# - Make the concurrency targets equal (10% in both groups) and re-run Step 1.
#   Does the naive model still diverge? Why not?