R/elnet_coord.R

In DavidJamesKent/stsci6520_hw2: STSCI6520 HW2

#' Elastic Net
#' @description Computes elastic net estimate of coefficients. Assumes columns of \code{X} are standardized.
#'
#' @param y Dependent variable.
#' @param X Design matrix
#' @param alpha,lambda Parameters in objective.
#' @return Vector of length \code{p=ncol(X)} containing elastic net estimates for the coefficients.
#' @export
#'
elnet_coord <- function(y,X,alpha,lambda) {

  # helper
  sthresh <- function(x,t)
    ifelse(abs(x)<=t,0,sign(x)*(abs(x)-t))

  p <- ncol(X)
  n <- nrow(X)
  hbeta <- rep(0,p)
  change <- Inf
  while( change >= 1e-5 ) {
    prev_est <- hbeta
    for(j in 1:p) {
      r <- y - (X[,-j] %*% hbeta[-j]) #partial resid
      # update
      hbeta[j] <- sthresh((X[,j] %*% r)/n,lambda*alpha/2)/( lambda*(1-alpha) + 1 )
    }
    change <- sum( (prev_est - hbeta)^2 )
    if( change > 1e3 ) # Abort if things go south
      return(rep(NA,p))
  }
  return(hbeta)
}