Predictive performance and increasing model complexity for nonlinear modelling approaches

library(brms)
library(loo)
library(tidyverse)
library(geomtextpath)
library(khroma)
library(MASS)
library(splines)
library(ggdist)

options(brms.backend = "cmdstanr")
seed_train <- 42424242 
seed_test <- seed_train / 2

# set theme and labels for plots ####
fontsize = 8 
theme_set(theme_bw() +
            theme(legend.text = element_text(size = fontsize),
                  legend.key.size = unit(0.7, "cm"),
                  legend.title = element_blank(), 
                  legend.position = "bottom",
                  legend.direction = "horizontal",
                  legend.box.spacing = unit(0, "pt"),
                  legend.margin = margin(0, 0, 0, 0),
                  legend.spacing.x = unit(0, "mm"),
                  legend.spacing.y = unit(0, "mm"),
                  panel.grid.major = element_blank(),
                  panel.grid.minor = element_blank(),
                  strip.background = element_blank(),
                  panel.background = element_blank(),
                  text = element_text(size = fontsize),
                  plot.title = element_text(size = fontsize),
                  axis.title = element_text(size = fontsize),
                  axis.text = element_text(size = fontsize),
                  strip.text = element_text(size = fontsize)))

custom_colours <- c("Polynomial (raw)" = "#2497DD", 
                    "Polynomial (raw, init=0.1)" = "#0072B2",
                    "Polynomial (orthogonal)" = "#66CCEE",
                    "Thin plate splines (fixed df)" = "#EE6677",
                    "Thin plate splines" = "#CC79A7",
                    "HSGP" = "#AA3377") 

degree_labels <- as_labeller(
  c(`3` = "k = 3",
    `5` = "k = 5",
    `7` = "k = 7",
    `9` = "k = 9",
    `11` = "k = 11",
    `13` = "k = 13", 
    `15` = "k = 15", 
    `17` = "k = 17",
    `19` = "k = 19",
    `20` = "k = 20"))

We adopt a similar setup as, for example, (Rasmussen and Ghahramani 2000) and Chapter 7 by (McElreath 2020), and model a nonlinear relationship between a single predictor and a response with an increasing polynomial degree or number of basis functions. We consider the following conditions and modelling approaches:

# number of obs 
n <- 100
n_test <- 2000 

# fixed range for x-values 
xmin <- -3 
xmax <- 3

selected_degrees <- c(3, 5, 7, 9, 11, 13, 15, 17, 19, 20)

models <- c("Polynomial (raw)", 
            "Polynomial (raw, init=0.1)", 
            "Polynomial (orthogonal)",
            "Thin plate splines (fixed df)",
            "Thin plate splines",
            "HSGP") 

conditions <- expand.grid(degree = selected_degrees, model = models) 

# helper function: create synthetic data similar to drowning.txt ####
get_synthetic_drowning_data <- function(original_data, 
                                        n, 
                                        noise_sd, 
                                        xmin, 
                                        xmax, 
                                        scale = FALSE,
                                        df_glm_bs = 30, 
                                        k = 60){
  
  nb_mod <- MASS::glm.nb(y ~ bs(x, df = df_glm_bs), data = original_data)
  theta <- nb_mod$theta
  
  keep_ids <- sample(1:(k*NROW(original_data)), n) # reduce to desired number of obs
  
  mu_hat <- fitted(nb_mod) # expected counts per year
  mu_rep <- rep(mu_hat, each = k) # k synthetic replicates per observed x
  mu_rep_reduced <- mu_rep[keep_ids]
  
  x <- rep(original_data$x, each = k)
  x_reduced <- x[keep_ids]
  
  synthetic_counts <- rnbinom(length(mu_rep_reduced), size = theta, mu = mu_rep_reduced)
  
  synthetic_df <- data.frame(
    x_unscaled = x_reduced + rnorm(n, sd = noise_sd), 
    y = synthetic_counts) |>
    # scale to ensure x in [xmin, xmax] as for the other dgps 
    mutate(x = ((2*xmax) * (x_unscaled - min(x_unscaled)) / (max(x_unscaled) - min(x_unscaled)) + xmin)) 
    
    d <- data.frame(y = synthetic_df$y, x = synthetic_df$x)
    d <- d[order(d$x), ]
    if(scale) {d$y = scale(d$y)}
    return(d)
}

