# COVID Demography Lab (Worked Solution)
#
# Accompanies the SISMID course "Network Modeling for Epidemics" (NME).
# Source page: mod13-COVID1-Lab.qmd
#
# The lab takes the COVID model from the Module 13 tutorial and adds a
# CATEGORICAL age attribute, age.grp, splitting the population into minors
# (under 18, coded 0) and adults (18 and over, coded 1). The point of the
# exercise is practice with the full life cycle of a nodal attribute: set it
# on the network, keep it current inside the modules as people age and as
# new people arrive, and then use it in the network model.
#
# This script PAIRS with mod13-COVID1-fx-Lab1.R, which holds the module
# functions. Keep both in the same working directory; the source() call
# below uses a relative path. The functions must be sourced before
# control.net, which captures them by value.
#
# The four sections below follow the lab's four steps. Run top to bottom.
#
# Note the one decision the lab deliberately leaves open is the target
# statistic for age-group mixing. Step 3 works through how to choose it.

library("EpiModel")

# The lab page asks you to start from a clean environment, so that nothing
# lingers from the tutorial. This matters here because the tutorial leaves
# behind an object named absdiff, and this script reuses the name edges.
rm(list = ls())
set.seed(12345)


# =====================================================================
# STEP 1: Add the age.grp attribute to the network
# =====================================================================

# Range of possible ages
ages <- 0:85

# Age-specific mortality rates
# Rates per 100,000 for age groups: <1, 1-4, 5-9, 10-14, 15-19, 20-24, 25-29,
#                                   30-34, 35-39, 40-44, 45-49, 50-54, 55-59,
#                                   60-64, 65-69, 70-74, 75-79, 80-84, 85+
departure_rate <- c(588.45, 24.8, 11.7, 14.55, 47.85, 88.2, 105.65, 127.2,
                    154.3, 206.5, 309.3, 495.1, 736.85, 1051.15, 1483.45,
                    2294.15, 3642.95, 6139.4, 13938.3)

# Per-capita daily death rate (divide by 365,000)
dr_pp_pd <- departure_rate / 1e5 / 365

# Create vector of daily death rates
age_spans <- c(1, 4, rep(5, 16), 1)
dr_vec <- rep(dr_pp_pd, times = age_spans)
data.frame(ages, dr_vec)

# Initialize network
n <- 1000
nw <- network_initialize(n)

# Set continuous age attribute, as in the tutorial. We keep this even though
# the lab adds a categorical version, because the mortality lookup in dfunc
# still needs age in years.
ageVec <- sample(ages, n, replace = TRUE)
nw <- set_vertex_attribute(nw, "age", ageVec)

# THE LAB'S STEP 1. Derive the binary age group from the continuous age with
# ifelse, and set it as a second vertex attribute. Minors are 0, adults 1.
# Setting it on the network (rather than only inside a module) is what makes
# it available to the formation formula in Step 3.
ageGrpVec <- ifelse(ageVec < 18, 0, 1)
nw <- set_vertex_attribute(nw, "age.grp", ageGrpVec)

# Confirm the split is what we expect before going further.
table(ageGrpVec)

# Create vector of disease statuses
statusVec <- rep("s", n)
init.latent <- sample(1:n, 50)
statusVec[init.latent] <- "e"
which(statusVec == "e")
table(statusVec)

# Set status attribute
nw <- set_vertex_attribute(nw, "status", statusVec)
nw


# =====================================================================
# STEP 2: Keep the attribute current inside the modules
# =====================================================================

# This step happens in the companion file, mod13-COVID1-fx-Lab1.R, which is
# sourced further down. Two changes there relative to the tutorial's module
# file, both of which you should read before running:
#
#   1. aging() recalculates age.grp from the updated age each step, so a
#      node crossing its 18th birthday changes group on its own.
#
#   2. afunc() appends age.grp = 0 for every arrival. Every attribute the
#      model tracks needs its own append_attr call; miss one and its vector
#      falls out of alignment with the others, silently.
#
# The lab also asks for a summary statistic validating the approach. The
# aging module now records prop.minor, the share of the population under 18,
# which we plot at the end.


# =====================================================================
# STEP 3: Use the attribute in the network model
# =====================================================================

# Review network object
nw

# The lab asks us to switch from absdiff("age") to a nodematch term on the
# categorical attribute. absdiff measured how far apart partners were in
# years; nodematch("age.grp") instead counts partnerships where both
# partners fall in the SAME age group, which is the natural way to express
# mixing on a categorical variable.
formation <- ~edges + degree(0) + nodematch("age.grp")

# Target statistics
mean_degree <- 2
edges <- mean_degree * (n/2)
isolates <- n * 0.08

