# Regularised Regression: A Machine Learning Toolkit for Econometrics
# Athanassios Stavrakoudis
# astavrak@uoi.gr
# with claude's assistance
#
# Complete R code for the Ridge part
# Standalone: libraries loaded and configuration hard-coded below.

suppressPackageStartupMessages({
  library(glmnet)
  library(plm)
  library(sandwich)
  library(lmtest)
  library(wooldridge)
  library(xtable)
  library(tidyverse)
  library(patchwork)
})

# Reproducibility and parallel configuration
SEED       <- 14159L
N_CORES    <- 6L
N_THREADS  <- 12L
N_CV_FOLDS <- 10L
N_ALPHAS   <- 100L
N_MC_REPS  <- 200L
N_TREES    <- 1000L
set.seed(SEED)

# Lecture colour palette and ggplot theme
col_main   <- "#185FA5"
col_accent <- "#D85A30"
col_ok     <- "#1D9E75"
col_muted  <- "#6B7280"
col_warn   <- "#B22222"
theme_lecture <- theme_minimal(base_size = 13)


# [r-ridge-s1]
library(wooldridge)

wp <- wagepan

# - Outcome and a deliberately collinear, high-dimensional control block
y_col   <- "lwage"
cand    <- c("exper", "expersq", "educ", "married", "union", "black", "hisp",
             "south", "nrthcen", "nrtheast", "rur",
             "agric", "bus", "construc", "ent", "fin", "manuf", "min",
             "pro", "pub", "tra", "trad",
             "occ1", "occ2", "occ3", "occ4", "occ5", "occ6", "occ7", "occ8", "occ9",
             "poorhlth", "hours")
cand    <- intersect(cand, names(wp))
dat     <- wp[stats::complete.cases(wp[, c(y_col, cand)]), c(y_col, cand)]

X_r <- as.matrix(dat[, cand])
y_r <- dat[[y_col]]
n   <- nrow(X_r)
p   <- ncol(X_r)

# Standardise BEFORE fitting — Ridge is scale-sensitive
X_s       <- scale(X_r)
kappa_val <- kappa(t(X_s) %*% X_s)
cat(sprintf("n = %d  |  p = %d controls\n", n, p))
cat(sprintf("κ(X'X) = %.0f  [>30 problematic, >1000 severe]\n", kappa_val))

# 70/30 train/test split (firewall principle: test set never used in fitting)
set.seed(SEED)
tr_idx <- sample(n, floor(0.7 * n))
te_idx <- setdiff(seq_len(n), tr_idx)
X_tr <- X_s[tr_idx, ]
y_tr <- y_r[tr_idx]
X_te <- X_s[te_idx, ]
y_te <- y_r[te_idx]
cat(sprintf("Train: %d  |  Test: %d\n", length(tr_idx), length(te_idx)))

# [r-ridge-s2]
# OLS — fitted on training sample; likely unstable due to high κ
ols_tr     <- lm(y_tr ~ X_tr)

# Ridge — 10-fold CV on training sample selects λ
# alpha = 0 → pure L2 penalty (Ridge); standardize=FALSE because X_tr already scaled
cv_ridge_r <- cv.glmnet(
  x           = X_tr,
  y           = y_tr,
  alpha       = 0,
  nfolds      = N_CV_FOLDS,
  standardize = FALSE
)
lam_min    <- cv_ridge_r$lambda.min    # λ minimising CV-MSE
lam_1se    <- cv_ridge_r$lambda.1se   # largest λ within 1 SE of minimum (sparser)

cat(sprintf("OLS df        : %d\n", p))
cat(sprintf("Ridge λ.min   : %.4f\n", lam_min))
cat(sprintf("Ridge λ.1se   : %.4f\n", lam_1se))

# Effective degrees of freedom: df(λ) = Σ d²/(d²+λ)
d2     <- svd(X_s)$d ^ 2
df_min <- sum(d2 / (d2 + lam_min))
df_1se <- sum(d2 / (d2 + lam_1se))
cat(sprintf("df(λ.min) = %.2f  |  df(λ.1se) = %.2f  (OLS df = %d)\n",
            df_min, df_1se, p))

# [r-ridge-s3]
rmse_fn <- function(y, yhat) sqrt(mean((y - yhat)^2))
r2_fn   <- function(y, yhat) 1 - sum((y - yhat)^2) / sum((y - mean(y))^2)

