MCMC Sampling for a Logistic Regression Model

# packages
if (!require("pacman")) utils::install.packages("pacman", dependencies = TRUE)

pacman::p_load(
    here,
    stats, cumstats,
    tidyverse, 
    skimr, glue, scales, gt,
    broom, 
    TeachingDemos
)

# helper function: string concatenation (ex: "act" %&% "5")
'%&%' <- function(x, y) paste0(x, y)

# not %in%
'%notin%' <- Negate('%in%')

# summarize a distribution (graphically and numerically)
dist_summary <- function(
    x, var_expression,
    pos_x = 0, pos_y = 1, adj = c(0, 1),
    conf.level = 0.95, digits = 4, plot = T, ...){
  
  # significance level
  alpha <- (1 - conf.level)/2
  
  # mean, median, and intervals
  mean <- mean(x)
  median <- median(x)
  sd <- sd(x)
  CI <- quantile(x, probs = c(alpha, 1 - alpha)) %>% `names<-`(vec_fmt_percent(c(alpha, 1 - alpha), decimals = 1))
  CI_hpd <- TeachingDemos::emp.hpd(x) %>% `names<-`(vec_fmt_percent(c(alpha, 1 - alpha), decimals = 1))

  if(plot){
    # plot distributions with annotations of mean, median and CIs
    hist_obj <- hist(
      x, 
      probability = T, 
      col = NULL,
      yaxt = "n",
      xlab = var_expression, 
      main = "Histogram")
    
    # add lines
    abline(v = mean, col = "red", lwd = 2)
    abline(v = median, col = "blue", lwd = 2)
    abline(v = CI)
    abline(v = CI_hpd, lty = 2, col = "green")
    
    # annotation
    text(
      quantile(hist_obj$breaks, probs = pos_x),
      quantile(hist_obj$density, probs = pos_y),
      glue(
        "Mean = {signif(mean, digits)}\n",
        "Median = {signif(median, digits)}\n",
        "Std = {signif(sd, digits)}\n",
        "{conf.level*100}% credible interval:\n",
        "  Equal-tail = ({toString(signif(CI, digits))})\n",
        "  HPD = ({toString(signif(CI_hpd, digits))})\n"
      ),
      adj = adj, ...)
  }
  
  # return summaries
  return(tibble(
    mean = mean,
    median = median,
    sd = sd,
    CI = list(CI),
    CI_hpd = list(CI_hpd)
  ))
}

diabetes <- read_csv("diabetes-dataset.csv", show_col_types = F)

# skim the data 
skim(diabetes)

# count table for the response variable
table(diabetes$Outcome)

# calculate p
expit <- function(X, beta0, beta1) {
  return( 1/(1 + exp(-beta0 - beta1*X)) )
}

# log prior for beta
logprior <- function(beta, mean = 0, sd = 15) {
  return( -0.5*(log(sd^2) + (beta - mean)^2/sd^2) )
}

# log posterior
logpost <- function(Y, X, beta0, beta1){
  P <- expit(X, beta0, beta1)
  return( sum(Y*log(P) + (1-Y)*log(1-P)) + logprior(beta0) + logprior(beta1) )
}

