R/ridge_reg.R

In casxue/bis557: Coursework for BIS 557

#' Fit a ridge regression model
#'
#' @description This function fits a ridge regression model.
#' @param formula a formula
#' @param lambda a numeric value
#' @param data a data.frame
#' @return a ridge_reg object
#' @importFrom stats model.frame model.matrix 
#' @examples
#' fit <- ridge_reg(Sepal.Length ~ . , 1.5, iris)
#' fit$coef
#' @export
ridge_reg <- function (formula, lambda, data) {
    y <- model.frame(formula, data)[,1]
    X <- model.matrix(formula, data)
    
    ## rvars <- colnames(X)
    ## omit <- rownames(alias(formula, data)$Complete)
    ## X <- X[, setdiff(rvars, omit)]

    svd_obj <- svd(X)
    U <- svd_obj[["u"]]
    V <- svd_obj[["v"]]
    svals <- svd_obj[["d"]]

    D <- diag(svals / (svals^2 + lambda))
    beta_hat <- V %*% D %*% t(U) %*% y
    
    ## names(coeff) <- setdiff(rvars, omit)
    ## beta_hat <- sapply(rvars, function (x) {coeff[match(x, names(coeff))]})
    ## names(beta_hat) <- rvars

    names(beta_hat) <- colnames(X)
    beta_hat <- list(coef = beta_hat)
    class(beta_hat) <- "ridge_reg"

    beta_hat
    
    return(beta_hat)
}