## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(collapse = TRUE, comment = "#>")
library(drmTMB)

## -----------------------------------------------------------------------------
set.seed(106)
n <- 360
soil_counts <- data.frame(
  habitat = factor(
    rep(c("degraded", "restored"), each = n / 2),
    levels = c("degraded", "restored")
  ),
  surface = factor(
    sample(c("litter", "bare"), n, replace = TRUE, prob = c(0.72, 0.28)),
    levels = c("litter", "bare")
  ),
  moisture = as.numeric(scale(runif(n, 0.15, 0.95))),
  trap_nights = sample(2:5, n, replace = TRUE)
)

restored <- as.numeric(soil_counts$habitat == "restored")
bare <- as.numeric(soil_counts$surface == "bare")

rate <- exp(log(2.4) + 0.35 * restored + 0.30 * soil_counts$moisture)
mu <- soil_counts$trap_nights * rate
sigma_nb2 <- exp(-0.75 + 0.40 * restored)
zi <- plogis(-3.2 + 1.6 * bare - 0.25 * restored)

structural_zero <- runif(n) < zi
soil_counts$springtails <- ifelse(
  structural_zero,
  0L,
  rnbinom(n, size = 1 / sigma_nb2^2, mu = mu)
)

head(soil_counts)

## -----------------------------------------------------------------------------
fit_nb2 <- drmTMB(
  bf(
    springtails ~ habitat + moisture + offset(log(trap_nights)),
    sigma ~ habitat
  ),
  family = nbinom2(),
  data = soil_counts
)

## -----------------------------------------------------------------------------
check_drm(fit_nb2)

## -----------------------------------------------------------------------------
coef(fit_nb2, "mu")
coef(fit_nb2, "sigma")

sigma_ratio <- exp(coef(fit_nb2, "sigma")["habitatrestored"])
c(
  sigma_ratio_restored_vs_degraded = sigma_ratio,
  variance_multiplier_at_same_mu = sigma_ratio^2
)

## -----------------------------------------------------------------------------
fit_zinb2 <- drmTMB(
  bf(
    springtails ~ habitat + moisture + offset(log(trap_nights)),
    sigma ~ habitat,
    zi ~ surface
  ),
  family = nbinom2(),
  data = soil_counts
)

check_drm(fit_zinb2)

## -----------------------------------------------------------------------------
coef(fit_zinb2, "zi")

zero_grid <- data.frame(
  habitat = factor(c("degraded", "degraded"), levels = levels(soil_counts$habitat)),
  surface = factor(c("litter", "bare"), levels = levels(soil_counts$surface)),
  moisture = 0,
  trap_nights = 3
)

data.frame(
  surface = zero_grid$surface,
  structural_zero_probability = predict(fit_zinb2, newdata = zero_grid, dpar = "zi")
)

## -----------------------------------------------------------------------------
new_traps <- data.frame(
  habitat = factor(c("degraded", "restored"), levels = levels(soil_counts$habitat)),
  surface = factor(c("litter", "litter"), levels = levels(soil_counts$surface)),
  moisture = c(0, 0),
  trap_nights = c(3, 3)
)

mu_hat <- predict(fit_zinb2, newdata = new_traps, dpar = "mu")
sigma_hat <- predict(fit_zinb2, newdata = new_traps, dpar = "sigma")
zi_hat <- predict(fit_zinb2, newdata = new_traps, dpar = "zi")

data.frame(
  habitat = new_traps$habitat,
  conditional_mean = mu_hat,
  sigma = sigma_hat,
  structural_zero_probability = zi_hat,
  unconditional_mean = (1 - zi_hat) * mu_hat,
  unconditional_variance =
    (1 - zi_hat) * (mu_hat + sigma_hat^2 * mu_hat^2) +
      zi_hat * (1 - zi_hat) * mu_hat^2
)

## ----count-model-parts-figure, eval=requireNamespace("ggplot2", quietly = TRUE), fig.width=7.4, fig.height=4.8, fig.cap="Response-scale model parts for the zero-inflated NB2 example. Points show fitted conditional counts, unconditional counts, NB2 extra-Poisson scale, and structural-zero probabilities on separate facets with their own x scales. No interval bars are drawn because this figure compares fitted components, not confidence intervals.", fig.alt="Faceted point plot for the zero-inflated NB2 soil-count example. Separate facets show conditional expected counts, unconditional expected counts, NB2 extra-Poisson sigma, and structural-zero probabilities for habitat and surface combinations."----
library(ggplot2)

count_plot_grid <- expand.grid(
  habitat = factor(levels(soil_counts$habitat), levels = levels(soil_counts$habitat)),
  surface = factor(levels(soil_counts$surface), levels = levels(soil_counts$surface))
)
count_plot_grid$moisture <- 0
count_plot_grid$trap_nights <- 3
count_plot_grid$conditional_mean <- predict(
  fit_zinb2,
  newdata = count_plot_grid,
  dpar = "mu"
)
count_plot_grid$sigma <- predict(fit_zinb2, newdata = count_plot_grid, dpar = "sigma")
count_plot_grid$structural_zero_probability <- predict(
  fit_zinb2,
  newdata = count_plot_grid,
  dpar = "zi"
)
count_plot_grid$unconditional_mean <-
  (1 - count_plot_grid$structural_zero_probability) *
  count_plot_grid$conditional_mean
count_plot_grid$row_label <- paste(count_plot_grid$habitat, count_plot_grid$surface)
count_plot_long <- rbind(
  data.frame(
    count_plot_grid[c("habitat", "row_label")],
    component = "Conditional mean",
    value = count_plot_grid$conditional_mean
  ),
  data.frame(
    count_plot_grid[c("habitat", "row_label")],
    component = "Unconditional mean",
    value = count_plot_grid$unconditional_mean
  ),
  data.frame(
    count_plot_grid[c("habitat", "row_label")],
    component = "NB2 sigma",
    value = count_plot_grid$sigma
  ),
  data.frame(
    count_plot_grid[c("habitat", "row_label")],
    component = "Structural-zero probability",
    value = count_plot_grid$structural_zero_probability
  )
)
count_plot_long$component <- factor(
  count_plot_long$component,
  levels = c(
    "Conditional mean",
    "Unconditional mean",
    "NB2 sigma",
    "Structural-zero probability"
  )
)
count_plot_long$row_label <- factor(
  count_plot_long$row_label,
  levels = rev(unique(count_plot_grid$row_label))
)

ggplot(count_plot_long, aes(value, row_label, colour = habitat)) +
  geom_point(size = 2.7) +
  facet_wrap(~component, scales = "free_x", ncol = 2) +
  scale_colour_manual(values = c("degraded" = "#D55E00", "restored" = "#009E73")) +
  labs(
    title = "Zero-inflated counts have several fitted pieces",
    subtitle = "Facets keep count means, NB2 sigma, and structural-zero probabilities separate",
    x = "Response-scale fitted value",
    y = NULL,
    colour = "Habitat"
  ) +
  theme_minimal(base_size = 11) +
  theme(
    panel.grid.minor = element_blank(),
    panel.grid.major.y = element_blank(),
    legend.position = "bottom",
    plot.title = element_text(face = "bold"),
    plot.subtitle = element_text(colour = "grey30"),
    strip.text = element_text(face = "bold")
  )

## -----------------------------------------------------------------------------
AIC(fit_nb2, fit_zinb2)