# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .r
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: R
#     language: R
#     name: ir
# ---

library(ggplot2)
library(tibble)
library(microbenchmark)

# ### Data generator

generate_data <- function(n=200, d=100, s=2, seed =42) 
{
    set.seed(seed)
    beta <- numeric(length = d)
    beta [1:s] <- 1
    X <- matrix(rnorm(n*d), nrow=n, ncol=d)
    y <- c(X %*% beta) + rnorm(n)
    list(X=X, y=y, beta=beta)
}
data <- generate_data(n=1000, d=400, s=4, seed =58)

# ### Helpers

linspace <- function(x1, x2, n=100) {
    stopifnot(is.numeric(x1), is.numeric(x2), length(x1)==1, length(x2)==1)
    n <- floor(n)
    if (n <= 1) x2
    else seq(x1, x2, length.out=n)
}
hess.loss.max <- function(X)
{
    hess.loss = 2 * (t(X) %*% X)
    hess.ev = eigen(hess.loss)$values
    return(max(hess.ev))
}
l1.norm <- function(a) {sum(abs(a))}
l2.norm <- function(a) {sqrt(sum(a*a))}
loss.lasso <- function(X, y, beta, lam){0.5*l2.norm(X%*%beta-y)**2+lam*l1.norm(beta)}
backtrack <- function(f, grad, x, eta, k) 
{
    i <- 1
    alpha <- c(1)
    search_vec <- -grad
    t <- -k * t(search_vec)%*%grad
    while (f(x)-f(x+alpha[i]*search_vec)<alpha[i]*t)
    {
        alpha <- c(alpha, eta*alpha[i])
        i <- i+1
    }
    return(tail(alpha, 1))
}

# ### GD steps

gd_step.sub_grad <- function(X, y, beta_, t, lam) {beta_-t* (t(X)%*%(X%*%beta_-y)+lam*sign(beta_))}

prox.l1 <- function(x, lam){sign(x)*pmax(abs(x)-lam, 0)}
gd_step.prox <- function(X, y, beta_, t, lam) {prox.l1(beta_-t*t(X)%*%(X%*%beta_-y), lam*t)}

gd_step.acc_prox <- function(X, y, beta_, beta_prev, t, lam, k)
{
    v <- beta_+(k/(k+3))*(beta_-beta_prev)
    return(prox.l1(v-t*t(X)%*%(X%*%v-y), lam*t))
}

# ### Fitters

fit.lasso <- function(alg, X, y, 
                      beta_true=NULL, lam=3, max_iter=50, early_stop=0, eta=0.1, k=0.3, back_track=FALSE, t=1)
{
    d <- dim(X)[2]
    beta = numeric(length=d)
    betas = list(beta, beta)
    loss <- NULL
    error <- NULL
    iter <- 0
    t<-t/hess.loss.max(X)
    while(iter<max_iter)
    {
        iter <- iter+1
        y_hat <- X%*%beta
        if(back_track)
            {t<-backtrack(function(b){loss.lasso(X, y, b, lam)}, t(X)%*%(X%*%beta-y)+lam*sign(beta), beta, eta, k)}
        if(alg=='prox'){beta <- gd_step.prox(X, y, beta, t, lam)}
        else if(alg=='sub_grad'){beta <- gd_step.sub_grad(X, y, beta, t, lam)}
        else if(alg=='acc_prox')
        {
            betas <- append(betas, list(gd_step.acc_prox(X, y, rev(betas)[[1]], rev(betas)[[2]], t, lam, iter)))
            if(length(betas)>2){betas<-betas[-1]}
            beta <- rev(betas)[[1]]
        }
        else {stop('invalid algorithm')}
        loss <- c(loss, loss.lasso(X, y, beta, lam))
        if(!is.null(beta_true))
        {
            log_loss <- log(tail(loss, 1))-log(loss.lasso(X, y, beta_true, lam))
            error <- c(error, log_loss)
        }
        if(iter==early_stop){break}
    }
    return(list(y_pred=X%*%beta, beta=beta, loss=loss/length(y),  error=error))
}

data <- generate_data(n=400, d=400, s=300, seed =58)
model.prox <- fit.lasso('prox', data$X, data$y, data$beta)
model.sub_grad <- fit.lasso('sub_grad', data$X, data$y, data$beta)
model.acc_prox <- fit.lasso('acc_prox', data$X, data$y, data$beta)

prox.overfit <- which.max(model.prox$error<0)
sub_grad.overfit <- which.max(model.sub_grad$error<0)
acc_prox.overfit <- which.max(model.acc_prox$error<0)

df <- tibble(n.iter = 1:length(model.prox$error) ,
prox.error = model.prox$error ,
sub_grad.error = model.sub_grad$error,
acc_prox.error = model.acc_prox$error)

# pdf('30conv4.pdf', width=10, height =6)
ggplot(data=df, aes(x=n.iter)) +
geom_line(aes(y=prox.error , color='proximal'), lwd=1.5) +
geom_line(aes(y=sub_grad.error , color='subgradient'), lwd=1.5) +
geom_line(aes(y=acc_prox.error, color='accelerated proximal'), lwd=1.5) +
ylab('Log error') +
xlab('epoch ') +
theme_bw() +
theme(legend.position='bottom', text = element_text(size = 16))
dev.off()

L <- hess.loss.max(data$X)
1/L

n.reps <- 1000
bench <- summary(microbenchmark(
    gd_step.prox(data$X, data$y, numeric(length=dim(data$X)[2]), 0.01, 1),
    gd_step.sub_grad(data$X, data$y, numeric(length=dim(data$X)[2]), 0.01, 1),
    gd_step.acc_prox(data$X, data$y, numeric(length=dim(data$X)[2]), numeric(length=dim(data$X)[2])+1/dim(data$X)[2],0.01, 0.1, 2),
    unit='ms',
    times=n.reps))

print(tibble(iter=bench$mean ,
overfit=c(prox.overfit , sub_grad.overfit, acc_prox.overfit),
total=c(prox.overfit *bench$mean[1],
sub_grad.overfit * bench$mean [2], acc_prox.overfit *bench$mean[3])))





StarcoderdataJupyterScriptsDedupFiltered