R/LinearModels_RandomEffects.R

In tommasorigon/PhDPack: Useful script to survive PhD

#' Estimates a Bayesian linear model with random intercept
#'
#' Bayesian linear modeling with random intercept
#' @param formula The formula of the model
#' @param strata A vector including the group membership of each observation
#' @param data The dataset
#' @export

Bayeslm.rintercept <- function(formula, strata, data = NULL, prior = list(Sigma.beta = 10^5, 
     a = 1e-05, b = 1e-05, beta0 = 0), R = 10^4, burn.in = 1000, verbose = TRUE) {
     
     require(spam)
     require(coda)
     
     mf <- model.frame(formula, data)
     y  <- model.response(mf)
     X  <- model.matrix(formula, mf)
     XtX <- crossprod(X)
     
     D <- as.spam(model.matrix(y ~ as.factor(strata) - 1, data = NULL))
     DtD <- crossprod.spam(D, D)
     
     n <- length(y)
     p <- NCOL(X)
     q <- NCOL(D)
     
     ### Prior Specification
     
     Sigma.beta <- diag(prior$Sigma.beta, p)
     beta0 <- rep(prior$beta0, p)
     a <- prior$a
     b <- prior$b
     # Initialization: marginal OLS estimate
     
     # Placing NULL vector and matrix
     
     Beta <- matrix(0, p, R)
     Mu <- matrix(0, q, R)
     Tau <- numeric(R)
     Tau.mu <- numeric(R)
     
     # OLS estimation as initialization
     beta <- solve(XtX) %*% t(X) %*% y
     # 
     mu <- solve(DtD) %*% t(D) %*% (y - X %*% beta)
     # Tau
     tau <- 1/mean((y - X %*% beta - D %*% mu)^2)
     # Tau.mu
     tau.mu <- rgamma(1, a + q/2, b + crossprod(mu)/2)
     
     # Gibbs sampling
     for (i in 1:R) {
          
          # Updating beta
          W <- solve(Sigma.beta) + tau * XtX
          M.beta <- solve(W) %*% (solve(Sigma.beta) %*% beta0 + tau * crossprod(X, 
               y - D %*% mu))
          # Generating Beta
          Beta[, i] <- beta <- t(rmvnorm(1, M.beta, solve(W)))
          
          # Updating mu
          W.mu <- tau.mu + tau * DtD
          M.mu <- solve(W.mu) %*% (tau * t(D) %*% (y - X %*% beta))
          # Generating Mu
          Mu[, i] <- mu <- t(rmvnorm(1, M.mu, solve(W.mu)))
          
          # Updating tau
          residuals <- y - D %*% mu - X %*% beta
          Tau[i] <- tau <- rgamma(1, a + n/2, b + sum(residuals^2)/2)
          
          # Updating tau.mu
          Tau.mu[i] <- tau.mu <- rgamma(1, a + q/2, b + crossprod(mu)/2)
          if (verbose) 
               if (i%%(R/100) == 0) 
                    cat(paste(i, " "))
     }
     burn.in <- min(R/2, burn.in)
     burn.in <- -(1:burn.in)
     list(Beta = as.mcmc(t(Beta[, burn.in])), Mu = as.mcmc(t(Mu[, burn.in])), var.res = as.mcmc(1/Tau[burn.in]), 
          var.ran = as.mcmc(1/Tau.mu[burn.in]))
}


#' Estimates a Bayesian linear
#'
#' Bayesian linear model
#' @param formula The formula of the model
#' @param data The dataset
#' @export

Bayeslm <- function(formula, data = NULL, R = 10^4, burn.in = 1000, verbose = FALSE, 
     scale = FALSE) {
     
     model <- model.frame(formula, data = data)
     y <- model.response(model)
     X <- model.matrix.lm(model)
     if (scale) {
          X[, -1] <- scale(X[, -1])
     }
     
     n <- length(y)
     p <- NCOL(X)
     
     ### Prior Specification
     Sigma.beta <- 10^5 * diag(p)
     beta0 <- rep(0, p)
     a.tau <- b.tau <- 1e-05
     a.beta <- b.beta <- 1e-05
     
     # Fixed quantities
     P <- solve(Sigma.beta) %*% beta0
     XtX <- crossprod(X)
     Xy <- crossprod(X, y)
     
     # Output
     beta.out <- matrix(0, R, p)
     tau.out <- numeric(R)
     predictive.out <- matrix(0, R, n)
     
     # OLS estimation as initialization
     beta <- solve(XtX) %*% t(X) %*% y
     # Tau
     tau <- 1/mean((y - X %*% beta)^2)
     # Tau.mu
     tau.beta <- rgamma(1, a.beta + p/2, b.tau + crossprod(beta)/2)
     
     # Gibbs sampling
     for (r in 1:R) {
          
          # Updating beta
          Sigma.tilde.beta <- solve(solve(Sigma.beta) + tau * XtX)
          beta0.tilde <- Sigma.tilde.beta %*% (P + tau * Xy)
          beta <- t(rmvnorm(1, beta0.tilde, Sigma.tilde.beta))
          
          # Updating tau
          residuals <- y - X %*% beta
          tau <- rgamma(1, a.tau + n/2, b.tau + sum(residuals^2)/2)
          predictive <- X %*% beta + rnorm(n)/sqrt(tau)
          
          beta.out[r, ] <- beta
          tau.out[r] <- tau
          predictive.out[r, ] <- predictive
          
          # Updating tau.mu
          if (verbose) 
               if (r%%(R/100) == 0) 
                    cat(paste(r, " "))
     }
     
     burn.in <- min(R/2, burn.in)
     burn.in <- -(1:burn.in)
     list(beta = beta.out[burn.in, ], tau = tau.out[burn.in], predictive = predictive.out[burn.in, 
          ])
}

