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

## ----library------------------------------------------------------------------
library(BATTS)

## ----data-location------------------------------------------------------------
set.seed(12345)
n1 <- 250L
loc <- list(x = rnorm(n1), y = rnorm(n1, mean = 0.5))
loc_truth <- function(t) dnorm(t, log = TRUE) - dnorm(t, mean = 0.5, log = TRUE)

## ----draw-samples-------------------------------------------------------------
col_x <- adjustcolor("firebrick", alpha.f = 0.6)
col_y <- adjustcolor("steelblue", alpha.f = 0.6)
draw_samples <- function(sample, main) {
  if (NCOL(sample$x) == 1L) {
    dens_x <- density(sample$x)
    dens_y <- density(sample$y)
    plot(dens_x, col = "firebrick", lwd = 2, main = main, xlab = "x",
         xlim = range(dens_x$x, dens_y$x), ylim = c(0, max(dens_x$y, dens_y$y)))
    lines(dens_y, col = "steelblue", lwd = 2, lty = 2)
    rug(sample$x, side = 1, col = col_x)
    rug(sample$y, side = 3, col = col_y)
    legend("topright", c("first sample (p)", "second sample (q)"),
           col = c("firebrick", "steelblue"), lty = c(1, 2), lwd = 2,
           bty = "n", cex = 0.8)
  } else {
    x <- rbind(sample$x, sample$y)
    first <- rep(c(TRUE, FALSE), c(nrow(sample$x), nrow(sample$y)))
    ord <- sample.int(nrow(x))
    plot(x[ord, ], col = ifelse(first[ord], col_x, col_y),
         pch = ifelse(first[ord], 16, 17), cex = 0.7, asp = 1,
         xlab = "x1", ylab = "x2", main = main)
    legend("bottom", c("first sample (p)", "second sample (q)"),
           col = c(col_x, col_y), pch = c(16, 17), horiz = TRUE, bty = "n",
           inset = c(0, 1), xpd = TRUE, cex = 0.8)
  }
}

## ----plot-location------------------------------------------------------------
draw_samples(loc, "Location shift")

## ----data-scale---------------------------------------------------------------
set.seed(23456)
scl <- list(x = rnorm(n1), y = rnorm(n1, sd = 1.5))
scl_truth <- function(t) dnorm(t, log = TRUE) - dnorm(t, sd = 1.5, log = TRUE)

## ----plot-scale---------------------------------------------------------------
draw_samples(scl, "Scale difference")

## ----controls-----------------------------------------------------------------
boost_ctl <- boost_control(depth = 2)
bart_ctl <- bart_control(num_trees = 100, burnin = 300, draws = 301,
                         keep_draws = TRUE)

## ----fit-1d-------------------------------------------------------------------
set.seed(34567)
loc_boost <- batts_boost(loc$x, y = loc$y, control = boost_ctl, quiet = TRUE)
loc_bart <- batts_bart(loc$x, y = loc$y, control = bart_ctl, quiet = TRUE)
scl_boost <- batts_boost(scl$x, y = scl$y, control = boost_ctl, quiet = TRUE)
scl_bart <- batts_bart(scl$x, y = scl$y, control = bart_ctl, quiet = TRUE)

## ----print--------------------------------------------------------------------
loc_boost
loc_bart

## ----summary------------------------------------------------------------------
summary(loc_boost)

## ----convergence-plot---------------------------------------------------------
old_par <- par(mfrow = c(1, 2), mar = c(4, 4, 2.5, 1))
plot(loc_boost)
plot(loc_bart)
par(old_par)

## ----second-chain-------------------------------------------------------------
set.seed(45678)
loc_bart2 <- batts_bart(loc$x, y = loc$y, control = bart_ctl, quiet = TRUE)
chains_at_0 <- cbind(chain1 = predict(loc_bart, 0, summary = "none")[1, ],
                     chain2 = predict(loc_bart2, 0, summary = "none")[1, ])
check_grid <- seq(-2, 2, length.out = 101)
chain_means <- cbind(chain1 = predict(loc_bart, check_grid),
                     chain2 = predict(loc_bart2, check_grid))

## ----second-chain-plot--------------------------------------------------------
old_par <- par(mfrow = c(1, 2), mar = c(4, 4, 2.5, 1))
matplot(chains_at_0, type = "l", lty = 1, col = c("black", "orange"),
        xlab = "saved draw", ylab = "log(p/q) at x = 0",
        main = "Draws from two chains")
matplot(check_grid, chain_means, type = "l", lty = 1, lwd = 2,
        col = c("black", "orange"), xlab = "x", ylab = "log(p/q)",
        main = "Posterior means")
legend("topright", c("chain 1", "chain 2"), col = c("black", "orange"),
       lty = 1, lwd = 2, bty = "n", cex = 0.8)
par(old_par)

## ----evaluate-1d--------------------------------------------------------------
ratio_curves <- function(sample, boost, bart, truth) {
  grid <- seq(quantile(c(sample$x, sample$y), 0.025),
              quantile(c(sample$x, sample$y), 0.975), length.out = 200)
  list(grid = grid, truth = truth(grid), boost = predict(boost, grid),
       band = predict(bart, grid, interval = TRUE, level = 0.95))
}
loc_curves <- ratio_curves(loc, loc_boost, loc_bart, loc_truth)
scl_curves <- ratio_curves(scl, scl_boost, scl_bart, scl_truth)
head(loc_curves$band, 3)

