---
title: "Stratified and partially-pooled CFR"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Stratified and partially-pooled CFR}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE}
knitr::opts_chunk$set(collapse = TRUE, comment = "#>")
# Fit chunks need CmdStan; evaluate them only where a toolchain is present so the
# vignette still builds on CmdStan-less builders (Windows R CMD check, CRAN) --
# there it shows the code without running it, while pkgdown renders in full.
knitr::opts_chunk$set(
  eval = requireNamespace("cmdstanr", quietly = TRUE) &&
    !is.null(tryCatch(cmdstanr::cmdstan_version(), error = function(e) NULL))
)
```

## Why stratify, and why pool

The CFR often varies by group -- age, treatment centre, calendar period -- and a
per-group estimate is frequently required. Two limiting cases bracket the
options. Fitting each group independently (*no pooling*) gives high-variance
estimates where a group's sample is small. Fitting one CFR for all groups
(*complete pooling*) ignores between-group variation.

*Partial pooling* lies between the two: each group's CFR is drawn from a common
distribution, so groups with more data are estimated close to independently while
groups with less data are shrunk toward the shared mean. `cfrnow` is an
[epidist](https://epidist.epinowcast.org/) model, so `prob` takes a `brms`
formula and pooling is a one-line change: `prob ~ (1 | group)`.

## A six-site line list

We simulate six treatment centres: two large (n = 450, 250) and four small
(n = 15--20), with true CFRs spread around a common mean. The large sites pin the
shared mean; the small ones carry little information on their own, so shrinkage
acts mainly on them. Site labels are synthetic.

```{r data}
library(cfrnow)
set.seed(20260714)

sites <- data.frame(
  site = c("A", "B", "C", "D", "E", "F"),
  n    = c(450, 250, 20, 15, 15, 18),
  cfr  = c(0.45, 0.50, 0.30, 0.65, 0.35, 0.60)
)

death_delay    <- LogNormal(mean = 11, sd = 6.5)
recovery_delay <- LogNormal(mean = 16, sd = 8)

ll <- do.call(rbind, Map(function(s, n, cfr) {
  cbind(
    simulate_linelist(
      n = n, cfr = cfr, delay = death_delay,
      recovery = recovery_delay, onset_days = 45
    ),
    site = s
  )
}, sites$site, sites$n, sites$cfr))

d <- prepare_cfr_data(ll, obs_time = as.Date("2026-02-20"), covariates = "site")
c(cases = d$n_cases, deaths = d$n_deaths)
```

## Three fits: complete, no, and partial pooling

The only thing that changes across the three fits is the CFR formula; the delay
and recovery delay are co-estimated in each. (Two short chains keep the build
time short; use more for inference.)

```{r fits, results = "hide", message = FALSE, warning = FALSE}
otd <- LogNormal(meanlog = Normal(2.4, 0.2), sdlog = Normal(0.5, 0.15))
otr <- LogNormal(meanlog = Normal(2.7, 0.2), sdlog = Normal(0.5, 0.15))
args <- list(
  delay = otd, recovery_delay = otr, prob_prior = Beta(1, 1),
  backend = "cmdstanr", chains = 2, iter = 1000, refresh = 0, seed = 1
)

# complete pooling: a single CFR
fit_pool    <- do.call(fit_cfr, c(list(d, formula = brms::bf(mu ~ 1, prob ~ 1)), args))

# no pooling: an independent CFR per site (one estimated logit-CFR per site)
fit_sep     <- do.call(fit_cfr, c(list(d, formula = brms::bf(mu ~ 1, prob ~ 0 + site)), args))

# partial pooling: per-site CFR shrunk to a common mean
fit_partial <- do.call(fit_cfr, c(list(d, formula = brms::bf(mu ~ 1, prob ~ (1 | site))), args))
```

## Reading the estimates

For the complete-pool fit, `summary()` reports a single `prob` row:

```{r summary-pool}
summary(fit_pool)
```

For a grouped fit it reports one `prob[<site>]` row per site:

```{r summary-grouped}
summary(fit_sep)
```

Collect those per-site CFR rows from the two grouped fits for the figure below;
the complete-pool fit supplies the shared mean.

```{r extract}
cfr_rows <- function(fit, approach) {
  s <- summary(fit)
  r <- s[grepl("^prob\\[", s$quantity), ] # the prob[<site>] rows
  data.frame(
    approach = approach,
    site = sub("^prob\\[(.*)\\]$", "\\1", r$quantity),
    median = r$q50, lower = r$q2.5, upper = r$q97.5
  )
}

est <- rbind(cfr_rows(fit_sep, "no pool"), cfr_rows(fit_partial, "partial pool"))
est <- merge(est, sites[, c("site", "n", "cfr")], by = "site")
names(est)[names(est) == "cfr"] <- "truth"

s_pool <- summary(fit_pool)
pool_mean <- s_pool$q50[s_pool$quantity == "prob"]
```

## Shrinkage

Each site shows its no-pool and partial-pool estimate against the truth (×) and
the complete-pool mean (dotted). For the large sites (A, B) the two estimates
nearly coincide. For the small sites (C--F) the no-pool interval is wide and the
estimate noisy; partial pooling pulls it toward the mean and narrows the interval
sharply -- lower variance at the cost of some bias.

```{r shrinkage, fig.width = 6.5, fig.height = 3.6}
library(ggplot2)
est$approach <- factor(est$approach, levels = c("no pool", "partial pool"))
est$site <- factor(est$site, levels = rev(sites$site))

ggplot(est, aes(median, site, colour = approach)) +
  geom_vline(xintercept = pool_mean, linetype = 3) +
  geom_point(aes(x = truth), shape = 4, size = 2.4, colour = "grey30") +
  geom_pointrange(aes(xmin = lower, xmax = upper), position = position_dodge(width = 0.55)) +
  scale_y_discrete(labels = function(s) paste0(s, "  (n=", sites$n[match(s, sites$site)], ")")) +
  labs(
    x = "CFR: median & 95% CrI   (× = truth, dotted = complete-pool mean)",
    y = NULL, colour = NULL
  ) +
  theme_minimal()
```

## Varying the delay by group

A pooled onset-to-death delay that appears irregular is often a mixture of
several populations. If sites resolve on different timescales, pool the delay by
site as well, so each retains its own onset-to-death distribution rather than
contributing to a single multimodal one:

```{r delay-pool, eval = FALSE}
fit_both <- do.call(fit_cfr, c(
  list(d, formula = brms::bf(mu ~ (1 | site), prob ~ (1 | site))), args
))
```

## Check the fit

Check that the parametric delay reproduces the observed delays; the real-time
correction depends on extrapolating that delay's tail.

```{r ppcheck, fig.width = 6.5, fig.height = 3.6, message = FALSE, warning = FALSE}
pp_check_cfr(fit_partial, type = "delay")
```