source('http://myweb.uiowa.edu/pbreheny/7600/s16/notes/fun.R')

###############################
#### Slides 16, 17, and 23 ####
###############################
scad <- function(theta,l,a=3.7) {
  theta <- abs(theta)
  (theta <= l) * theta * l + (theta > l & theta <= a * l) * (a*l*theta - (theta^2+l^2)/2)/(a-1) + (theta > a*l) * (l^2*(a^2-1))/(2*(a-1))
}
dSCAD <- function(theta,l,a=3.7) {
  theta = abs(theta)
  (theta < l)*l + ((theta > l) & (theta < a*l))*((a*l-theta)/(a-1))
}
Fs <- function(z,l,g=3.7){
  (abs(z) > g*l)*z + (abs(z) > 2*l & abs(z) <= g*l)*(S(z,g*l/(g-1))/(1-1/(g-1))) + (abs(z) <= 2*l)*S(z,l)
}
mcp <- function(theta,l,a=3.7) {
  T <- length(theta)
  val <- numeric(T)
  for (i in 1:T) {
    x <- abs(theta[i])
    val[i] <- (x < a*l)*(l*x - x^2/(2*a)) + (x >= a*l)*(1/2)*a*l^2
  }
  val
}
dMCP <- function(theta,l,a=3.7) {
  theta = abs(theta)
  (theta <= a*l) * (l - theta/a)
}
Fm <- function(z,l,g=3.7){
  (abs(z) > g*l)*z + (abs(z) <= g*l)*(S(z,l)/(1-1/g))
}
lasso <- function(theta,l){l*abs(theta)}
dLasso <- function(theta,l){rep(l,length(theta))}
S <- function(z,l) {
  (abs(z) > l)*(z-l*sign(z))
}

res <- 101
x <- seq(-4,4,len=res)
xx <- seq(0.0001,4,len=res)
g <- 3
col <- c("#FF4E37FF", "#00B500FF", "#008DFFFF")

Y <- cbind(lasso(x, 1), mcp(x, 1, g), scad(x, 1, g))
matplot(x, Y, type='l', bty='n', lty=1, lwd=3, col=col, las=1,
        xlab=expression(beta), ylab=expression(P(beta*'|'*lambda,gamma)))
text(2.3, 3.5, "Lasso")
text(3.8, 2.4, "SCAD")
text(3.8, 1.1, "MCP")

Y <- cbind(dLasso(xx, 1), dMCP(xx, 1, g), dSCAD(xx, 1, g))
matplot(xx, Y, type='l', bty='n', lty=1, lwd=3, col=col, las=1,
        xlab=expression(beta), ylab=expression(dot(P)(beta*'|'*lambda,gamma)))
text(3.5, 0.9, "Lasso")
text(2.25, 0.7, "SCAD")
text(0.9, 0.48, "MCP")

Y <- cbind(S(x,1), Fm(x,1,g), Fs(x,1,g))
matplot(x, Y, type='l', bty='n', lty=1, lwd=3, col=col, las=1,
        xlab="z", ylab=expression(hat(beta)))
lines(c(-4,4),c(-4,4), col="gray70", lwd=1)
text(3.7, 1.2, "Lasso")
text(3, 4.4, "SCAD", xpd=1)
text(1.4, 2.1, "MCP")

######################
#### Slides 25-29 ####
######################

# Gen data
set.seed(105)
n <- 200; p <- 1000
X <- matrix(rnorm(n*p), nrow=n, ncol=p)
z1 <- rnorm(n); z2 <- rnorm(n)
X[,1:4]   <- X[,1:4]+z1
X[,5]     <- X[,5]+2*z1
X[,6]     <- X[,6]+1.5*z1
X[,7:20]  <- X[,7:20]+0.5*z1
X[,21:40] <- X[,21:40]+0.5*z2
beta <- c(4, 2, -4,-2, rep(0, 996))
y <- rnorm(n, X%*%beta, sd=1.5)
require(ncvreg)

# Lasso
ylim <- c(-4,4)
fit <- ncvreg(X, y, penalty='lasso', lambda.min=0.002)
plot(fit, col=pal(4), lwd=2, shade=FALSE, bty='n', ylim=ylim)