## ----evaluate-1d-plot---------------------------------------------------------
draw_ratio <- function(curves, main) {
  band <- curves$band
  plot(curves$grid, band[, "fit"], type = "n",
       ylim = range(band, curves$truth, curves$boost),
       xlab = "x", ylab = "log(p/q)", main = main)
  polygon(c(curves$grid, rev(curves$grid)),
          c(band[, "lower"], rev(band[, "upper"])), col = "grey85", border = NA)
  abline(h = 0, col = "grey60")
  lines(curves$grid, curves$truth, col = "red", lwd = 2, lty = 2)
  lines(curves$grid, curves$boost, col = "darkgreen", lwd = 2, lty = 3)
  lines(curves$grid, band[, "fit"], lwd = 2)
}

old_par <- par(mfrow = c(1, 2), mar = c(4, 4, 2.5, 1))
draw_ratio(loc_curves, "Location shift")
legend("topright", c("truth", "boosting", "posterior mean", "95% band"),
       col = c("red", "darkgreen", "black", "grey85"), lty = c(2, 3, 1, 1),
       lwd = c(2, 2, 2, 8), bty = "n", cex = 0.8)
draw_ratio(scl_curves, "Scale difference")
par(old_par)

## ----scales-------------------------------------------------------------------
density_ratio(loc_boost, c(-1, 0, 1))
predict(loc_bart, c(-1, 0, 1), type = "weight")

## ----data-global--------------------------------------------------------------
log_dnorm2 <- function(x, mean, sigma) {
  z <- sweep(x, 2, mean) %*% solve(chol(sigma))
  -0.5 * rowSums(z^2) - log(2 * pi) - 0.5 * log(det(sigma))
}
rnorm2 <- function(n, mean, sigma) {
  sweep(matrix(rnorm(2 * n), n, 2) %*% chol(sigma), 2, mean, "+")
}

sigma_g <- matrix(c(1, 0.5, 0.5, 1), 2, 2)
shift_g <- c(0.75, 0.75)
set.seed(56789)
n2 <- 1000L
glb <- list(x = rnorm2(n2, c(0, 0), sigma_g), y = rnorm2(n2, shift_g, sigma_g))
glb_truth <- function(x) {
  log_dnorm2(x, c(0, 0), sigma_g) - log_dnorm2(x, shift_g, sigma_g)
}

## ----plot-global, fig.height = 4.4--------------------------------------------
draw_samples(glb, "Global shift")

## ----data-local---------------------------------------------------------------
centres <- rbind(c(-2.5, 1), c(1, -2), c(2, 2.5))
rmix <- function(n, centres) {
  k <- sample.int(3L, n, replace = TRUE)
  centres[k, ] + matrix(rnorm(2 * n, sd = 0.3), n, 2)
}
dmix <- function(x, centres) {
  dens <- sapply(1:3, function(k) {
    dnorm(x[, 1], centres[k, 1], 0.3) * dnorm(x[, 2], centres[k, 2], 0.3)
  })
  rowMeans(dens)
}

shifted <- centres
shifted[3, ] <- centres[3, ] + c(-0.3, 0)
set.seed(67890)
lcl <- list(x = rmix(n2, centres), y = rmix(n2, shifted))
lcl_truth <- function(x) log(dmix(x, centres)) - log(dmix(x, shifted))

## ----plot-local, fig.height = 4.4---------------------------------------------
draw_samples(lcl, "Local shift")

## ----fit-2d-------------------------------------------------------------------
bart_ctl_2d <- bart_control(num_trees = 100, burnin = 300, draws = 301)
set.seed(78901)
glb_boost <- batts_boost(glb$x, y = glb$y, control = boost_ctl, quiet = TRUE)
glb_bart <- batts_bart(glb$x, y = glb$y, control = bart_ctl_2d, quiet = TRUE)
lcl_boost <- batts_boost(lcl$x, y = lcl$y, control = boost_ctl, quiet = TRUE)
lcl_bart <- batts_bart(lcl$x, y = lcl$y, control = bart_ctl_2d, quiet = TRUE)

## ----draw-2d------------------------------------------------------------------
diverging <- hcl.colors(101, "Blue-Red 3")
sequential <- hcl.colors(101, "YlOrRd", rev = TRUE)
draw_points <- function(x, value, limits, palette, main) {
  index <- round((pmin(pmax(value, limits[1]), limits[2]) - limits[1]) /
                   diff(limits) * 100) + 1
  ord <- order(abs(value - mean(limits)))
  plot(x[ord, ], col = palette[index[ord]], pch = 16, cex = 0.5, asp = 1,
       xlab = "x1", ylab = "x2", main = main)
}
draw_scenario <- function(sample, boost, bart, truth) {
  x <- rbind(sample$x, sample$y)
  true_value <- truth(x)
  limits <- c(-1, 1) * max(abs(true_value))
  draws <- fitted(bart, summary = "none")
  old_par <- par(mfrow = c(2, 2), mar = c(4, 4, 2.5, 1))
  on.exit(par(old_par))
  draw_points(x, true_value, limits, diverging, "true log(p/q)")
  draw_points(x, fitted(boost), limits, diverging, "boosting")
  draw_points(x, rowMeans(draws), limits, diverging, "posterior mean")
  sds <- apply(draws, 1, sd)
  draw_points(x, sds, c(0, max(sds)), sequential, "posterior SD")
}

## ----ratio-global, fig.height = 6.4-------------------------------------------
draw_scenario(glb, glb_boost, glb_bart, glb_truth)

## ----ratio-local, fig.height = 6.4--------------------------------------------
draw_scenario(lcl, lcl_boost, lcl_bart, lcl_truth)