---
title: "Lecture 7: Bias-Variance Decomposition"
output:
  pdf_document: default
  html_notebook: default
---

# Setup

```{r setup, message=FALSE}
library(tidyverse)
library(FNN)

# Set seed for reproducibility
set.seed(226)
```

# Part 1: Flight Delays Example

## Load the data

```{r}
# Set seed for reproducibility
set.seed(226)

# Load raw flight data
flights_raw <- readRDS("../../data/flights_2024/flights_2024_top_15_raw.rds")

# Filter to valid flights (completed with delay data)
flights_valid <- flights_raw %>%
  filter(!is.na(arr_time) & !is.na(arr_delay))

head(flights_valid %>% select(carrier, origin, dest, arr_delay))
```

## Population parameters

```{r}
mu <- mean(flights_valid$arr_delay)
sigma <- sd(flights_valid$arr_delay)
n_population <- nrow(flights_valid)

cat(sprintf("Population mean (mu): %.2f minutes\n", mu))
cat(sprintf("Population standard deviation (sigma): %.2f minutes\n", sigma))
cat(sprintf("Population size: %d flights\n", n_population))
```

## Sampling distribution of the sample mean (n=10)

Generate the sampling distribution by repeatedly sampling n=10 flights and computing the sample mean.

We use an efficient algorithm: sample all data at once, then reshape into a matrix where each row is one simulation.

```{r}
n_sample <- 10
n_simulations <- 10000000

# Efficient sampling: one big sample, then reshape
all_samples <- sample(flights_valid$arr_delay,
                      size = n_simulations * n_sample,
                      replace = TRUE)
sample_matrix <- matrix(all_samples, nrow = n_simulations, ncol = n_sample)
sample_means <- rowMeans(sample_matrix)

# Compute properties of sampling distribution
sampling_mean <- mean(sample_means)
sampling_sd <- sd(sample_means)
sampling_bias <- sampling_mean - mu

cat(sprintf("Sampling distribution mean: %.3f\n", sampling_mean))
cat(sprintf("Sampling distribution SD (standard error): %.3f\n", sampling_sd))
cat(sprintf("Bias: %.3f\n", sampling_bias))
cat(sprintf("\nTheoretical standard error: sigma/sqrt(n) = %.3f\n", sigma/sqrt(n_sample)))
```

## Visualize sampling distribution

```{r fig.width=8, fig.height=5}
ggplot(data.frame(sample_mean = sample_means), aes(x = sample_mean)) +
  geom_density(fill = "steelblue", alpha = 0.6) +
  geom_vline(xintercept = mu, color = "red", linetype = "dashed", size = 1) +
  geom_vline(xintercept = sampling_mean, color = "darkblue", linetype = "solid", size = 1) +
  labs(
    title = "Sampling Distribution of Sample Mean (n=10)",
    subtitle = sprintf("Over %d simulations", n_simulations),
    x = "Sample Mean Arrival Delay (minutes)",
    y = "Density"
  ) +
  xlim(-60, 120) +   
  annotate("text", x = mu + 5, y = 0.03,
           label = sprintf("True mu = %.2f", mu),
           color = "red", hjust = 0) +
  annotate("text", x = sampling_mean - 5, y = 0.025,
           label = sprintf("E[Ybar] = %.2f", sampling_mean),
           color = "darkblue", hjust = 1) +
  theme_minimal(base_size = 14)
```


# Part 2: k-Nearest Neighbors Example

## Generate synthetic data

```{r}
# Set seed for reproducibility
set.seed(226)

n_total <- 1000
n_train <- 500
n_test <- 500

# Generate covariates
X1 <- rnorm(n_total, mean = 0, sd = 1)
X2 <- rnorm(n_total, mean = 0, sd = 1)

# Generate outcomes: Y = 1 + 2*X1 + 3*X2 + epsilon
epsilon <- rnorm(n_total, mean = 0, sd = sqrt(5))
Y <- 1 + 2*X1 + 3*X2 + epsilon

# True conditional expectation function
f_star <- function(x1, x2) {
  return(1 + 2*x1 + 3*x2)
}

# Create data frame
synthetic_data <- data.frame(X1 = X1, X2 = X2, Y = Y)

# Split into training and test
train_data <- synthetic_data[1:n_train, ]
test_data <- synthetic_data[(n_train+1):n_total, ]

cat(sprintf("Training set size: %d\n", nrow(train_data)))
cat(sprintf("Test set size: %d\n", nrow(test_data)))
```