pred_ols_tr  <- fitted(ols_tr)
pred_ols_te  <- predict(ols_tr, newdata = data.frame(X_tr = X_te))
pred_rdg_tr  <- as.numeric(predict(cv_ridge_r, X_tr, s = "lambda.min"))
pred_rdg_te  <- as.numeric(predict(cv_ridge_r, X_te, s = "lambda.min"))

gof_tbl <- tibble(
  Estimator    = c("OLS", "Ridge (λ.min)"),
  `Train RMSE` = c(rmse_fn(y_tr, pred_ols_tr), rmse_fn(y_tr, pred_rdg_tr)),
  `Test RMSE`  = c(rmse_fn(y_te, pred_ols_te), rmse_fn(y_te, pred_rdg_te)),
  `Test R²`    = c(r2_fn(y_te, pred_ols_te),   r2_fn(y_te, pred_rdg_te))
) %>%
  mutate(`Overfit ratio` = round(`Train RMSE` / `Test RMSE`, 3))

print(gof_tbl, n = Inf)

# Higher overfit ratio → more overfitting. OLS should be >> 1; Ridge closer to 1.

# [r-ridge-cvshape]
# Conceptual U-curve: illustrates the SHAPE of CV-MSE(λ); not a fit to data.
log_lam <- seq(-4, 6, length.out = 200)
# Flat-ish at small λ (OLS variance), rises at large λ (bias); gentle dip between.
cv_mse  <- 0.9 + 0.18 * (log_lam - 1)^2 / (1 + 0.15 * pmax(log_lam, 0)) -
           0.15 * exp(-(log_lam - 1)^2)
lam_min_x <- log_lam[which.min(cv_mse)]
ggplot(tibble(log_lam, cv_mse), aes(log_lam, cv_mse)) +
  geom_line(colour = col_main, linewidth = 1) +
  geom_vline(xintercept = lam_min_x, linetype = "dashed", colour = col_accent) +
  annotate("text", x = lam_min_x + 0.4, y = max(cv_mse),
           label = "λ.min", colour = col_accent, size = 3.5) +
  labs(x = expression(log(lambda)), y = "CV-MSE",
       title = "Cross-validation error is U-shaped in log λ") +
  theme_lecture + NULL

# [r-ridge-sfshape]
# Conceptual: shrinkage factor across ordered PC directions for two λ values.
d_sq  <- (seq(1, 0.04, length.out = 20))^2   # decreasing singular values²
sf <- bind_rows(
  tibble(PC = 1:20, sf = d_sq / (d_sq + 0.05), rule = "small λ"),
  tibble(PC = 1:20, sf = d_sq / (d_sq + 0.50), rule = "large λ")
)
ggplot(sf, aes(PC, sf, colour = rule)) +
  geom_hline(yintercept = 1, colour = "grey60", linetype = "dashed") +
  geom_point(size = 2.4) + geom_line(linewidth = 0.7) +
  scale_colour_manual(values = c("small λ" = col_main, "large λ" = col_accent)) +
  scale_y_continuous(limits = c(0, 1.05)) +
  labs(x = "PC direction (high variance → low variance)",
       y = expression(d[j]^2 / (d[j]^2 + lambda)),
       title = "Near-collinear directions are shrunk most", colour = NULL) +
  theme_lecture + NULL

# [r-ridge-traceshape]
# Conceptual trace: smooth shrinkage of several coefficients toward 0.
log_lam <- seq(-3, 5, length.out = 120)
starts  <- c(1.2, -0.9, 0.7, -0.5, 0.4, 0.25, -0.2)
trace_df <- bind_rows(lapply(seq_along(starts), function(k)
  tibble(log_lam, beta = starts[k] / (1 + exp(log_lam - 1)), j = factor(k))))
ggplot(trace_df, aes(log_lam, beta, colour = j)) +
  geom_hline(yintercept = 0, colour = "grey75") +
  geom_line(linewidth = 0.8) +
  geom_vline(xintercept = 1, linetype = "dashed", colour = "grey40") +
  annotate("text", x = 1.4, y = max(trace_df$beta) * 0.9,
           label = "CV λ", colour = "grey30", size = 3.2) +
  scale_colour_manual(
    values = colorRampPalette(c(col_main, col_accent, col_muted))(length(starts)),
    guide = "none") +
  labs(x = expression(log(lambda)), y = expression(hat(beta)[j](lambda)),
       title = "Ridge coefficients shrink smoothly — none hit exactly zero") +
  theme_lecture + NULL

