Time-Varying Covariates (GEE)

### Load necessary libraries
library(NBDCtools)    # ABCD data access helper
library(arrow)        # Read Parquet files
library(tidyverse)    # Data wrangling & visualization
library(gtsummary)    # Summary tables
library(rstatix)      # Statistical tests in tidy format
library(geepack)      # Generalized Estimating Equations (GEE) analysis
library(broom)        # Organizing model outputs
library(gt)           # Presentation-ready display tables

### Load harmonized ABCD data required for this analysis
requested_vars <- c(
    "ab_g_dyn__design_site",
    "ab_g_stc__design_id__fam",
    "ph_p_sds__dims_001",
    "mh_p_cbcl__dsm__anx_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
)

# Data wrangling: clean, restructure, and recode sleep variable
df_long <- abcd_data %>%
  filter(session_id %in% c("ses-00A", "ses-01A", "ses-02A", "ses-03A")) %>%  # Keep Baseline - Year 3
  arrange(participant_id, session_id) %>%
   mutate(
    participant_id = factor(participant_id), # Convert participant_id to a factor
    session_id = factor(session_id,
                        levels = c("ses-00A", "ses-01A", "ses-02A", "ses-03A"),
                        labels = c("Baseline", "Year_1", "Year_2", "Year_3")),  # Label sessions

    ab_g_dyn__design_site = factor(ab_g_dyn__design_site),  # Convert site to a factor
    ab_g_stc__design_id__fam = factor(ab_g_stc__design_id__fam), # Convert family id to a factor
    sleep_hours = as.numeric(ph_p_sds__dims_001),
    anxiety = as.numeric(mh_p_cbcl__dsm__anx_tscore)  # Convert to numeric
  ) %>%
  filter(sleep_hours != 999, anxiety != 999) %>%  # Remove "Don't know" responses
  mutate(
    sleep_binary = ifelse(sleep_hours == 1, 1, 0)  # Recode: 9-11 hrs = 1, others = 0
) %>%
  rename(  # Rename for simplicity
    site = ab_g_dyn__design_site,
    family_id = ab_g_stc__design_id__fam,
  ) %>%
  drop_na()  # Remove missing values

# Summarize 'sleep_binary' and 'strength_exercise' across session_ids

descriptives_table <- df_long %>%
    select(session_id, sleep_binary, anxiety) %>%
    tbl_summary(
        by = session_id,
        missing = "no",
        label = list(
          sleep_binary ~ "Sufficient Sleep (9-11 hrs)",
          anxiety ~ "Anxiety"
        ),
        statistic = list(
            all_categorical() ~ "{n} ({p}%)",
            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 GEE model with time-varying anxiety covariate
model <- geeglm(sleep_binary ~ session_id + site + anxiety,
  id = participant_id,
  data = df_long,
  family = binomial(link = "logit"),
  corstr = "exchangeable"
)

# Generate summary table from model coefficients
coefs <- as.data.frame(summary(model)$coefficients)
model_summary <- data.frame(
  Parameter = rownames(coefs),
  Estimate = coefs$Estimate,
  SE = coefs$Std.err,
  Wald = coefs$Wald,
  p_value = coefs[, "Pr(>|W|)"]
) %>%
  gt() %>%
  tab_header(title = "GEE Model: Sleep ~ Time + Site + Anxiety (TVC)") %>%
  fmt_number(columns = c(Estimate, SE, Wald), decimals = 3) %>%
  fmt_number(columns = p_value, decimals = 4) %>%
  cols_label(
    Parameter = "Parameter",
    Estimate = "Estimate",
    SE = "Robust SE",
    Wald = "Wald",
    p_value = "p-value"
  )

model_summary

# Save as standalone HTML
gt::gtsave(model_summary, filename = "model_summary.html")

# Create GEE diagnostics data
diagnostics_data <- data.frame(
  Characteristic = c(
    "Correlation Structure",
    "Correlation Parameter (α)",
    "Scale Parameter",
    "Number of Participants",
    "Total Observations"
  ),
  Value = c(
    model$corstr,
    round(model$geese$alpha, 3),
    round(summary(model)$dispersion$Estimate, 3),
    length(unique(df_long$participant_id)),
    nrow(df_long)
  )
)

# Format diagnostics table
diagnostics_table <- diagnostics_data %>%
  gt::gt() %>%
  gt::tab_header(title = "GEE Model Diagnostics")

diagnostics_table

# Save diagnostics table
gt::gtsave(diagnostics_table, filename = "model_diagnostics.html")

# Generate predicted probabilities across anxiety range
# Use the anxiety coefficient to show marginal effect at reference levels
coefs <- summary(model)$coefficients
intercept <- coefs["(Intercept)", "Estimate"]
beta_anxiety <- coefs["anxiety", "Estimate"]
se_anxiety <- coefs["anxiety", "Std.err"]

anxiety_range <- seq(min(df_long$anxiety, na.rm = TRUE),
                     max(df_long$anxiety, na.rm = TRUE), length.out = 50)
linpred <- intercept + beta_anxiety * anxiety_range
se_pred <- se_anxiety * abs(anxiety_range - mean(anxiety_range))

preds <- data.frame(
  anxiety = anxiety_range,
  predicted = plogis(linpred),
  conf.low = plogis(linpred - 1.96 * se_pred),
  conf.high = plogis(linpred + 1.96 * se_pred)
)

# Plot predicted probabilities
visualization <- ggplot(preds, aes(x = anxiety, y = predicted)) +
  geom_line(color = "darkblue", linewidth = 1) +
  geom_ribbon(aes(ymin = conf.low, ymax = conf.high), alpha = 0.2, fill = "darkblue") +
  labs(title = "Effect of Anxiety on Sleep Sufficiency",
       x = "Anxiety Level",
       y = "Predicted Probability of Sufficient Sleep") +
  theme_classic() +
  theme(plot.title = element_text(face = "bold", size = 14),
        axis.title = element_text(size = 12))

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