# helper function: get model specification based on modelling approach+degree/number of basis functions 
get_model_spec <- function(degree, model) {
  stopifnot(degree >= 3)
  
  prior_poly <- c(
    prior(normal(0, 1), class = "Intercept"), 
    prior(normal(0, 1), class = "b"),
    prior(exponential(1), class = "sigma")
  )
  
  prior_gp <- c(
    prior(normal(0, 1), class = "Intercept"),
    prior(inv_gamma(2, 1), class = "lscale", coef = "gpx"),
    prior(inv_gamma(2, 1), class = "sdgp"),
    prior(inv_gamma(2, 1), "sigma")
  )
  
  init_value <- NULL
  
  if (model == "Polynomial (raw)") {
    model_formula  <- paste0("y ~ poly(x,", degree, ", raw = TRUE)")
    prior_setting <- prior_poly
    
  } else if (model == "Polynomial (raw, init=0.1)") {
    model_formula  <- paste0("y ~ poly(x,", degree, ", raw = TRUE)")
    prior_setting <- prior_poly
    init_value <- 0.1
    
  } else if (model == "Polynomial (orthogonal)") {
    model_formula  <- paste0("y ~ poly(x,", degree, ")")
    prior_setting <- prior_poly
    
  } else if (model == "Thin plate splines (fixed df)") {
    model_formula  <- paste0("y ~ x + s(x, fx = TRUE, k = ", degree, ")")
    prior_setting <- prior_poly
    
  } else if (model == "Thin plate splines") {
    model_formula  <- paste0("y ~ x + s(x, k = ", degree, ")")
    prior_setting <- prior_poly
    
  } else if (model == "HSGP") {
    model_formula  <- paste0("y ~ gp(x, k = ", degree, ", c = 3/2)")
    prior_setting <- prior_gp
  } else {
    stop("Unknown model: ", model)
  }
  
  list(model_formula = model_formula, 
       prior_setting = prior_setting, 
       init_value = init_value)
}

# helper function: get model fit objects given degrees and models
get_model_fit <- function(degree, model, dataset = data){
  stopifnot(degree >= 3)
  
  # get model specification 
  model_specs <- get_model_spec(degree, model)
  
  # fit model 
  model_fit_object <- brm(formula = model_specs$model_formula, 
                          data = dataset,
                          prior = model_specs$prior_setting,
                          init = model_specs$init_value,
                          backend = "cmdstanr", 
                          chains = 4, 
                          cores = 4, 
                          refresh = 0,
                          silent = 2)
  
  return(model_fit_object)
}

get_model_fit_prior_only <- function(degree, model, dataset = data){
  stopifnot(degree >= 3)
  
  # get model specification 
  model_specs <- get_model_spec(degree, model)
  
  model_fit_object_prior_only <- brm(formula = model_specs$model_formula, 
                                     data = dataset,
                                     prior = model_specs$prior_setting,
                                     init = model_specs$init_value,
                                     sample_prior = "only",
                                     backend = "cmdstanr", 
                                     chains = 4, 
                                     cores = 4, 
                                     refresh = 0,
                                     silent = 2)
  
  return(model_fit_object_prior_only)
}

# helper function: extract posterior predictive samples
get_posterior_predictions <- function(model_fit) {
  posterior_predictive_draws <- posterior_predict(model_fit)
  pred_df <- tibble(
    posterior_mean = apply(posterior_predictive_draws, 2, mean),
    posterior_sd = apply(posterior_predictive_draws, 2, sd),
    posterior_lower_50 = apply(posterior_predictive_draws, 2, quantile, probs = 0.25),
    posterior_upper_50 = apply(posterior_predictive_draws, 2, quantile, probs = 0.75),
    posterior_lower_90 = apply(posterior_predictive_draws, 2, quantile, probs = 0.05),
    posterior_upper_90 = apply(posterior_predictive_draws, 2, quantile, probs = 0.95)
  )
  return(pred_df)
}

# helper function: get elpd test 
get_elpd_test <- function(model_object, test_data){
  elpd_test <- elpd(log_lik(model_object, newdata = test_data))
  return(elpd_test)
}