# MCMC MH sampler - univariate manner
mcmc_mh <- function(
    Y, X, N = 1000 * 10, 
    beta0_init = 0, proposed_sd0 = 0.01, 
    beta1_init = 0, proposed_sd1 = 0.01, 
    seed = 2022){
  
  # make the simulation reproducible
  if(!is.null(seed)) set.seed(seed)
  
  # start timer
  proc_timer <- proc.time()
  
  # set up a data frame to store beta0 and beta1
  df <-  tibble(
    beta0 = beta0_init, 
    beta1 = beta1_init, 
    proc_time = (proc.time() - proc_timer)["elapsed"],
    accept = 0
  )
  
  # current step - initialize
  beta0 <- beta0_init
  beta1 <- beta1_init
  
  # log posterior for the current step
  l <- logpost(Y, X, beta0, beta1)
  
  # iteration
  for (i in 2:N) {
    # reset accept score
    accept <- 0
    
    # beta0: pick a candidate 
    beta0_new <- rnorm(1, beta0, proposed_sd0)
    # new log posterior
    l_new <- logpost(Y, X, beta0_new, beta1)
    # decide whether to accept or reject
    lnR <- l_new - l
    lnU <- log(runif(1))
    if(lnR > lnU){ 
      beta0 <- beta0_new 
      l <- l_new
      accept <- accept + 0.5
    }
    
    # beta1: pick a candidate 
    beta1_new <- rnorm(1, beta1, proposed_sd1)
    # new log posterior
    l_new <- logpost(Y, X, beta0, beta1_new)
    lnR <- l_new - l
    lnU <- log(runif(1))
    if(lnR > lnU){ 
      beta1 <- beta1_new 
      l <- l_new
      accept <- accept + 0.5
    }
    
    # save betas and elapsed time
    df <- df %>% add_row(
      beta0 = beta0, 
      beta1 = beta1,
      proc_time = (proc.time() - proc_timer)["elapsed"],
      accept = accept
    )
  }
  
  print(glue(
    "Last values: (beta0, beta1) = ({beta0}, {beta1})\n",
    "Process time [s]: {tail(df$proc_time, 1)}\n",
    "Acceptance rate: {mean(df$accept)}"
  ))
  
  # return all betas
  return(df)
}

mcmc_summary <- function(df, burn_in, ...){
  # add cumulative mean of acceptance (acceptance rate)
  df <- mutate(df, cum_accept_rate = cummean(accept))
  
  N <- nrow(df)
  proc_time <- tail(df$proc_time, 1)
  acceptance_rate <- tail(df$cum_accept_rate, 1)
  
  # plot algorithm performance
  par(mfrow = c(1,3))
  # plot 1: trace
  plot(df$beta0, df$beta1, type = "l", xlab = expression(beta[0]), ylab = expression(beta[1]))
  text(df$beta0[1], df$beta1[1], "Start", col = "blue")
  text(tail(df$beta0, 1), tail(df$beta1, 1), "End", col = "red")
  # plot 2: process time
  plot(df$proc_time, type = "l", xlab = "Iteration", ylab = "Process Time [s]")
  text(N/2, proc_time, round(proc_time, 2) %&% " [s]", col = "red")
  # plot 3: acceptance rate
  plot(df$cum_accept_rate, type = "l", xlab = "Iteration", ylab = "Acceptance Rate")
  text(N*0.9, acceptance_rate, vec_fmt_percent(acceptance_rate, 2), col = "red")
  
  # remove burn-in rows
  df2 <- tail(df, - burn_in)

  # plot distributions of betas with annotations of mean, median and CIs
  par(mfrow = c(1,2))
  beta0 <- dist_summary(df2$beta0, var_expression = expression(beta[0]), ...)
  beta1 <- dist_summary(df2$beta1, var_expression = expression(beta[1]), ...)
  
  # reset par options
  par(mfrow = c(1,1))
  
  # combine betas and report intervals, process time and acceptance rate
  bind_rows(list(beta0 = beta0, beta1 = beta1), .id = "term") %>% 
    # unpack list columns for intervals 
    unnest_wider(c(CI, CI_hpd), names_sep = "_") %>% 
    mutate(
      proc_time = proc_time,
      acceptance_rate = acceptance_rate
    )
}

# simulation setup
N <- 1000 * 100
burn_in <- N/10

# MCMC MH  - univariate manner
df_univar <- mcmc_mh(
  Y = diabetes$Outcome, X = diabetes$Glucose, N = N,
  beta0_init = 0, proposed_sd0 = 0.01, 
  beta1_init = 0, proposed_sd1 = 0.001) %>% 
  mcmc_summary(burn_in)

df_univar

# fit a GLM logistic regression model
proc_timer <- proc.time()
fit <- glm(Outcome ~ Glucose, family = binomial(link = 'logit'), data = diabetes)
proc_time <- proc.time() - proc_timer