# [r-ridge-bv-concept]
# Conceptual decomposition: illustrates the SHAPE of each component vs log λ.
# Not a Monte Carlo and not a fit to data — the analytic forms below are the
# textbook shapes (bias² rising, variance falling, MSE their U-shaped sum).
log_lam <- seq(-4, 5, length.out = 200)
variance <- 1.1 / (1 + exp(1.1 * (log_lam + 0.5)))   # high at small λ, →0
bias_sq  <- 0.9 / (1 + exp(-1.3 * (log_lam - 1.2)))   # ~0 at small λ, rising
mse      <- bias_sq + variance
lam_star <- log_lam[which.min(mse)]

bv <- bind_rows(
  tibble(log_lam, value = bias_sq, component = "Bias²"),
  tibble(log_lam, value = variance, component = "Variance"),
  tibble(log_lam, value = mse,      component = "MSE")
)
p_bv <- ggplot(bv, aes(log_lam, value, colour = component)) +
  geom_vline(xintercept = lam_star, linetype = "dashed",
             colour = "grey50", linewidth = 0.7) +
  geom_line(linewidth = 1.05) +
  scale_colour_manual(values = c("Bias²" = col_accent,
                                 "Variance" = col_main, "MSE" = col_ok)) +
  annotate("text", x = lam_star + 0.35, y = max(mse) * 0.96,
           label = "optimal λ", colour = "grey40", size = 3.5) +
  labs(x = expression(log(lambda)), y = "Error (conceptual units)",
       title = "Bias² rises, variance falls — MSE is their U-shaped sum",
       subtitle = "Ridge trades a little bias for a large variance reduction; optimal λ at the trough",
       colour = NULL) +
  theme_lecture
print(p_bv)

# [r-ridge-attenuation]
# Conceptual: the per-direction shrinkage factor d²/(d²+λ) as the singular value
# d falls. High-variance directions (large d) are barely shrunk; near-collinear
# directions (small d) are shrunk toward zero. Illustrative — not a fit to data.
d_vals <- seq(3, 0.1, length.out = 20)        # singular values, large → small
tibble(
  PC      = 1:20,
  small_l = d_vals^2 / (d_vals^2 + 0.3),
  large_l = d_vals^2 / (d_vals^2 + 3.0)
) %>%
  pivot_longer(c(small_l, large_l), names_to = "rule", values_to = "sf") %>%
  mutate(rule = recode(rule, small_l = "small λ", large_l = "large λ")) %>%
  ggplot(aes(PC, sf, colour = rule)) +
    geom_hline(yintercept = 1, colour = "grey60", linetype = "dashed") +
    geom_point(size = 2.4) + geom_line(linewidth = 0.8) +
    scale_colour_manual(values = c("small λ" = col_main, "large λ" = col_accent)) +
    scale_y_continuous(limits = c(0, 1.05),
      name = expression(d[j]^2 / (d[j]^2 + lambda))) +
    labs(x = "PC direction (high variance → low variance)",
         title = "Shrinkage factor by direction",
         subtitle = "Near-collinear directions (right) are damped most; larger λ shrinks all directions further",
         colour = NULL) + theme_lecture + NULL

# [r-ridge-wagepan]
library(wooldridge)
library(plm)
library(sandwich)

# - Load wagepan from the wooldridge package (hard fail if not installed)
wp2 <- wagepan

# - Candidate controls for the union-premium specification
ctrl <- c("exper", "expersq", "married", "educ", "black", "hisp", "south")
wp2  <- wp2[, c("nr", "year", "lwage", "union", ctrl)]

# - Within-transform (remove individual fixed effect)
wp2_dm <- wp2 %>%
  group_by(nr) %>%
  mutate(across(where(is.numeric),
                ~ . - mean(., na.rm = TRUE) + mean(wp2[[cur_column()]], na.rm = TRUE))) %>%
  ungroup()

# - Control matrix: demeaned covariates + year dummies
yr_dummies <- model.matrix(~ factor(year) - 1, data = wp2_dm)[, -1]
num_cols   <- setdiff(names(wp2_dm), c("nr", "year", "lwage", "union"))
ctrl_r2    <- cbind(as.matrix(wp2_dm[, num_cols]), yr_dummies)
y_r2       <- wp2_dm$lwage
D_r2       <- wp2_dm$union