# helper function: get basis function values for different modelling approaches 
get_basis_values_long <- function(
    degree, 
    model,
    x_grid = seq(min(data$x), max(data$x), length.out = 500), 
    n_basis = 6) {
  
  # get model specifications 
  spec <- get_model_spec(degree, model)
  
  # make_standata needs a data.frame with all variables in the formula
  newdat <- tibble(
    x = x_grid,
    y = 0 # dummy, not used for basis extraction
  )
  
  sd <- make_standata(
    formula = as.formula(spec$model_formula),
    data = newdat,
    prior = spec$prior_setting
  )
  
  X <- sd$X
  col_names <- colnames(X)

  if (model == "Thin plate splines (fixed df)"){
    
    Xs <- sd$Xs
    col_names_xs <- colnames(Xs)
    X <- cbind(X, Xs)
    col_names <- c(col_names, col_names_xs)
  
  } else if (model == "Thin plate splines"){
    
    Xs <- sd$Xs
    col_names_xs <- colnames(Xs)
    
    Zs <- sd$Zs_1_1
    penalised_tps_names <- rep(paste0("smooth_term_x", 2:(NCOL(Zs)+1)))
    colnames(Zs) <- penalised_tps_names
    col_names_zs <- colnames(Zs)
    
    X <- cbind(X, Xs, Zs)
    col_names <- c(col_names, col_names_xs, col_names_zs)
    
  } else if (model == "HSGP"){
    
    X_gp <- sd$Xgp_1
    gp_names <- rep(paste0("gp_term_x", 1:(NCOL(X_gp))))
    colnames(X_gp) <- gp_names
    col_names_x_gp <- colnames(X_gp)
    X <- X_gp
    col_names <- col_names_x_gp
    
  }
  
  if (model %in% c("Polynomial (raw)", "Polynomial (raw, init=0.1)", "Polynomial (orthogonal)")){
    basis_cols <- seq_len(NCOL(X))[-1]  # drop intercept column
    basis_cols <- head(basis_cols, n_basis) # only keep the first n_basis columns
    basis_mat <- X[, basis_cols, drop = FALSE]
  } else if (model %in% c("Thin plate splines (fixed df)", "Thin plate splines")){
    basis_cols <- seq_len(NCOL(X))[-1]  # drop intercept column
    basis_cols <- head(basis_cols, n_basis + 1)
    basis_mat <- X[, basis_cols, drop = FALSE]
  } else if (model == "HSGP"){
    basis_cols <- seq_len(NCOL(X))[-1]  # drop intercept column
    basis_cols <- head(basis_cols, n_basis)
    basis_mat <- X[, basis_cols]
  } 
  
  as_tibble(basis_mat) |>
    mutate(x = x_grid) |>
    pivot_longer(cols = -x,
                 names_to = "basis_name",
                 values_to = "value") |>
    mutate(basis_id = match(basis_name, unique(basis_name)))
}

First, we generate data \(\{x_i, y_i\}_{i=1, \cdots, n}\) with \(n=100\) as well as independent test data with \(n_{\text{test}}=2000\). We generate the data such that it is similar to a real data set that contains number of drownings observed in 34 different years from 1980 to 2013. This data has been used, for example, in BDA 3 and is available in the accompanying GitHub repository for BDA R demos. Details of the data-generating process can be found in the helper function get_synthetic_drowning_data() in the above folded code chunk.

# generate data similar to drowning data from BDA ####
original_data <- read.table(
  here::here("data", "model-complexity-basis-functions", "drowning.txt"), 
  header = FALSE, 
  stringsAsFactors = FALSE) |>
  rename(x = V1, y = V2) 

data <- withr::with_seed(
  seed_train, 
  get_synthetic_drowning_data(
    original_data = original_data, 
    n = n,
    xmin = xmin, 
    xmax = xmax, 
    scale = TRUE, 
    noise_sd = 0.5)) 

data_test <- withr::with_seed(
  seed_test, 
  get_synthetic_drowning_data(
    original_data = original_data, 
    n = n_test,
    xmin = xmin, 
    xmax = xmax,
    scale = TRUE, 
    noise_sd = 0.5))

We can visualise outcome vs. covariate of our generated dataset with a scatterplot.

# plot one realisation of the true DGP ####
p_scatter <- ggplot(
  data = data, aes(x = x, y = y)) +
  geom_point() +
  ggtitle(paste0("DGP: Like BDA drowning data"))
p_scatter

Now, we fit (1) a raw polynomial, (2) a raw polynomial with adjusted initial value, (3) an orthogonal polynomial, a thin plate spline (4) with and (5) without fixed degrees of freedom, and (6) a Hilbert-space approximated Gaussian process. We repeat this for degrees/number of basis functions in \(\{3, 5, 7, 10, 12, 15, 17, 20\}\).

# add model fit objects given conditions
df <- conditions |>
  mutate(model_fit_object = purrr::map2(degree, model, get_model_fit), 
         model_fit_object_prior_only = map2(degree, model, get_model_fit_prior_only)) 

We plot the observed data versus the posterior predictions we obtain for each modelling approach and degree/number of basis functions.

# add summaries based on draws from posterior predictive ####
df_with_posterior_results <- df |> 
  mutate(posterior_pred_stats = map(model_fit_object, get_posterior_predictions)) |> 
  unnest(posterior_pred_stats) 

