## ----setup, include=FALSE-----------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 4.5
)

## ----data---------------------------------------------------------------------
library(ZINB.GP)

data_dir <- system.file("extdata", "oregon", package = "ZINB.GP")
if (data_dir == "") {
  data_dir <- file.path("..", "inst", "extdata", "oregon")
}

observations <- as.matrix(read.table(
  file.path(data_dir, "landslides_county_year.dat")
))
county_data <- read.csv(file.path(data_dir, "counties.csv"))
years <- read.csv(file.path(data_dir, "Years.csv"))[[2]]

dim(observations)
mean(observations == 0)

## ----design-------------------------------------------------------------------
panel <- make_y_Vs_Vt(observations)
y <- panel$y
Vs <- panel$Vs
Vt <- panel$Vt

stopifnot(
  length(y) == nrow(Vs),
  nrow(Vs) == nrow(Vt),
  ncol(Vs) == nrow(observations) - 1,
  ncol(Vt) == ncol(observations) - 1
)

## ----covariate----------------------------------------------------------------
elevation_sd <- county_data$county_elev_std
X <- cbind(
  "(Intercept)" = 1,
  elevation_sd = as.vector(
    cbind(1, Vs) %*% elevation_sd
  )
)

head(X)

## ----county-summary-----------------------------------------------------------
county_mean <- rowMeans(observations)
county_nonzero <- rowMeans(observations > 0)

par(mfrow = c(1, 2), mar = c(4, 4, 2, 1))
plot(
  elevation_sd, log1p(county_mean),
  xlab = "County elevation variability",
  ylab = "log(1 + average annual count)",
  pch = 19, col = "#2C7FB8"
)
plot(
  elevation_sd, county_nonzero,
  xlab = "County elevation variability",
  ylab = "Fraction of years with a recorded event",
  pch = 19, col = "#D95F0E"
)
par(mfrow = c(1, 1))

## ----heatmap------------------------------------------------------------------
image(
  x = seq_along(years),
  y = seq_len(nrow(observations)),
  z = t(log1p(observations)),
  axes = FALSE,
  xlab = "Year",
  ylab = "County",
  main = "Observed county-year counts: log(1 + count)",
  col = hcl.colors(25, "YlOrRd", rev = TRUE)
)
year_ticks <- pretty(seq_along(years))
year_ticks <- year_ticks[year_ticks >= 1 & year_ticks <= length(years)]
axis(1, at = year_ticks, labels = years[year_ticks])
axis(2, at = seq_len(nrow(observations)), labels = FALSE)
box()

## ----overdispersion-----------------------------------------------------------
county_variance <- apply(observations, 1, var)
plot(
  county_mean, county_variance,
  log = "xy",
  xlab = "County mean",
  ylab = "County variance",
  pch = 19, col = "#756BB1",
  main = "Overdispersion by county"
)
abline(a = 0, b = 1, lty = 2)

## ----distances----------------------------------------------------------------
Ds <- 100 * as.matrix(read.table(
  file.path(data_dir, "county_dist_scale.dat")
))
Dt <- as.matrix(read.table(file.path(data_dir, "time_dist.dat")))
county_coords <- as.matrix(read.table(
  file.path(data_dir, "counties_coords.dat")
))

stopifnot(
  nrow(Ds) == ncol(Vs) + 1,
  nrow(Dt) == ncol(Vt) + 1,
  nrow(county_coords) == nrow(Ds)
)

## ----fit, eval=FALSE----------------------------------------------------------
#  lt_prior <- list(max = 125, mh_sd = 0.5, a = 1, b = 0.001)
#  ls_prior <- list(max = 250, mh_sd = 3, a = 1, b = 0.001)
#  
#  fit <- ZINB_GP(
#    X = X,
#    y = y,
#    coords = county_coords,
#    Vs = Vs,
#    Vt = Vt,
#    Ds = Ds,
#    Dt = Dt,
#    nsim = 162000,
#    burn = 20000,
#    thin = 100,
#    save_ypred = FALSE,
#    print_progress = TRUE,
#    use_count_gp = TRUE,
#    use_inflation_gp = TRUE,
#    ltPrior = lt_prior,
#    lsPrior = ls_prior
#  )

## ----cached-fit---------------------------------------------------------------
fit <- readRDS(file.path(data_dir, "posterior_draws.rds"))
names(fit)

## ----fixed-effects------------------------------------------------------------
interval <- function(x) {
  quantile(x, probs = c(0.025, 0.5, 0.975), names = FALSE)
}

fixed_effect_intervals <- rbind(
  "At-risk: intercept" = interval(fit$Alpha[, 1]),
  "At-risk: elevation variability" = interval(fit$Alpha[, 2]),
  "Count: intercept" = interval(fit$Beta[, 1]),
  "Count: elevation variability" = interval(fit$Beta[, 2])
)
colnames(fixed_effect_intervals) <- c("2.5%", "median", "97.5%")
round(fixed_effect_intervals, 3)

## ----recenter-----------------------------------------------------------------
recenter_draws <- function(reduced_effect, intercept) {
  full_effect <- cbind(0, reduced_effect)
  shift <- rowMeans(full_effect)
  list(
    effect = sweep(full_effect, 1, shift),
    intercept = intercept + shift
  )
}

at_risk_space <- recenter_draws(fit$A, fit$Alpha[, 1])
count_space <- recenter_draws(fit$C, fit$Beta[, 1])
at_risk_time <- recenter_draws(fit$B, at_risk_space$intercept)
count_time <- recenter_draws(fit$D, count_space$intercept)