# summary           
summary(fit)

# estimate, SE and 95% CI
df_glm <- tidy(fit, conf.int = T) %>% 
  mutate(proc_time = proc_time["elapsed"])

df_glm

# MCMC MH sampler - bivariate manner
mcmc_mh2 <- function(
    Y, X, N = 1000 * 10, 
    beta0_init = 0, proposed_sd0 = 0.01, 
    beta1_init = 0, proposed_sd1 = 0.01, 
    seed = 2022){
  
  # make the simulation reproducible
  if(!is.null(seed)) set.seed(seed)
  
  # start timer
  proc_timer <- proc.time()
  
  # set up a data frame to store beta0 and beta1
  df <-  tibble(
    beta0 = beta0_init, 
    beta1 = beta1_init, 
    proc_time = (proc.time() - proc_timer)["elapsed"],
    accept = 0
  )
  
  # current step - initialize
  beta0 <- beta0_init
  beta1 <- beta1_init
  
  # log posterior for the current step
  l <- logpost(Y, X, beta0, beta1)
  
  # iteration
  for (i in 2:N) {
    # reset accept score
    accept <- 0
    
    # pick candidates 
    beta0_new <- rnorm(1, beta0, proposed_sd0)
    beta1_new <- rnorm(1, beta1, proposed_sd1)
    # candidate's log posterior
    l_new <- logpost(Y, X, beta0_new, beta1_new)
    # decide whether to accept or reject
    lnR <- l_new - l
    lnU <- log(runif(1))
    if(lnR > lnU){ 
      beta0 <- beta0_new 
      beta1 <- beta1_new 
      l <- l_new
      accept <- 1
    }
    
    # save betas and elapsed time
    df <- df %>% add_row(
      beta0 = beta0, 
      beta1 = beta1,
      proc_time = (proc.time() - proc_timer)["elapsed"],
      accept = accept
    )
  }
  
  print(glue(
    "Last values: (beta0, beta1) = ({beta0}, {beta1})\n",
    "Process time [s]: {tail(df$proc_time, 1)}\n",
    "Acceptance rate: {mean(df$accept)}"
  ))
  
  # return all betas
  return(df)
}

# MCMC MH - bivariate manner
df_bivar <- mcmc_mh2(
  Y = diabetes$Outcome, X = diabetes$Glucose, N = N,
  beta0_init = 0, proposed_sd0 = 0.01, 
  beta1_init = 0, proposed_sd1 = 0.001) %>% 
  mcmc_summary(burn_in)

df_bivar

# A naive way to compare efficiency of process time usage
glue(
  "Process Time Wasted [s] := Process Time * (1 - Acceptance Rate)\n",
  "MCMC-MH univariate: {round(df_univar$proc_time[1] * (1 - df_univar$acceptance_rate[1]), 2)} [s]\n",
  "MCMC-MH bivariate:  {round(df_bivar$proc_time[1] * (1 - df_bivar$acceptance_rate[1]), 2)} [s]"
)

# combine all results for comparison
bind_rows(
  # from glm
  df_glm %>% 
    rename(
      "sd" = "std.error", 
      `CI_2.5%` = conf.low, 
      `CI_97.5%` = conf.high) %>% 
    mutate(
      term = if_else(term == "(Intercept)", "beta0", "beta1"),
      method = "MLE"
    ),
  # from MCMC: univariate and bivariate
  df_univar %>% mutate(method = "MCMC_MH_univar"), 
  df_bivar %>% mutate(method = "MCMC_MH_bivar")
  ) %>% 
  # calculate width of intervals
  mutate(CI_width = `CI_97.5%` - `CI_2.5%`) %>% 
  # select columns of interest
  select(term, method, estimate, median, sd, `CI_2.5%`, `CI_97.5%`, CI_width, proc_time, acceptance_rate) %>% 
  # make it easy to read
  mutate(across(where(is.numeric), ~ round(.x, 4))) %>% 
  arrange(term)