---
title: "Randomized LASSO for a Gaussian linear model"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Randomized LASSO for a Gaussian linear model}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE}
knitr::opts_chunk$set(collapse = TRUE, comment = "#>",
                      eval = requireNamespace("glmnet", quietly = TRUE))
```

This vignette shows the basic workflow of `lassoinf`. We fit a LASSO to a randomized
response, then use `LassoInference` for confidence intervals and p-values for the
selected coefficients that are valid after selection.

## Randomization and the proportion used for selection

In data splitting, a proportion $\pi$ of the data is used to select the model and the rest
for inference. For the score $Z = X'y$, with variance $\Sigma = \sigma^2 X'X$, the score of
a $\pi$ fraction of the data has variance $\Sigma / \pi$.

Randomization instead selects with $Z_{noisy} = Z + \omega$, $\omega \sim N(0, \Omega)$. Matching
$\text{Var}(Z_{noisy}) = \Sigma + \Omega = \Sigma / \pi$ gives

$$
\Omega = \frac{1 - \pi}{\pi} \Sigma.
$$

For least squares, this is the same as fitting the LASSO to $y_{noisy} = y + \epsilon$ with
$\epsilon \sim N(0, \frac{1-\pi}{\pi}\sigma^2 I)$. As $\pi \to 1$ almost all the information is
used for selection and little is left for inference. Smaller $\pi$ leaves more.

## Data

We use $n = 100$ and $p = 20$, with one strong effect, two weaker ones and $\sigma^2 = 2$.
Selection uses $\pi = 0.9$.

```{r}
library(lassoinf)
set.seed(0)

n <- 100; p <- 20
signal_strength <- 4
X <- matrix(rnorm(n * p), n, p)
beta <- numeric(p)
beta[1:3] <- c(2, 1, -1) * signal_strength / sqrt(n)
sigma2 <- 2
y <- drop(X %*% beta) + rnorm(n, sd = sqrt(sigma2))
```

With $n > p$ we estimate $\sigma^2$ from the full least squares fit, and build $\Sigma$
and $\Omega$:

```{r}
pi_selection <- 0.9
Q <- crossprod(X)
Z_full <- drop(crossprod(X, y))
sigma2_hat <- sum(lm.fit(X, y)$residuals^2) / (n - p)
Sigma <- sigma2_hat * Q
sigma2_rand <- (1 - pi_selection) / pi_selection * sigma2_hat
Sigma_noise <- sigma2_rand * Q
y_noisy <- y + rnorm(n, sd = sqrt(sigma2_rand))
c(sigma2_hat = sigma2_hat, sigma2_rand = sigma2_rand)
```

## Fit the LASSO on the randomized response

We fit the LASSO $\frac{1}{2}\|y_{noisy} - Xb\|^2 + \lambda \|b\|_1$. glmnet scales the loss by
$1/n$, so we pass it $\lambda / n$. `LassoInference` needs the solution, the gradient of
the smooth part at it, $G = X'(X\hat{\beta} - y_{noisy})$, and its Hessian, $Q = X'X$.

```{r}
lam <- 2 * sqrt(n)
fit <- glmnet::glmnet(X, y_noisy, lambda = lam / n, intercept = FALSE,
                      standardize = FALSE, control = list(thresh = 1e-14))
beta_hat <- as.numeric(coef(fit))[-1]
G_hat <- drop(crossprod(X, X %*% beta_hat - y_noisy))
which(beta_hat != 0)
```

## Inference

The arguments are the solution, gradient and Hessian, the penalty weights $D$, the bounds
$L$, $U$ (none here), the full-data score and its covariance, and the randomization
covariance.

```{r}
li <- LassoInference$new(beta_hat, G_hat, Q, rep(lam, p), NULL, NULL,
                         Z_full, Sigma, Sigma_noise, level = 0.9)
summ <- li$summary()
summ
```

`index` is 0-based. `beta_hat` is the LASSO solution, which is shrunk toward zero and can
lie outside its interval. The intervals are for the coefficients
of the selected model, $\beta_E = (X_E'X_E)^{-1} X_E' \mu$ with $\mu = E[y]$. They come
from the contrasts in `li$contrasts`, so we can compute the targets and check coverage:

```{r}
Z_mu <- drop(Q %*% beta)
summ$truth <- sapply(as.character(summ$index), function(j) sum(li$contrasts[[j]] * Z_mu))
summ$cover <- summ$lower_conf <= summ$truth & summ$truth <= summ$upper_conf
summ
```

`summary()` also stores naive intervals, which ignore selection, and data-splitting intervals
in `li$"_naive"` and `li$"_splitting"`:

```{r}
li$`_naive`
li$`_splitting`
```

## Coverage over repeated experiments

The randomized intervals should cover their targets 90% of the time. The naive intervals
should not.

```{r}
one_replication <- function() {
  X <- matrix(rnorm(n * p), n, p)
  y <- drop(X %*% beta) + rnorm(n, sd = sqrt(sigma2))
  Q <- crossprod(X)
  sigma2_hat <- sum(lm.fit(X, y)$residuals^2) / (n - p)
  sigma2_rand <- (1 - pi_selection) / pi_selection * sigma2_hat
  y_noisy <- y + rnorm(n, sd = sqrt(sigma2_rand))
  fit <- glmnet::glmnet(X, y_noisy, lambda = lam / n, intercept = FALSE,
                        standardize = FALSE, control = list(thresh = 1e-14))
  beta_hat <- as.numeric(coef(fit))[-1]
  if (all(beta_hat == 0)) return(NULL)
  G_hat <- drop(crossprod(X, X %*% beta_hat - y_noisy))
  li <- LassoInference$new(beta_hat, G_hat, Q, rep(lam, p), NULL, NULL,
                           drop(crossprod(X, y)), sigma2_hat * Q, sigma2_rand * Q, level = 0.9)
  summ <- li$summary()
  truth <- sapply(as.character(summ$index), function(j) sum(li$contrasts[[j]] * drop(Q %*% beta)))
  covers <- function(df) df$lower_conf <= truth & truth <= df$upper_conf
  data.frame(method = rep(c("randomized", "naive", "splitting"), each = nrow(summ)),
             cover = c(covers(summ), covers(li$`_naive`), covers(li$`_splitting`)),
             length = c(summ$upper_conf - summ$lower_conf,
                        li$`_naive`$upper_conf - li$`_naive`$lower_conf,
                        li$`_splitting`$upper_conf - li$`_splitting`$lower_conf))
}

set.seed(1)
results <- do.call(rbind, replicate(200, one_replication(), simplify = FALSE))
data.frame(coverage = tapply(results$cover, results$method, mean),
           median_length = tapply(results$length, results$method, median))
```

The randomized intervals use the information that selection left over, so they are
typically shorter than data splitting while keeping the nominal coverage.