#' Estimates a Bayesian linear model with ridge penalization
#'
#' Bayesian linear modeling with ridge penalization
#' @param formula The formula of the model
#' @param data The dataset
#' @export

BayesRidge <- function(formula, data = NULL, R = 10^4, burn.in = 1000, verbose = FALSE, 
     scale = TRUE) {
     
     model <- model.frame(formula, data = data)
     y <- model.response(model)
     X <- model.matrix.lm(model)
     if (scale) {
          X[, -1] <- scale(X[, -1])
     }
     
     n <- length(y)
     p <- NCOL(X)
     
     ### Prior Specification
     Sigma.beta <- 10^5 * diag(p)
     beta0 <- rep(0, p)
     a.tau <- b.tau <- 1e-05
     a.beta <- b.beta <- 1e-05
     
     # Fixed quantities
     XtX <- crossprod(X)
     Xy <- crossprod(X, y)
     
     # Output
     beta.out <- matrix(0, R, p)
     tau.out <- numeric(R)
     tau.beta.out <- numeric(R)
     predictive.out <- matrix(0, R, n)
     
     # Initialization
     beta <- rep(0, p)
     # Tau
     tau <- 1/mean((y - X %*% beta)^2)
     # Tau.mu
     tau.beta <- rgamma(1, a.beta + p/2, b.tau + crossprod(beta)/2)
     
     # Gibbs sampling
     for (r in 1:R) {
          
          # Updating beta
          Sigma.tilde.beta <- solve(solve(Sigma.beta) + tau * XtX)
          beta0.tilde <- Sigma.tilde.beta %*% (tau * Xy)
          beta <- t(rmvnorm(1, beta0.tilde, Sigma.tilde.beta))
          
          # Updating tau
          residuals <- y - X %*% beta
          tau <- rgamma(1, a.tau + n/2, b.tau + sum(residuals^2)/2)
          
          # Updating tau.beta - Attention to not penalize the intercept!
          tau.beta <- rgamma(1, a.beta + (p - 1)/2, b.tau + crossprod(beta[-1])/2)
          diag(Sigma.beta)[-1] <- 1/tau.beta
          
          predictive <- X %*% beta + rnorm(n)/sqrt(tau)
          
          beta.out[r, ] <- beta
          tau.out[r] <- tau
          tau.beta.out[r] <- tau.beta
          predictive.out[r, ] <- predictive
          if (verbose) 
               if (r%%(R/100) == 0) 
                    cat(paste(r, " "))
     }
     
     burn.in <- min(R/2, burn.in)
     burn.in <- -(1:burn.in)
     list(beta = beta.out[burn.in, ], tau = tau.out[burn.in], tau.beta = tau.beta.out[burn.in], 
          predictive = predictive.out[burn.in, ])
}


#' Estimates a Bart model with random intercept
#'
#' Bayesian linear modeling with random intercept
#' @param formula The formula of the model
#' @param strata A vector including the group membership of each observation
#' @param data The dataset
#' @export