# CHOOSING THE NODEMATCH TARGET. The lab says there is no single right
# answer here but warns that extreme values can strain the fit. The way to
# reason about it is to start from the RANDOM-MIXING baseline, which is not
# zero and is easy to overlook.
#
# Ages are drawn uniformly from 0:85, so minors (0 through 17) are 18 of the
# 86 values, about 0.209, and adults about 0.791. If partners paired at
# random, the share of matched partnerships would already be
#
#     0.209^2 + 0.791^2 = 0.044 + 0.626 = about 0.67
#
# So a target of 0.67 is not "no homophily expressed", it IS random mixing.
# Anything above it represents genuine assortativity. We pick 0.80, which is
# meaningfully assortative while leaving plenty of cross-group ties. A value
# like 0.96 would mean near-total age segregation, which is both hard to
# reconcile with a mean degree of 2 and the kind of extreme the lab warns
# about. Try other values and watch the diagnostics.
nodematch.stat <- edges * 0.80

target.stats <- c(edges, isolates, nodematch.stat)
target.stats

# Dissolution model
coef.diss <- dissolution_coefs(~offset(edges), 80, mean(dr_vec)*2)
coef.diss

# Fit the model
est <- netest(nw, formation, target.stats, coef.diss)

# Model diagnostics. Note the monitoring formula now tracks the nodematch
# term rather than absdiff, so we can see whether the mixing target holds.
dx <- netdx(est, nsims = 10, ncores = 5, nsteps = 500,
            nwstats.formula = ~edges + nodematch("age.grp") + degree(0:3),
            set.control.tergm =
              control.simulate.formula.tergm(MCMC.burnin.min = 10000))
print(dx)
plot(dx)


# =====================================================================
# STEP 4: Run the model
# =====================================================================

# Parameterize model
param <- param.net(inf.prob = 0.1,
                   act.rate = 3,
                   departure.rates = dr_vec,
                   departure.disease.mult = 100,
                   arrival.rate = 1/(365*85),
                   ei.rate = 0.05, ir.rate = 0.05)

# Initialize model
init <- init.net()

# Control settings (source the lab's function script, not the tutorial's)
source("mod13-COVID1-fx-Lab1.R")

# The lab asks for a single simulation, which is the right first move: it is
# quick and it is enough to confirm the new attribute is being maintained.
# Raise nsims to 5 or 10 once you are satisfied, if you want interval bands
# on the plots below.
control <- control.net(type = NULL,
                       nsims = 1,
                       ncores = 1,
                       nsteps = 300,
                       infection.FUN = infect,
                       progress.FUN = progress,
                       aging.FUN = aging,
                       departures.FUN = dfunc,
                       arrivals.FUN = afunc,
                       resimulate.network = TRUE,
                       tergmLite = TRUE,
                       set.control.tergm =
                         control.simulate.formula.tergm(MCMC.burnin.min = 5000),
                       nwstats.formula = ~edges + meandeg + degree(0:2) +
                         nodematch("age.grp"))

# Simulate model
sim <- netsim(est, param, init, control)
sim


# =====================================================================
# Checking that it behaved
# =====================================================================

# The three fitted terms over time. Expect the simulated statistics to sit
# somewhat BELOW their targets rather than exactly on them: as the tutorial
# discusses, the edges-only dissolution approximation combined with
# demography pulls the network off target once vital dynamics are running.
# A modest, stable gap is the expected result here; a statistic collapsing
# toward zero or drifting without settling would be the thing to worry about.
plot(sim, type = "formation",
     stats = c("edges", "nodematch.age.grp", "degree0"), plots.joined = FALSE)

# The lab's validation statistic. prop.minor starts near 0.21 (18 of the 86
# uniform ages are under 18) and drifts up to roughly 0.23 over the run.
#
# Worth being precise about WHY, because the obvious answer is the wrong one.
# It is not age-0 arrivals accumulating: over 300 steps this model sees only
# about 9 arrivals against about 58 departures, and the population actually
# SHRINKS from 1000 to around 950. The rise comes mostly from who is dying.
# Mortality is steeply age-graded, and departure.disease.mult multiplies it
# by 100 for infected nodes, so the COVID deaths (about 51 of the 58) fall
# overwhelmingly on adults. Removing adults faster than minors raises the
# minor share. Aging contributes essentially nothing here: 300 steps is
# 0.82 years, so almost nobody crosses the age-18 boundary during the run.
plot(sim, y = "prop.minor")

# Mean age. Aging pushes this up by 1/365 per step, while arrivals at age 0
# and age-skewed mortality push it down; over 300 steps the downward forces
# win, so expect a gentle decline rather than a rise.
plot(sim, y = "meanAge")

# Compare network stats for mean degree
plot(sim, type = "formation", stats = "meandeg", qnts = 1, ylim = c(1, 3))
abline(h = 2, lty = 2, lwd = 2)

# Simulation plot
plot(sim, qnts = 1)

# Disease incidence
plot(sim, y = "se.flow")

# Export data to data frame
df <- as.data.frame(sim, out = "mean")
head(df$meanAge)
tail(df$meanAge)
head(df$prop.minor)
tail(df$prop.minor)

# Population size over time
plot(sim, y = "num")

# Deaths per day
plot(sim, y = c("total.deaths", "covid.deaths"), qnts = FALSE, legend = TRUE)

# Calculate sum of COVID deaths
sum(df$covid.deaths, na.rm = TRUE)