# Adaptive lasso
fit <- ncvreg(X, y, penalty='lasso', lambda.min=0.002)
beta <- fit$beta
L <- length(fit$lambda)
for (l in 1:L) {
  w <- pmin(1e6, 1/abs(fit$beta[-1,l]))
  fit.al <- ncvreg(X, y, lambda=fit$lambda[1:l], penalty='lasso', penalty.factor=w)
  beta[,l] <- fit.al$beta[,l]
}
fit$beta <- beta
nz <- which(apply(beta[-1,], 2, function(x) any(x!=0)))
xlim <- c(fit$lambda[nz[1]-1], fit$lambda[100])
plot(fit, col=pal(4), lwd=2, shade=FALSE, xlim=xlim, bty='n', ylim=ylim)

# MCP
fit <- ncvreg(X, y, gamma=3)
plot(fit, col=pal(4), lwd=2, shade=FALSE, bty='n', ylim=ylim)

# SCAD
fit <- ncvreg(X, y, gamma=4, penalty='SCAD')
plot(fit, col=pal(4), lwd=2, shade=FALSE, bty='n', ylim=ylim)

##################
#### Slide 33 ####
##################
S <- function(z,l) {
  (abs(z) > l)*(z-l*sign(z))
}
Fm <- function(z,l,g=3.7){
  (abs(z) > g*l)*z + (abs(z) <= g*l)*(S(z,l)/(1-1/g))
}
Fs <- function(z,l,g=3.7){
  (abs(z) > g*l)*z + (abs(z) > 2*l & abs(z) <= g*l)*(S(z,g*l/(g-1))/(1-1/(g-1))) + (abs(z) <= 2*l)*S(z,l)
}

# MCP
gam1 <- seq(1+1e-6, 9, len=99)
G <- length(gam1)
Bias <- Var <- numeric(G)
for (i in 1:G) {
  N <- 100000
  z <- rnorm(N, mean=1, sd=0.6)
  b <- Fm(z, 1, gam1[i])
  Bias[i] <- mean(b) - 1
  Var[i] <- var(b)
}
MCP <- cbind(Bias^2, Var, Bias^2 + Var)

# SCAD
gam2 <- seq(2+1e-6, 10, len=99)
G <- length(gam2)
Bias <- Var <- numeric(G)
for (i in 1:G) {
  N <- 100000
  z <- rnorm(N, mean=1, sd=0.6)
  b <- Fs(z, 1, gam2[i])
  Bias[i] <- mean(b) - 1
  Var[i] <- var(b)
}
SCAD <- cbind(Bias^2, Var, Bias^2 + Var)

par(mfrow=c(1,2), mar=c(5,5,1.5,0.5), oma=c(0,0,2,0))
matplot(gam1, MCP, type='l', lwd=3, col=pal(3), lty=1, las=1, ylab='', xlab=expression(gamma), bty='n', xaxt='n')
axis(1, at=seq(1,9,2))
mtext('MCP', line=0.5)
matplot(gam2, SCAD, type='l', lwd=3, col=pal(3), lty=1, las=1, ylab='', xlab=expression(gamma), bty='n')
mtext('SCAD', line=0.5)
toplegend(legend=c(expression("Bias"^2), "Var", "MSE"), lwd=3, col=pal(3))

##################
#### Slide 34 ####
##################

# Gen data
set.seed(105)
n <- 200; p <- 1000
X <- matrix(rnorm(n*p), nrow=n, ncol=p)
z1 <- rnorm(n); z2 <- rnorm(n)
X[,1:4]   <- X[,1:4]+z1
X[,5]     <- X[,5]+2*z1
X[,6]     <- X[,6]+1.5*z1
X[,7:20]  <- X[,7:20]+0.5*z1
X[,21:40] <- X[,21:40]+0.5*z2
beta <- c(4, 2, -4,-2, rep(0, 996))
y <- rnorm(n, X%*%beta, sd=1.5)
require(ncvreg)

par(mfrow=c(1,3), mar=c(4, 4, 2, 0), cex=1)
ylim <- c(-4,4)
col <- c("#FF4E37FF", "#5FA600FF", "#00C1C9FF", "#D63EFFFF")

# 1.2
fit <- ncvreg(X, y, gamma=1.5)
plot(fit, col=col, lwd=2, shade=FALSE, bty='n', ylim=ylim)
mtext(expression(gamma==1.5), line=0.5)

# 2.7
fit <- ncvreg(X, y, gamma=2.7)
plot(fit, col=col, lwd=2, shade=FALSE, bty='n', ylim=ylim)
mtext(expression(gamma==2.7), line=0.5)

# 6
fit <- ncvreg(X, y, gamma=6)
plot(fit, col=col, lwd=2, shade=FALSE, bty='n', ylim=ylim)
mtext(expression(gamma==6), line=0.5)
