[R] ADMM for lasso

This is a short code for studying admm for lasso.

#--------------------------- functions to be used ----------------------

## the main function f = g + h
f <- function(x, A, b, lambda){
        1/2*norm(A %*% x - b, "2")^2 + lambda*sum(abs(x))
}
## smooth function g
g <- function(x, A, b){
        1/2*norm(A %*% x - b, "2")^2
}

## proximal function of none smooth function h
prox_L1 <- function(x, lambda){
        sign(x) * pmax(abs(x) - lambda, 0)
}

#------------------------------------- generate data --------------------------------------

# dimension
n <- 500                ## number of samples
p <- 1000               ## number of features
s <- 100                ## number of inportant features
## coefficients
beta <- c(runif(s, -2, 2), rep(0, p- s))
A <- matrix(rnorm(p*p), p, p)
A <- scale(A, TRUE, TRUE)
b <-  A %*% beta + 0.1*rnorm(n)
# tunning paramter
lambda <- 0.1*max(abs(t(A) %*% b))

# term will be used
solATA <- solve(t(A) %*% A + rho*diag(p))
ATb <- t(A) %*% b

#----------------------------------- ADMM ------------------------------------------------
# initialization
x <- y <- z <- rep(0, p)

maxiter <- 4000
rho <- 0.5
rrs <- rep(0, maxiter)
for(i in 1:maxiter){
        # store temp y
        y_old <- y

        # update x
        x <- solATA %*% (ATb + rho* (y-z))

        # update y
        y <- prox_L1(x + z, lambda/rho)

        # update z
        z <- z + x - y

        # residual
        rr <- norm(x - y, "2")
        ss <- rho*norm(y - y_old, "2")

        cat(rr, ss, fill = TRUE)

        rrs[i] <- rr
}

#------------------------ accelerated ADMM ---------------------------------------------------
# initialization
x <- y <- z <- rep(0, p)
hat_y <- hat_z <- rep(0, p)

maxiter <- 4000
rho <- 0.5
rrs_fast <- rep(0, maxiter)

alpha <- c(1, rep(0, maxiter - 1))
for(i in 1:maxiter){
        # store temp z
        z_old <- z
        y_old <- y

        # update x
        x <- solATA %*% (ATb + rho* (hat_y-hat_z))

        # update y
        y <- prox_L1(x + hat_z, lambda/rho)

        # update z
        z <- hat_z + x - y

        # residual
        rr <- norm(x - y, "2")
        ss <- rho*norm(y - y_old, "2")

        cat(rr, ss, fill = TRUE)

        rrs_fast[i] <- rr

        # accelerated stage
        alpha[i+1] <- (1+sqrt(1+4*alpha[i]))/2
        hat_y <- y + (alpha[i] - 1)/alpha[i+1] * (y - y_old)
        hat_z <- z + (alpha[i] - 1)/alpha[i+1] * (z - z_old)
}

#------------------------------------ plot -----------------------------------------
library(ggplot2)
data <- data.frame(iter=seq(maxiter), rrs=rrs, rrs_fast=rrs_fast)
# gather data (gather in tidyr) or melt data (melt in reshape2)
library(reshape2)
test_data <- melt(data, id = "iter")
windows()
ggplot(test_data, aes(x=iter, y=value, colour=variable)) + geom_line()