Bart.rintercept <- function(formula, strata, data = NULL, prior = list(a = 1e-05, 
     b = 1e-05), R = 10^4, burn.in = 1000, verbose = TRUE) {
     require(coda)
     require(dbarts)
     
     frame <- model.frame(formula, data = data)
     X <- model.matrix(frame, data = data)
     y <- model.response(frame)
     strata <- as.factor(strata)
     n.strata <- as.numeric(table(strata))
     strata.index <- as.numeric(strata)
     
     n <- length(y)
     p <- NCOL(X)
     q <- length(n.strata)
     
     # Values for the prior
     a <- prior$a
     b <- prior$b
     
     # Initialization: marginal OLS estimate
     
     # Placing NULL vector and matrix
     alpha.out <- matrix(0, R, q)  # Output for the random effects
     fit.out <- matrix(0, R, n)  #
     fit.bart.out <- matrix(0, R, n)
     sigma.out <- numeric(R)  # Variance. Replaced by BART routines.
     sigma.alpha.out <- numeric(R)
     
     # Initialization
     mbart <- dbarts(X, y)
     alpha <- tapply(y - mean(y), strata, mean)
     response.alpha <- alpha[strata.index]
     
     # Gibbs sampling
     for (r in 1:R) {
          
          # Updating BART
          mbart$setResponse(y - response.alpha)  # Updating the response
          update.bart <- mbart$run(numBurnIn = 0, numSamples = 1)
          
          # Updating response
          response.bart <- c(update.bart$train)
          # 'Updating' tau
          tau <- 1/update.bart$sigma^2
          
          # Updating tau.mu
          tau.alpha <- rgamma(1, a + q/2, b + crossprod(alpha)/2)
          
          # Updating alpha
          W.alpha <- tau.alpha + tau * n.strata
          alpha.mean <- (tau * tapply(y - response.bart, strata, sum))/W.alpha
          alpha <- rnorm(q, alpha.mean, sqrt(1/W.alpha))
          response.alpha <- c(alpha[strata.index])
          
          if (verbose) 
               if (r%%(R/100) == 0) 
                    cat(paste(r, " "))
          
          alpha.out[r, ] <- alpha
          fit.out[r, ] <- response.bart + response.alpha
          fit.bart.out[r, ] <- response.bart
          sigma.out[r] <- update.bart$sigma
          sigma.alpha.out[r] <- sqrt(1/tau.alpha)
          
     }
     
     burn.in <- min(R/2, burn.in)
     burn.in <- -(1:burn.in)
     
     list(fit = as.mcmc(fit.out[burn.in, ]), fit.bart = as.mcmc(fit.bart.out[burn.in, 
          ]), alpha = as.mcmc(alpha.out[burn.in, ]), sigma = as.mcmc(sigma.out[burn.in]), 
          sigma.alpha = as.mcmc(sigma.alpha.out[burn.in]))
}

#' Estimates a Bart model with random intercept
#'
#' BACT modeling with random intercept
#' @param formula The formula of the model
#' @param strata A vector including the group membership of each observation
#' @param data The dataset
#' @export

Bart.probit.rintercept <- function(formula, strata, data = NULL, prior = list(a = 1e-05, 
     b = 1e-05), R = 10^4, burn.in = 1000, verbose = TRUE) {
     
     require(coda)
     require(dbarts)
     
     frame <- model.frame(formula, data = data)
     X <- model.matrix(frame, data = data)
     y <- model.response(frame)
     strata <- as.factor(strata)
     n.strata <- as.numeric(table(strata))
     strata.index <- as.numeric(strata)
     
     n <- length(y)
     p <- NCOL(X)
     q <- length(n.strata)
     
     # Values for the prior
     a <- prior$a
     b <- prior$b
     
     # Initialization: marginal OLS estimate
     
     # Placing NULL vector and matrix
     alpha.out <- matrix(0, R, q)  # Output for the random effects
     fit.out <- matrix(0, R, n)  #
     fit.bart.out <- matrix(0, R, n)
     sigma.alpha.out <- numeric(R)
     
     # Initialization
     z <- rnorm(n, 0)
     z[y == 0] <- pmin(z[y == 0], 0)
     z[y == 1] <- pmax(z[y == 1], 0)
     alpha <- tapply(z - mean(z), strata, mean)
     response.alpha <- c(alpha[strata.index])
     mbart <- dbarts(X, y)
     
     # Gibbs sampling
     for (r in 1:R) {
          
          # Updating BART
          mbart$setOffset(response.alpha)  # Updating the response
          update.bart <- mbart$run(numBurnIn = 0, numSamples = 1)
          
          # Updating response
          response.bart <- c(update.bart$train) - response.alpha
          
          # Updating latent level
          z <- rnorm(n, response.alpha + response.bart)
          z[y == 0] <- pmin(z[y == 0], 0)
          z[y == 1] <- pmax(z[y == 1], 0)
          
          # Updating tau.mu
          tau.alpha <- rgamma(1, a + q/2, b + crossprod(alpha)/2)
          
          # Updating alpha
          W.alpha <- tau.alpha + n.strata
          alpha.mean <- tapply(z - response.bart, strata, sum)/W.alpha
          alpha <- rnorm(q, alpha.mean, sqrt(1/W.alpha))
          response.alpha <- c(alpha[strata.index])
          
          if (verbose) 
               if (r%%(R/100) == 0) 
                    cat(paste(r, " "))
          
          alpha.out[r, ] <- alpha
          fit.out[r, ] <- response.bart + response.alpha
          fit.bart.out[r, ] <- response.bart
          sigma.alpha.out[r] <- sqrt(1/tau.alpha)
          
     }
     
     burn.in <- min(R/2, burn.in)
     burn.in <- -(1:burn.in)
     
     list(fit = as.mcmc(fit.out[burn.in, ]), fit.bart = as.mcmc(fit.bart.out[burn.in, 
          ]), alpha = as.mcmc(alpha.out[burn.in, ]), sigma.alpha = as.mcmc(sigma.alpha.out[burn.in]))
}