# add observations for plotting ####
df_with_obs <- df_with_posterior_results |>
  mutate(x = rep(data$x, length(selected_degrees) * length(models)),
         y = rep(data$y, length(selected_degrees) * length(models)))

plot_data_vs_posterior_predictions_n_100 <- 
  ggplot(data = df_with_obs |> filter(degree %in% c(3, 5, 7, 11, 13, 17, 20)), # only select some degrees for plot 
         aes(x = x, y = y, color = model, fill = model)) +
  geom_point(color = "black", alpha = 0.5) + # observed data points
  geom_line(aes(y = posterior_mean)) + # posterior mean
  geom_ribbon(aes(ymin = posterior_lower_50, ymax = posterior_upper_50), alpha = 0.4, colour = NA) +
  geom_ribbon(aes(ymin = posterior_lower_90, ymax = posterior_upper_90), alpha = 0.2, colour = NA) +   
  scale_colour_manual(values = custom_colours, labels = models) +
  scale_fill_manual(values = custom_colours, labels = models) +
  coord_cartesian(ylim = c(-3, 3)) + # with scaled y 
  facet_grid(degree ~ model, labeller = labeller(degree = degree_labels)) +
  labs(x = "x", y = "y", subtitle = paste0("n = ", n)) +
  theme(strip.text.x = element_blank())

plot_data_vs_posterior_predictions_n_100

Now, we use the test data to evaluate the predictive performance on previously unobserved data.

# check predictive performance on test data ####
df_with_elpd_test <- df |>
  group_by(degree, model) |> 
  mutate(elpd_test_object = map(model_fit_object, ~get_elpd_test(.x, data_test))) |>
  mutate(elpd_test = map_dbl(elpd_test_object, function(x) x$estimates["elpd", "Estimate"]))  

# plot elpd test across degrees for the different modelling approaches 
ggplot(data = df_with_elpd_test, aes(x = degree, y = elpd_test, group = model, colour = model)) + 
  geom_point() + 
  geom_line() +
  scale_colour_manual(values = custom_colours, labels = models) + 
  coord_cartesian(ylim = c(-2300, -1800))

We investigate differences in prior predictive variance between the modelling approaches for an increasing degree/number of basis functions.

# visualise prior predictive variance ####
df_with_prior_predictions <- df |>
  mutate(yrep_prior = map(model_fit_object_prior_only, posterior_predict), 
         prior_pred_vars = map(yrep_prior, ~apply(.x, 2, var))) 

ggplot(df_with_prior_predictions |> 
         group_by(model, degree) |> 
         mutate(degree = as.factor(degree)) |>
         unnest(c(prior_pred_vars)), 
       aes(x = degree, y = prior_pred_vars, group = degree, colour = model, fill = model)) +
  geom_boxplot(alpha = 0.3) + 
  facet_grid(cols = vars(model)) + 
  scale_y_continuous(trans = "log10") +
  scale_colour_manual(values = custom_colours, labels = models) +
  scale_fill_manual(values = custom_colours, labels = models) +
  labs(
    y = "Prior predictive variance per observation (log10-transformed)",
    title = paste0("One repetition, DGP: Like BDA drowning data, n = ", n)) +
  theme(strip.text.x = element_blank())

ggplot(df_with_prior_predictions |> 
         group_by(model, degree) |> 
         mutate(degree = as.factor(degree)) |>
         filter(model != "Polynomial (raw, init=0.1)") |>
         unnest(c(prior_pred_vars)), 
       aes(x = degree, y = prior_pred_vars, group = degree, colour = model, fill = model)) +
  geom_boxplot(alpha = 0.3) + 
  facet_grid(cols = vars(model)) + 
  scale_y_continuous(trans = "log10") +
  scale_colour_manual(
    values = custom_colours[names(custom_colours) != "Polynomial (raw, init=0.1)"], 
    labels = models[!models == "Polynomial (raw, init=0.1)"]) +
  scale_fill_manual(
    values = custom_colours[names(custom_colours) != "Polynomial (raw, init=0.1)"],
    labels = models[!models == "Polynomial (raw, init=0.1)"]) +
  coord_cartesian(ylim = c(2.5, 200)) + # NOTE: limit view to focus on alternatives to raw polynomial
  labs(
    y = "Prior predictive variance per observation (log10-transformed)",
    title = paste0("One repetition, DGP: Like BDA drowning data, n = ", n)) +
  theme(strip.text.x = element_blank())

Since the prior predictive variance for the raw polynomials explodes, we can also “zoom in” to focus on the alternative modelling approaches.