## Visualize the synthetic data (scatter plot)

```{r fig.width=8, fig.height=6}
ggplot(train_data, aes(x = X1, y = X2, color = Y)) +
  geom_point(size = 2, alpha = 0.6) +
  scale_color_gradient2(low = "blue", mid = "white", high = "red",
                        midpoint = median(train_data$Y)) +
  labs(
    title = "Training Data: X1 vs X2",
    subtitle = "Color indicates Y value",
    x = "X1",
    y = "X2"
  ) +
  theme_minimal(base_size = 14) +
  coord_fixed()
```

## k-NN: RMSE as a function of k

```{r fig.width=8, fig.height=6}
k_values <- c(1, 2, 3, 5, 10, 15, 20, 30, 50, 75, 100, 150, 200, 300, 500)
rmse_values <- numeric(length(k_values))

for (i in seq_along(k_values)) {
  k <- k_values[i]

  # Fit k-NN model
  train_X <- train_data[, c("X1", "X2")]
  test_X <- test_data[, c("X1", "X2")]

  # Get predictions
  predictions <- knn.reg(train = train_X,
                         test = test_X,
                         y = train_data$Y,
                         k = k)$pred

  # Compute RMSE
  rmse_values[i] <- sqrt(mean((test_data$Y - predictions)^2))
}

knn_results <- data.frame(k = k_values, RMSE = rmse_values)

ggplot(knn_results, aes(x = k, y = RMSE)) +
  geom_line(color = "steelblue", size = 1) +
  geom_point(color = "steelblue", size = 3) +
  labs(
    title = "k-NN: RMSE vs k",
    subtitle = "Test set performance",
    x = "k (number of neighbors)",
    y = "RMSE"
  ) +
  theme_minimal(base_size = 14) +
  scale_x_log10()
```

## k-NN: Prediction errors across multiple samples

We'll pick a specific test point (X1=1, X2=1) and see how predictions vary across different training samples.

```{r fig.width=8, fig.height=6}
# Set seed for reproducibility
set.seed(226)

# Target point
target_X1 <- 1
target_X2 <- 1
true_f <- f_star(target_X1, target_X2)

# Range of k values to test
k_values_detail <- c(1, 2, 5, 10, 20, 50, 100, 200, 500)
n_repetitions <- 10

# Store results
prediction_errors <- data.frame()

for (rep in 1:n_repetitions) {
  # Generate new training data
  X1_train <- rnorm(n_train, mean = 0, sd = 1)
  X2_train <- rnorm(n_train, mean = 0, sd = 1)
  epsilon_train <- rnorm(n_train, mean = 0, sd = sqrt(5))
  Y_train <- 1 + 2*X1_train + 3*X2_train + epsilon_train

  train_matrix <- cbind(X1_train, X2_train)
  target_matrix <- matrix(c(target_X1, target_X2), nrow = 1)

  for (k in k_values_detail) {
    # Get prediction at target point
    pred <- knn.reg(train = train_matrix,
                    test = target_matrix,
                    y = Y_train,
                    k = k)$pred

    # Compute error: f*(X) - prediction
    error <- true_f - pred

    prediction_errors <- rbind(prediction_errors,
                               data.frame(k = k,
                                         repetition = rep,
                                         error = error,
                                         prediction = pred))
  }
}

# Plot prediction errors
ggplot(prediction_errors, aes(x = k, y = error, color = factor(repetition), group = repetition)) +
  geom_line(alpha = 0.7, size = 1) +
  geom_point(alpha = 0.7, size = 2) +
  geom_hline(yintercept = 0, linetype = "dashed", color = "black", size = 1) +
  labs(
    title = "k-NN: Prediction Error at (X1=1, X2=1)",
    subtitle = sprintf("True f*(X) = %.0f; Each color is one repetition", true_f),
    x = "k (number of neighbors)",
    y = "Error: f*(X) - Prediction",
    color = "Repetition"
  ) +
  theme_minimal(base_size = 14) +
  scale_x_log10() +
  theme(legend.position = "none")
```

