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

## ----setup--------------------------------------------------------------------
library(nlfh)

data(acs_dat)
acs_small <- as.data.frame(acs_dat[1:80, ])

fh_formula <- MedInc ~ SNAPRate + PovRate + White + Black + Hispanic + Asian
sampling_variance <- acs_small$MedIncSE^2

fit_control <- list(
  n_iter = 120,
  burn_in = 60,
  progress = FALSE
)

## ----fit-models---------------------------------------------------------------
set.seed(1)

fit_linear <- fit_fh(
  formula = fh_formula,
  sampling_variance = sampling_variance,
  data = acs_small,
  method = "linear",
  control = fit_control
)

fit_rnn <- fit_fh(
  formula = fh_formula,
  sampling_variance = sampling_variance,
  data = acs_small,
  method = "rnn",
  control = c(fit_control, list(n_hidden = 20))
)

fit_bart <- fit_fh(
  formula = fh_formula,
  sampling_variance = sampling_variance,
  data = acs_small,
  method = "bart",
  control = c(fit_control, list(n_bart_samples = 3, n_trees = 10))
)

## ----compare-dic--------------------------------------------------------------
fits <- list(
  linear = fit_linear,
  rnn = fit_rnn,
  bart = fit_bart
)

dic_table <- data.frame(
  model = names(fits),
  dic = vapply(fits, function(x) x$dic, numeric(1)),
  row.names = NULL
)
dic_table <- dic_table[order(dic_table$dic), ]
dic_table$delta_dic <- dic_table$dic - min(dic_table$dic)

knitr::kable(dic_table, digits = 2)

## ----summary-print------------------------------------------------------------
summary(fit_linear)

## ----variance-summaries-------------------------------------------------------
variance_table <- do.call(
  rbind,
  lapply(names(fits), function(model) {
    out <- summary(fits[[model]])$variance
    out$model <- model
    out[, c("model", "parameter", "mean", "sd", "median", "q2.5", "q97.5")]
  })
)

knitr::kable(variance_table, digits = 2)

## ----bart-variable-importance-------------------------------------------------
knitr::kable(
  data.frame(
    variable = names(fit_bart$variable_importance),
    importance = as.numeric(fit_bart$variable_importance)
  ),
  digits = 3
)

## ----coefficient-summaries----------------------------------------------------
linear_summary <- summary(fit_linear)
knitr::kable(linear_summary$coefficients, digits = 2)

## ----area-summaries-----------------------------------------------------------
area_summary <- fitted(fit_bart, full = TRUE)
knitr::kable(head(area_summary, 8), digits = 2)

## ----model-interval-plot, fig.height = 5--------------------------------------
interval_fits <- list(
  Linear = fit_linear,
  RNN = fit_rnn,
  BART = fit_bart
)
plot_areas <- seq_len(12)
model_offsets <- c(-0.22, 0, 0.22)
model_cols <- c("#1f5f99", "#7c3aed", "#b45309")

interval_summaries <- lapply(interval_fits, fitted, full = TRUE)
mean_range <- range(
  unlist(lapply(interval_summaries, function(x) {
    c(x$q2.5[plot_areas], x$q97.5[plot_areas])
  })),
  na.rm = TRUE
)

plot(
  NA,
  xlim = c(0.5, length(plot_areas) + 0.5),
  ylim = mean_range,
  xaxt = "n",
  xlab = "Area",
  ylab = "Posterior estimate of theta"
)
axis(1, at = plot_areas, labels = plot_areas)

for (i in seq_along(interval_summaries)) {
  out <- interval_summaries[[i]][plot_areas, ]
  x_pos <- plot_areas + model_offsets[i]
  segments(
    x0 = x_pos,
    y0 = out$q2.5,
    x1 = x_pos,
    y1 = out$q97.5,
    col = model_cols[i],
    lwd = 2
  )
  points(
    x_pos,
    out$mean,
    pch = 19,
    col = model_cols[i]
  )
}

legend(
  "topright",
  legend = names(interval_fits),
  col = model_cols,
  pch = 19,
  lwd = 2,
  bty = "n"
)

## ----fitted-plot--------------------------------------------------------------
plot(
  acs_small$MedInc,
  fitted(fit_bart),
  xlab = "Direct estimate",
  ylab = "Posterior mean of theta",
  pch = 19,
  col = "#2f6f8f"
)
abline(0, 1, col = "#9a3412", lwd = 2)

## ----posterior-draws----------------------------------------------------------
theta_draws <- posterior_draws(fit_bart, variable = "theta")
A_draws <- posterior_draws(fit_bart, variable = "A")

dim(theta_draws)
head(theta_draws[, 1:4])
head(A_draws)