Time-Varying Covariates (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(lme4)         # Linear mixed-effects models (LMMs)
library(lmerTest)     # P-values for lmer models
library(broom.mixed)  # Tidy output for mixed models
library(gt)           # For creating formatted tables

### 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__fluid__fullcorr_tscore",  # Fluid cognition (outcome)
  "ph_y_anthr__height_mean"                     # Height in inches (TVC)
)

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,
  value_to_na = TRUE,
  add_labels = TRUE
)

# Prepare data for LMM analysis with time-varying covariate
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,
    family_id = ab_g_stc__design_id__fam,
    cognition = nc_y_nihtb__comp__fluid__fullcorr_tscore,
    height = ph_y_anthr__height_mean
  ) %>%
  # Keep only participants with at least 2 non-missing observations
  group_by(participant_id) %>%
  filter(sum(!is.na(cognition) & !is.na(height)) >= 2) %>%
  ungroup() %>%
  drop_na(cognition, height) %>%
  # Create numeric time variable from session_id
  mutate(time = as.numeric(session_id) - 1) %>%
  # CRITICAL: Person-mean centering of time-varying covariate
  group_by(participant_id) %>%
  mutate(
    # Between-person component: person's average height across all waves
    height_pm = mean(height, na.rm = TRUE),
    # Within-person component: deviation from person's own average
    height_cwc = height - height_pm
  ) %>%
  ungroup()

# Verify centering worked correctly
df_long %>%
  group_by(participant_id) %>%
  summarise(
    mean_cwc = mean(height_cwc, na.rm = TRUE),
    .groups = "drop"
  ) %>%
  summary()  # Mean of cwc should be ~0 for each person

# Create descriptive summary table by wave
descriptives_table <- df_long %>%
  select(session_id, cognition, height) %>%
  tbl_summary(
    by = session_id,
    missing = "no",
    label = list(
      cognition ~ "Fluid Cognition",
      height ~ "Height (inches)"
    ),
    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 baseline model with just time
m1 <- lmer(
  cognition ~ time + (1 + time | participant_id),
  data = df_long,
  REML = TRUE
)

summary(m1)

# Add raw (uncentered) height as predictor
# WARNING: This conflates within- and between-person effects
m2 <- lmer(
  cognition ~ time + height + (1 + time | participant_id),
  data = df_long,
  REML = TRUE
)

summary(m2)

# Add decomposed TVC: within-person (cwc) + between-person (pm)
# This is the recommended approach for causal inference
m3 <- lmer(
  cognition ~ time + height_cwc + height_pm + (1 + time | participant_id),
  data = df_long,
  REML = TRUE
)

# Generate summary
summary(m3)

# Create formatted output table using broom.mixed + gt
model_summary_table <- broom.mixed::tidy(m3, effects = "fixed") %>%
  select(term, estimate, std.error, statistic, p.value) %>%
  gt() %>%
  tab_header(title = "Model 3: Decomposed Time-Varying Covariate") %>%
  fmt_number(columns = c(estimate, std.error, statistic, p.value), decimals = 3) %>%
  cols_label(
    term = "Term",
    estimate = "Estimate",
    std.error = "Std. Error",
    statistic = "t-value",
    p.value = "P-Value"
  )

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

# Generate alternative summary with custom labels using broom.mixed + gt
lmm_tvc_table <- broom.mixed::tidy(m3, effects = "fixed") %>%
  mutate(term = case_when(
    term == "(Intercept)" ~ "Intercept",
    term == "time" ~ "Time",
    term == "height_cwc" ~ "Height (Within)",
    term == "height_pm" ~ "Height (Between)",
    TRUE ~ term
  )) %>%
  select(term, estimate, std.error, statistic, p.value) %>%
  gt() %>%
  tab_header(title = "LMM with Time-Varying Covariate") %>%
  fmt_number(columns = c(estimate, std.error, statistic, p.value), decimals = 3) %>%
  cols_label(
    term = "Predictor",
    estimate = "Estimate",
    std.error = "SE",
    statistic = "t-value",
    p.value = "P-Value"
  )

gt::gtsave(lmm_tvc_table, filename = "lmm_tvc_results.html")

# Compare models using likelihood ratio test and information criteria
model_comparison <- anova(m1, m2, m3)

# Create comparison table
comparison_df <- data.frame(
  Model = c("M1: Time Only", "M2: Time + Raw TVC", "M3: Time + Decomposed TVC"),
  AIC = c(AIC(m1), AIC(m2), AIC(m3)),
  BIC = c(BIC(m1), BIC(m2), BIC(m3)),
  LogLik = c(logLik(m1), logLik(m2), logLik(m3))
)

comparison_table <- comparison_df %>%
  gt() %>%
  tab_header(title = "Model Comparison: Time-Varying Covariate Specifications") %>%
  fmt_number(columns = c(AIC, BIC, LogLik), decimals = 2)

gt::gtsave(comparison_table, filename = "model_comparison.html")

# Create visualization comparing within and between effects
# First, get predicted values
df_long <- df_long %>%
  mutate(predicted = predict(m3))

# Sample participants for visualization
set.seed(123)
sample_ids <- sample(unique(df_long$participant_id), min(100, length(unique(df_long$participant_id))))
df_plot <- df_long %>% filter(participant_id %in% sample_ids)

# Create spaghetti plot colored by height level
visualization <- ggplot(df_plot, aes(x = time, y = cognition, group = participant_id)) +
  geom_line(aes(color = height_pm), alpha = 0.5) +
  geom_point(aes(color = height_pm), alpha = 0.3, size = 1) +
  scale_color_viridis_c(name = "Person-Mean\nHeight", option = "plasma") +
  geom_smooth(aes(group = 1), method = "lm", color = "black", linewidth = 1.5, se = TRUE) +
  scale_x_continuous(
    breaks = 0:3,
    labels = c("Baseline", "Year 2", "Year 4", "Year 6")
  ) +
  labs(
    title = "Cognition Trajectories by Average Height",
    subtitle = "Line color indicates person-mean height (between-person effect)",
    x = "Assessment Wave",
    y = "Fluid Cognition"
  ) +
  theme_minimal() +
  theme(legend.position = "right")

visualization

ggsave(
  filename = "visualization.png",
  plot = visualization,
  width = 10, height = 6, dpi = 300
)