Random Slopes (LMM)

### Load necessary libraries
library(NBDCtools)  # ABCD data access helper
library(arrow)       # For reading Parquet files
library(tidyverse)   # For data manipulation & visualization
library(gtsummary)   # For generating publication-quality summary tables
library(rstatix)     # Provides tidy-format statistical tests
library(lme4)        # Linear mixed-effects models (LMMs)
library(kableExtra)  # Formatting & styling in HTML/Markdown reports
library(performance) # Useful functions for model diagnostics & comparisons
library(lmerTest)    # Provides p-values for lmer models
library(sjPlot)      # Visualization of mixed models

### Load harmonized ABCD data required for this analysis
requested_vars <- c(
  "ab_g_dyn__design_site",
  "ab_g_stc__design_id__fam",
  "nc_y_nihtb__comp__cryst__fullcorr_tscore"
)

data_dir <- Sys.getenv("ABCD_DATA_PATH", "/path/to/abcd/6_0/phenotype")

abcd_data <- create_dataset(
  dir_data = data_dir,
  study = "abcd",
  vars = requested_vars,
  release = "6.0",
  format = "parquet",
  categ_to_factor = TRUE,   # Convert categorical variables to factors
  value_to_na = TRUE,        # Convert missing codes (222, 333, etc.) to NA
  add_labels = TRUE          # Add variable and value labels
)

### Create a long-form dataset with relevant columns
df_long <- abcd_data %>%
  # Filter to available annual assessments using NBDCtools
  filter_events_abcd(conditions = c("annual")) %>%
  # Rename variables for clarity
  rename(
    site = ab_g_dyn__design_site,              # site already a factor from NBDCtools
    family_id = ab_g_stc__design_id__fam,
    cognition = nc_y_nihtb__comp__cryst__fullcorr_tscore  # cognition already numeric from NBDCtools
  ) %>%
  # Keep only participants with at least 2 non-missing cognition scores
  group_by(participant_id) %>%
  filter(sum(!is.na(cognition)) >= 2) %>%
  ungroup() %>%
  drop_na(cognition) %>%  # Remove rows with missing outcome only (lmer handles missing predictors)
  # Create numeric time variable from session_id
  mutate(time = as.numeric(session_id) - 1)  # Convert session to numeric time (0, 1, 2, 3...)

### Create descriptive summary table
descriptives_table <- df_long %>%
  select(session_id, cognition) %>%
  tbl_summary(
    by = session_id,
    missing = "no",
    label = list(
      cognition ~ "Cognition"
    ),
    statistic = list(all_continuous() ~ "{mean} ({sd})")
  ) %>%
  modify_header(all_stat_cols() ~ "**{level}**<br>N = {n}") %>%
  modify_spanning_header(all_stat_cols() ~ "**Assessment Wave**") %>%
  bold_labels() %>%
  italicize_levels()

### Apply compact styling
theme_gtsummary_compact()

descriptives_table <- as_gt(descriptives_table)

### Save the table as HTML
gt::gtsave(descriptives_table, filename = "descriptives_table.html")

### Print the table
descriptives_table

### Fit a Linear Mixed Model (LMM) with random intercepts and slopes
model <- lmerTest::lmer(
    cognition ~ time + (1 + time | participant_id), # Fixed effect (time), random intercept & slope (participant_id)
    data = df_long # Dataset containing repeated measures of cognition
)

### Generate a summary table for the LMM model (gtsummary format)
model_summary_table <- gtsummary::tbl_regression(model,
    digits = 3,
    intercept = TRUE
) %>%
  gtsummary::as_gt()

### Save the gt table
gt::gtsave(
  data = model_summary_table,
  filename = "model_summary.html",
  inline_css = FALSE
)

# Generate an alternative summary table with random effects (sjPlot format)
# This provides additional details on variance components not shown in gtsummary
sjPlot::tab_model(model,
    show.se = TRUE, show.df = FALSE, show.ci = FALSE,
    digits = 3,
    pred.labels = c("Intercept", "Time"),  # Adjust predictor labels
    dv.labels = c("Random Intercept & Slope LMM"),  # Update model label
    string.se = "SE",
    string.p = "P-Value",
    file = "lmm_model_results.html"
)

# Generate quick diagnostic plot using sjPlot
sjPlot::plot_model(model,
  type = "pred",
  terms = c("time"),
  title = "Predicted Cognition Scores Over Time",
  axis.title = c("Time", "Predicted Cognition"),
  show.data = TRUE,
  ci.lvl = 0.95,
  dot.size = 2,
  line.size = 1.2,
  jitter = 0.2,
  file = "lmm_sjPlot_results.html"
)

# Generate model predictions for custom visualization
df_long <- df_long %>%
  mutate(
    predicted_fixed = predict(model, re.form = NA),
    predicted_random = predict(model, re.form = ~ (1 + time | participant_id))
  )

# Select a subset of participant IDs for visualization
set.seed(124)
sample_ids <- sample(unique(df_long$participant_id),
                     size = min(250, length(unique(df_long$participant_id))))

# Filter dataset to include only sampled participants
df_long_sub <- df_long %>%
  filter(participant_id %in% sample_ids)

# Create the visualization of individual and overall cognition trajectories
visualization <- ggplot(df_long_sub, aes(x = time, y = cognition, group = participant_id)) +

  # Plot observed data (individual trajectories)
  geom_line(aes(color = "Observed Data"), alpha = 0.3) +
  geom_point(alpha = 0.3) +

  # Plot model-predicted values with random intercepts and slopes
  geom_line(aes(y = predicted_random, color = "Model Predictions"), alpha = 0.7) +
  geom_line(aes(y = predicted_fixed, group = 1, color = "Population Mean"), linewidth = 1.2) +

  # Customize colors for clarity
  scale_color_manual(values = c(
    "Observed Data" = "red",
    "Model Predictions" = "grey40",
    "Population Mean" = "blue"
  )) +

  # Format x-axis labels
  scale_x_continuous(
    breaks = 0:3,
    labels = c("Baseline", "Year 2", "Year 4", "Year 6")
  ) +

  # Add labels and title
  labs(
    title = "Random Intercept-Slope Model: Individual Trajectories",
    x = "Assessment Wave",
    y = "Cognition",
    color = "Trajectory Type"
  ) +

  # Apply theme
  theme_minimal() +
  theme(legend.position = "bottom")

# Display the plot
visualization

# Save the plot
ggsave(
  filename = "visualization.png",
  plot = visualization,
  width = 8, height = 6, dpi = 300
)