ggplot(df_with_prior_predictions |> 
         group_by(model, degree) |> 
         mutate(degree = as.factor(degree)) |>
         filter(model != "Polynomial (raw, init=0.1)") |>
         unnest(c(prior_pred_vars)), 
       aes(x = degree, y = prior_pred_vars, group = degree, colour = model, fill = model)) +
  stat_dots(quantiles = 100, 
            layout = "swarm", 
            alpha = 0.8) + 
  facet_grid(cols = vars(model)) + 
  scale_y_continuous(trans = "log10") +
  coord_cartesian(ylim = c(2.5, 200)) + # NOTE: this is currently a limited view to focus on alternatives to raw polynomial
   scale_colour_manual(
    values = custom_colours[names(custom_colours) != "Polynomial (raw, init=0.1)"], 
    labels = models[!models == "Polynomial (raw, init=0.1)"]) +
  scale_fill_manual(
    values = custom_colours[names(custom_colours) != "Polynomial (raw, init=0.1)"],
    labels = models[!models == "Polynomial (raw, init=0.1)"]) +
  labs(
    y = "Prior predictive variance per observation (log10-transformed)") +
  theme(strip.text.x = element_blank())

It might also be interesting to focus, for example, only on the thin-plate splines and the HSGP.

ggplot(df_with_prior_predictions |> 
         group_by(model, degree) |> 
         mutate(degree = as.factor(degree)) |>
         filter(model %in% c("Thin plate splines (fixed df)",
                             "Thin plate splines",
                             "HSGP")) |>
         unnest(c(prior_pred_vars)), 
       aes(x = degree, y = prior_pred_vars, group = degree)) +
  geom_boxplot(alpha = 0.3) + 
  facet_grid(cols = vars(model)) + 
  scale_y_continuous(trans = "log10") +
  labs(
    y = "Prior predictive variance per observation (log10-transformed)",
    title = paste0("One repetition, DGP: Like BDA drowning data, n = ", n)) 

ggplot(df_with_prior_predictions |> 
         group_by(model, degree) |> 
         mutate(degree = as.factor(degree)) |> 
         filter(model %in% c("Thin plate splines (fixed df)",
                             "Thin plate splines",
                             "HSGP")) |>
         unnest(c(prior_pred_vars)), 
       aes(x = degree, y = prior_pred_vars, group = degree)) +
  stat_dots(quantiles = 100, 
            layout = "swarm", 
            alpha = 0.8) + 
  facet_grid(cols = vars(model)) + 
  scale_y_continuous(trans = "log10") +
  labs(
    y = "Prior predictive variance per observation (log10-transformed)")

Lastly, we can also investigate the different basis function values for each of the modelling approaches for a chosen degree/number of basis functions. Below, we visualise the first five basis values for (1) raw polynomials, (2) orthogonal polynomials, (3) thin plate splines with fixed penalisation, (4) penalised thin plate spline, and (5) HSGP’s.

n_basis_to_plot <- 5
models <- c(
  "Polynomial (raw)",
  "Polynomial (orthogonal)",
  "Thin plate splines (fixed df)",
  "Thin plate splines", 
  "HSGP"
  )
selected_degrees <- c(7, 20) 

basis_df <- expand.grid(degree = selected_degrees, model = models) |>
  as_tibble() |>
  mutate(
    basis_data = map2(degree, model, ~get_basis_values_long(.x, .y, n_basis = n_basis_to_plot))) |>
  unnest(basis_data)

plot_basis_functions <- ggplot(
  basis_df,
  aes(x = x, y = value, group = basis_id, colour = factor(basis_id))) +
  geom_line() +
  scale_colour_light() +
  facet_wrap(
    nrow = 2, 
    ncol = 5, 
    degree ~ model,
    labeller = labeller(degree = degree_labels), 
    scales = "free_y"
  ) +
  labs(
    x = "x",
    y = "Basis value",
    title = paste0("First ", n_basis_to_plot, " basis functions"),
    subtitle = paste0("DGP: Like BDA drowning data")
  ) +
  theme(legend.position = "bottom")

plot_basis_functions

We see that the range of the basis function values for the first five basis functions differs largely between the different modelling approaches. Again, the raw polynomial does not handle increasing complexity well since, already for the fifth basis function, the range of basis function values for the raw polynomial increases to around -200 to 200. For all other modelling approaches, the range of basis function values is much lower.

References

McElreath, Richard. 2020. Statistical Rethinking: A Bayesian Course with Examples in R and Stan. Chapman; Hall/CRC.

Rasmussen, Carl, and Zoubin Ghahramani. 2000. “Occam’ s Razor.” In Advances in Neural Information Processing Systems. Vol. 13. MIT Press. https://papers.nips.cc/paper/2000/hash/0950ca92a4dcf426067cfd2246bb5ff3-Abstract.html.