As k increases, variance goes down (lines get closer together) while bias goes up (systematic shift away from 0).

## Summary statistics for bias and variance

```{r}
# Compute bias and variance for each k
bias_variance_summary <- prediction_errors %>%
  group_by(k) %>%
  summarize(
    mean_prediction = mean(prediction),
    bias = mean_prediction - true_f,
    variance = var(prediction),
    mse = mean((prediction - true_f)^2),
    .groups = "drop"
  )

print(bias_variance_summary)
```

# Part 3: Double Descent with Lasso

## Generate synthetic data for double descent

In this example, we explore the "double descent" phenomenon with lasso regression.

The data generating process is unusual:
1. First, we generate outcomes Y as N(0,1) random variables
2. Then, we generate features X that are correlated with Y plus noise: X_ij = Y_i + noise
3. Every additional feature brings new information about the outcome

We fit lasso models using increasing numbers of features (p = 100, 200, ..., 2000) and test them on a held-out test set.

```{r}
# Load the glmnet package for lasso
library(glmnet)

# Set seed for reproducibility
set.seed(226)

# Define the number of observations
n <- 1000
n_test <- 5000

# Define the range of p values (number of features)
p_min <- 100
p_max <- 2000
p_step <- 100

# Standard deviation of noise added to features
sd_noise <- 20

# Generate Y (outcomes) and X (features)
Y <- rnorm(n + n_test)
X_full <- matrix(rnorm((n + n_test) * p_max, 0, sd_noise), ncol = p_max) + Y

# Standardize Y and X
Y = scale(Y)
X_full = scale(X_full)

# Split into train and test sets
Ytrain = Y[1:n, ]
Ytest = Y[(n + 1):(n + n_test), ]
X_full_train = X_full[1:n, ]
X_full_test = X_full[(n + 1):(n + n_test), ]

cat(sprintf("Training set size: %d\n", n))
cat(sprintf("Test set size: %d\n", n_test))
cat(sprintf("Maximum number of features: %d\n", p_max))
cat(sprintf("Noise standard deviation: %.1f\n", sd_noise))
```

## Fit lasso models with varying number of features

```{r}
# Generate the values of p
p_values <- seq(from = p_min, to = p_max, by = p_step)

# Store the MSE results
train_MSE <- c()
test_MSE <- c()

# Loop over p (number of features)
for (p in p_values) {

  cat(sprintf("Fitting lasso with p = %d features...\n", p))

  # Select first p features
  Xtrain = X_full_train[, 1:p]
  Xtest = X_full_test[, 1:p]

  # Fit the lasso model (alpha = 1 means lasso, lambda = 0.001)
  fit <- glmnet(x = Xtrain, y = Ytrain, alpha = 1, lambda = 0.001)

  # Make predictions on training data
  predictions <- predict(fit, newx = Xtrain)
  # Calculate the training MSE
  MSE <- mean((predictions - Ytrain)^2)
  # Store the MSE
  train_MSE <- c(train_MSE, MSE)

  # Make predictions on test data
  predictions <- predict(fit, newx = Xtest)
  # Calculate the test MSE
  MSE <- mean((predictions - Ytest)^2)
  # Store the MSE
  test_MSE <- c(test_MSE, MSE)
}
```

## Visualize the results (sd_noise = 20)

```{r fig.width=8, fig.height=6}
# Prepare data for plotting
MSE = c(train_MSE, test_MSE)
Error = factor(c(rep("Train", length(p_values)),
          rep("Test", length(p_values))))
p_values_plot = c(p_values, p_values)

df = data.frame(p_values = p_values_plot, MSE = MSE, Error = Error)

# Plot the results
ggplot(df, aes(x = p_values, y = MSE, color = Error)) +
  geom_line(size = 1) +
  geom_point(size = 2) +
  labs(
    title = "Lasso Model: MSE vs Number of Features",
    subtitle = sprintf("lambda = 0.001, noise sd = %.0f", sd_noise),
    x = "Number of Features (p)",
    y = "Mean Squared Error (MSE)"
  ) +
  theme_minimal(base_size = 14)
```