# - TWFE baseline via plm (two-way within, cluster-robust SE)
pdat2    <- pdata.frame(wp2, index = c("nr", "year"))
twfe_m2  <- plm(lwage ~ union + exper + expersq + married + educ,
                data = pdat2, model = "within", effect = "twoways")
b_twfe2  <- coef(twfe_m2)["union"]
se_twfe2 <- sqrt(vcovHC(twfe_m2, cluster = "group")["union", "union"])

cat(sprintf("TWFE union premium: %.4f  (clustered SE: %.4f)\n", b_twfe2, se_twfe2))
cat(sprintf("Control matrix: %d obs × %d predictors\n", nrow(ctrl_r2), ncol(ctrl_r2)))

# Ridge with 10-fold CV — D included in the penalty (shrink all controls)
# For the union coefficient we use the post-Ridge strategy:
# partial out controls via Ridge, then OLS of residuals on D
set.seed(SEED)
cv_r_wp  <- cv.glmnet(
  x           = ctrl_r2,
  y           = y_r2,
  alpha       = 0,
  nfolds      = N_CV_FOLDS,
  standardize = TRUE
)
cv_d_wp  <- cv.glmnet(
  x           = ctrl_r2,
  y           = D_r2,
  alpha       = 0,
  nfolds      = N_CV_FOLDS,
  standardize = TRUE
)

# Partialled residuals
yres_r2 <- y_r2 - predict(cv_r_wp, ctrl_r2, s="lambda.min")
Dres_r2 <- D_r2 - predict(cv_d_wp, ctrl_r2, s="lambda.min")

# OLS of y-residuals on D-residuals (FWL / Robinson 1988)
ridge_ols2 <- lm(yres_r2 ~ Dres_r2)
b_ridge2   <- coef(ridge_ols2)["Dres_r2"]
se_ridge2  <- sqrt(vcovHC(ridge_ols2,"HC3")["Dres_r2","Dres_r2"])

cat(sprintf("Ridge (partial-out) union premium: %.4f  (HC3 SE: %.4f)\n",
            b_ridge2, se_ridge2))

# Effective df at lambda.min
d2_r2 <- svd(ctrl_r2, nu=0, nv=0)$d^2
df_r2 <- sum(d2_r2/(d2_r2 + cv_r_wp$lambda.min))
cat(sprintf("Ridge effective df: %.1f  (OLS df = %d)\n", df_r2, ncol(ctrl_r2)))

# - Which controls Ridge shrinks most — the numbers behind the shrinkage
# Ridge never zeroes; we show OLS vs Ridge coefficients on the SAME controls so
# the shrinkage is visible. Largest |Ridge coef| = controls Ridge leans on most.
b_ridge_ctrl <- as.numeric(coef(cv_r_wp, s = "lambda.min"))[-1]   # drop intercept
ols_full     <- lm(y_r2 ~ ctrl_r2)
b_ols_ctrl   <- coef(ols_full)[-1]
shrink_tbl <- tibble(
  Control      = colnames(ctrl_r2),
  `OLS b`      = round(as.numeric(b_ols_ctrl), 4),
  `Ridge b`    = round(b_ridge_ctrl, 4),
  `Shrunk by %`= round(100 * (1 - b_ridge_ctrl / as.numeric(b_ols_ctrl)), 1)
)
shrink_tbl <- arrange(shrink_tbl, desc(abs(`Ridge b`)))
cat("\nTop controls by |Ridge coefficient| (OLS vs Ridge):\n")
print(as.data.frame(head(shrink_tbl, 10)), row.names = FALSE)

tibble(
  Estimator = c("TWFE", "Ridge (partial-out)"),
  `Union premium` = c(b_twfe2, b_ridge2),
  `SE`            = c(se_twfe2, se_ridge2),
  `λ*`            = c("—", sprintf("%.4f", cv_r_wp$lambda.min)),
  `eff. df`       = c("5 (OLS)", sprintf("%.1f", df_r2))
) %>%
  kbl(caption="wagepan: Union Premium — TWFE vs Ridge (partial-out)", digits=4) %>%
  kable_styling(font_size=20, full_width=TRUE) %>%
  row_spec(2, bold=TRUE, color="white", background=col_main)