# One numerical check: centering preserves the spatial contribution.
j <- 1
s <- 10
baseline_scale <- fit$Alpha[j, 1] + c(0, fit$A[j, ])[s]
centered <- at_risk_space$intercept[j] + at_risk_space$effect[j, s]
stopifnot(isTRUE(all.equal(baseline_scale, centered)))

## ----spatial-effects----------------------------------------------------------
plot_spatial_effect <- function(value, title) {
  palette <- hcl.colors(30, "Blue-Red 3")
  bins <- cut(value, breaks = 30, include.lowest = TRUE)
  plot(
    county_coords[, 1], county_coords[, 2],
    asp = 1, pch = 21, cex = 1.8,
    bg = palette[as.integer(bins)],
    xlab = "Scaled east-west coordinate",
    ylab = "Scaled north-south coordinate",
    main = title
  )
}

par(mfrow = c(1, 2))
plot_spatial_effect(
  colMeans(at_risk_space$effect),
  "At-risk spatial effect"
)
plot_spatial_effect(
  colMeans(count_space$effect),
  "Conditional-count spatial effect"
)
par(mfrow = c(1, 1))

## ----temporal-effects---------------------------------------------------------
par(mfrow = c(1, 2), mar = c(4, 4, 2, 1))
plot(
  years, colMeans(at_risk_time$effect),
  type = "l", lwd = 2, col = "#D95F0E",
  xlab = "Year", ylab = "Centered effect",
  main = "At-risk temporal effect"
)
abline(h = 0, lty = 2)
plot(
  years, colMeans(count_time$effect),
  type = "l", lwd = 2, col = "#2C7FB8",
  xlab = "Year", ylab = "Centered effect",
  main = "Conditional-count temporal effect"
)
abline(h = 0, lty = 2)
par(mfrow = c(1, 1))

## ----fitted-means-------------------------------------------------------------
draw_id <- seq(1, nrow(fit$Alpha), by = 10)

eta1 <- fit$Alpha[draw_id, , drop = FALSE] %*% t(X) +
  fit$A[draw_id, , drop = FALSE] %*% t(Vs) +
  fit$B[draw_id, , drop = FALSE] %*% t(Vt)
eta2 <- fit$Beta[draw_id, , drop = FALSE] %*% t(X) +
  fit$C[draw_id, , drop = FALSE] %*% t(Vs) +
  fit$D[draw_id, , drop = FALSE] %*% t(Vt)

fitted_mean <- plogis(eta1) * exp(eta2)
fitted_mean <- sweep(fitted_mean, 1, fit$R[draw_id], "*")
posterior_mean <- colMeans(fitted_mean)

observed_by_year <- colSums(observations)
fitted_by_year <- colSums(matrix(
  posterior_mean,
  nrow = nrow(observations)
))

plot(
  years, log1p(observed_by_year),
  type = "l", lwd = 2, col = "black",
  xlab = "Year", ylab = "log(1 + annual total)",
  main = "Observed and posterior fitted annual totals"
)
lines(years, log1p(fitted_by_year), lwd = 2, col = "#2C7FB8")
legend(
  "topleft",
  legend = c("Observed", "Posterior fitted mean"),
  col = c("black", "#2C7FB8"), lwd = 2, bty = "n"
)

## ----coda---------------------------------------------------------------------
scalar_draws <- cbind(
  alpha_intercept = fit$Alpha[, 1],
  alpha_elevation = fit$Alpha[, 2],
  beta_intercept = fit$Beta[, 1],
  beta_elevation = fit$Beta[, 2],
  range_time_at_risk = fit$L1t,
  range_time_count = fit$L2t,
  range_space_at_risk = fit$L1s,
  range_space_count = fit$L2s,
  dispersion = fit$R
)

chain <- coda::mcmc(scalar_draws)
round(coda::effectiveSize(chain), 1)

## ----trace-acf----------------------------------------------------------------
par(mfrow = c(1, 2))
coda::traceplot(chain[, "range_space_count"], main = "Trace: count spatial range")
coda::autocorr.plot(
  chain[, "range_space_count"],
  auto.layout = FALSE,
  main = "Autocorrelation"
)
par(mfrow = c(1, 1))

## ----multiple-chains, eval=FALSE----------------------------------------------
#  extract_scalars <- function(x) {
#    cbind(
#      alpha_intercept = x$Alpha[, 1],
#      alpha_elevation = x$Alpha[, 2],
#      beta_intercept = x$Beta[, 1],
#      beta_elevation = x$Beta[, 2],
#      range_time_at_risk = x$L1t,
#      range_time_count = x$L2t,
#      range_space_at_risk = x$L1s,
#      range_space_count = x$L2s,
#      dispersion = x$R
#    )
#  }
#  
#  # fits must come from independent runs with different random-number seeds.
#  chains <- coda::mcmc.list(lapply(fits, function(x) {
#    coda::mcmc(extract_scalars(x))
#  }))
#  coda::gelman.diag(chains, multivariate = FALSE)
#  coda::effectiveSize(chains)

## ----posterior-wrapper, eval=FALSE--------------------------------------------
#  draws <- posterior::as_draws_matrix(scalar_draws)
#  posterior::summarise_draws(
#    draws,
#    posterior::mean,
#    posterior::sd,
#    posterior::ess_bulk,
#    posterior::ess_tail
#  )