The test error exhibits "double descent": it first decreases, then increases around p = 1000 (where p equals n, the training set size), and then decreases again as p continues to grow. This happens because:

- As p approaches n, lasso gets "confused": the noise is high and lambda is small, so it fits the training data well but generalizes poorly
- With more features beyond n, this "confusion" is overcome, and the additional features continue to bring useful information

## Compare with lower noise (sd_noise = 2)

Now let's repeat the experiment with lower noise (sd = 2 instead of 20). With lower noise, each feature brings clearer information about the outcome.

```{r}
# Set seed for reproducibility
set.seed(226)

# Standard deviation of noise added to features (lower noise)
sd_noise_low <- 2

# Generate Y (outcomes) and X (features)
Y_low <- rnorm(n + n_test)
X_full_low <- matrix(rnorm((n + n_test) * p_max, 0, sd_noise_low), ncol = p_max) + Y_low

# Standardize Y and X
Y_low = scale(Y_low)
X_full_low = scale(X_full_low)

# Split into train and test sets
Ytrain_low = Y_low[1:n, ]
Ytest_low = Y_low[(n + 1):(n + n_test), ]
X_full_train_low = X_full_low[1:n, ]
X_full_test_low = X_full_low[(n + 1):(n + n_test), ]

# Store the MSE results
train_MSE_low <- c()
test_MSE_low <- c()

# Loop over p (number of features)
for (p in p_values) {

  cat(sprintf("Fitting lasso with p = %d features (low noise)...\n", p))

  # Select first p features
  Xtrain_low = X_full_train_low[, 1:p]
  Xtest_low = X_full_test_low[, 1:p]

  # Fit the lasso model (alpha = 1 means lasso, lambda = 0.001)
  fit_low <- glmnet(x = Xtrain_low, y = Ytrain_low, alpha = 1, lambda = 0.001)

  # Make predictions on training data
  predictions_low <- predict(fit_low, newx = Xtrain_low)
  MSE_low <- mean((predictions_low - Ytrain_low)^2)
  train_MSE_low <- c(train_MSE_low, MSE_low)

  # Make predictions on test data
  predictions_low <- predict(fit_low, newx = Xtest_low)
  MSE_low <- mean((predictions_low - Ytest_low)^2)
  test_MSE_low <- c(test_MSE_low, MSE_low)
}
```

## Visualize the results (sd_noise = 2)

```{r fig.width=8, fig.height=6}
# Prepare data for plotting
MSE_low = c(train_MSE_low, test_MSE_low)
Error_low = factor(c(rep("Train", length(p_values)),
          rep("Test", length(p_values))))
p_values_plot_low = c(p_values, p_values)

df_low = data.frame(p_values = p_values_plot_low, MSE = MSE_low, Error = Error_low)

# Plot the results
ggplot(df_low, aes(x = p_values, y = MSE, color = Error)) +
  geom_line(size = 1) +
  geom_point(size = 2) +
  labs(
    title = "Lasso Model: MSE vs Number of Features",
    subtitle = sprintf("lambda = 0.001, noise sd = %.0f", sd_noise_low),
    x = "Number of Features (p)",
    y = "Mean Squared Error (MSE)"
  ) +
  theme_minimal(base_size = 14)
```

With lower noise, the test error continues to drop as p increases - there is no "double descent" phenomenon. Why? Because every additional feature brings clear, novel information about the outcome. The features are strongly correlated with Y (low noise), so lasso does not get "confused" even when p approaches n.

# Summary

This notebook generates all the examples and visualizations for Lecture 7:

1. **NBA Rookies**: Population parameters, sampling distribution of sample mean
2. **k-NN Example**: Synthetic data, RMSE vs k, prediction errors across samples
3. **Double Descent with Lasso**: Demonstrating how test error can decrease, increase, then decrease again as model complexity grows

The key insights:
- Sample mean is unbiased with variance $\sigma^2/n$
- As k increases in k-NN, variance decreases (predictions more stable) but bias increases (predictions systematically off)
- In the right data context (where new features bring novel information), both bias and variance can decrease with increasing model complexity, leading to the "double descent" phenomenon