Introduction

This project develops a supervised learning model to predict ISNCSCI motor score at 12 months post-injury using data from the first 30 days post-injury on an EMSCI-like simulated dataset. Limitations are discussed at the end.

## [1] "Packages loaded."

## [1] "Registered 12 workers for parallel processing."

Simulated data

The simulated data is a cohort based on descriptive statistics of EMSCI subjects reported by Bourguignon et al., 2022 (BMC Med). We aim to model realistic heterogeneity by creating several recovery phenotypes with high variance and potential plateauing that consider known SCI pathophysiology and recovery patterns.

## [1] "EMSCI-like cohort data simulation complete."

## [1] "Total patients: 1500"

## [1] "Total observations: 5516"

## [1] "Mean assessments per patient: 3.68"

Feature consolidation

We pass all features to the stacked ensemble, relying on the inherent feature selection capabilities of the base learners.

## [1] "Modeling data created with static and trajectory features."

Cross-validation strategy and preprocessing recipe

We define a grouped 5-fold cross-validation using 3 repeats and a preprocessing recipe to clean the predictors.

# Use 5-fold grouped CV by center, repeated 3 times for robust variance estimation
set.seed(123)
cv_splits <- group_vfold_cv(
  modeling_data,
  group = center_id,
  v = 5, # 10-fold may also work if the partitition is balanced
  repeats = 3 
)

# Create preprocessing recipe
emsci_recipe <- recipe(motor_score_12m ~ ., data = modeling_data) %>%
  update_role(patient_id, new_role = "ID") %>%
  update_role(patient_id_global, new_role = "ID") %>%
  update_role(center_id, new_role = "ID") %>%
  step_dummy(all_nominal_predictors(), one_hot = FALSE) %>% 
  step_normalize(all_numeric_predictors(), -all_of(c("patient_id", "patient_id_global", "center_id"))) %>% 
  step_integer(all_logical_predictors()) %>%
  step_zv(all_predictors())

print("Cross-validation and preprocessing recipe defined.")

## [1] "Cross-validation and preprocessing recipe defined."

We define a set of diverse base-learner models and computationally-efficient hyperparameter tuning grids for each.

# Define model specifications

# 1. ElasticNet (linear regularized)
set.seed(123)
elnet_spec <- linear_reg(
  penalty = tune(), 
  mixture = tune()
) %>%
  set_engine("glmnet") %>%
  set_mode("regression")

# 2. Random Forest
set.seed(123)
rf_spec <- rand_forest(
  mtry = tune(),
  min_n = tune(),
  trees = 2000                # Reduced for resource efficiency; increase if possible
) %>%
  set_engine("ranger", num.threads = 1) %>% 
  set_mode("regression")

# 3. XGBoost
# Early stopping and validation reduces unnecessary computation
set.seed(123)
xgb_spec <- boost_tree(
  trees = 500,                # Reduced for tuning efficiency
  tree_depth = tune(),
  learn_rate = tune(),  
  min_n = tune(),
  loss_reduction = tune(),
  sample_size = tune()
) %>%
  set_engine(
    "xgboost", 
    nthread = 1,
    validation = 0.2,          # Early stopping uses validation split
    early_stop = 50            # Stop after 50 rounds without improvement
  ) %>%
  set_mode("regression")

# 4. SVM (Radial Basis Function kernel)
set.seed(123)
svm_spec <- svm_rbf(cost = tune(), rbf_sigma = tune()) %>%
  set_engine("kernlab") %>%
  set_mode("regression")

# 5. MARS (Multivariate Adaptive Regression Splines)
set.seed(123)
mars_spec <- mars(num_terms = tune(), prod_degree = tune()) %>%
  set_engine("earth") %>%
  set_mode("regression")

print("Base-learner specifications set.")

## [1] "Base-learner specifications set."

# Random grids for high-dimensional or sensitive hyperparameters (ElasticNet, XGBoost, SVM)
# Regular grids for robust, fewer hyperparameters (RF, MARS)

# ElasticNet: 15-point random grid
set.seed(123)
elnet_grid <- grid_random(
  penalty(range = c(-5, 1)),
  mixture(range = c(0, 1)),
  size = 15                  # reduced for computation efficiency, increase if possible for all models
)

# Random Forest: Systematic grid (15 points)
set.seed(123)
rf_grid <- grid_regular(
  mtry(range = c(5, 20)),
  min_n(range = c(5, 15)),
  levels = c(mtry = 5, min_n = 3)
)

# XGBoost: 15-point random grid
set.seed(123)
xgb_grid <- grid_random(
  tree_depth(range = c(3, 10)),
  min_n(range = c(2, 20)),
  learn_rate(range = c(-3, -0.5)), # Tuning log10(learning_rate) from 0.001 to 0.3
  loss_reduction(range = c(-5, 1)),
  sample_prop(range = c(0.5, 1.0)),
  size = 15
)

# SVM: 15-point random grid
set.seed(123)
svm_grid <- grid_random(
  cost(range = c(-5, 5)),
  rbf_sigma(range = c(-5, 0)),
  size = 15
)

# MARS: Systematic grid (15 points)
set.seed(123)
mars_grid <- grid_regular(
  num_terms(range = c(5, 25)),
  prod_degree(range = c(1L, 2L)),
  levels = c(num_terms = 5, prod_degree = 3)
)

print("Hyperparameter grids finalized.")

## [1] "Hyperparameter grids finalized."

We used a workflow set to tune all five models across our grouped CV splits and save the out-of-fold predictions to be used inputs for the final ensemble.

## Base learners tuning: 526.64 sec elapsed

## [1] "Base learners tuned."

Stacked ensemble creation

We create the stacked ensemble by adding the tuned base learners as candidates and blending their out-of-fold predictions using a non-negative LASSO meta-learner to find optimal, interpretable weights. Then, we fit the complete stacked model to the full training set.

# Create and fit the stacked ensemble
tic("Stacking")

# Initialize stack
set.seed(123)
emsci_stack <- stacks() %>%
  add_candidates(model_set_results) # Add all tuned models at once

print(paste("Stack initialized with", ncol(emsci_stack), "candidate models"))

## [1] "Stack initialized with 70 candidate models"

# Blend the stack (fit meta-learner with non-negative LASSO)
set.seed(123)
emsci_stack_blended <- emsci_stack %>%
  blend_predictions(
    penalty = 10^seq(-4, -0.5, length.out = 20),
    mixture = 1,  # LASSO
    non_negative = TRUE,                # enforce interpretability
    metric = metric_set(rmse)           # RMSE prioritizes predictive accuracy
  )

print("Stack blended successfully")

## [1] "Stack blended successfully"

# Fit the final stacked model on the full training data
set.seed(123)
emsci_stack_fitted <- emsci_stack_blended %>%
  fit_members()

toc()

## Stacking: 58.93 sec elapsed

print("Final stacked ensemble fitted.")

## [1] "Final stacked ensemble fitted."

Meta-learner evaluation (two-stage CV)

We perform a two-stage evaluation using an independent set of 5-fold grouped CV splits. We use the base-learner predictions as inputs and the optimal non-negative LASSO penalty found during blending and generate the final out-of-sample performance metrics.

# Obtain out-of-fold predictions and performance metrics
tic("Model evaluation")

# Create the L1 dataset (base learner predictions + IDs)
l1_data <- emsci_stack %>%
  mutate(.row = row_number()) %>%
  left_join(
    modeling_data %>% 
      mutate(.row = row_number()) %>% 
      select(.row, patient_id_global, center_id),
    by = ".row"
  )

print(paste("L1 dataset created:", nrow(l1_data), "patients"))

## [1] "L1 dataset created: 1500 patients"

# Define NEW, INDEPENDENT CV folds for meta-learner evaluation
set.seed(456) # Use a different seed for independence
l1_folds <- group_vfold_cv(l1_data, group = center_id, v = 5)

# Get optimal penalty from Stage 1
best_penalty <- emsci_stack_blended$penalty$penalty

print(paste("Using meta-learner penalty from Stage 1:", best_penalty))

## [1] "Using meta-learner penalty from Stage 1: 0.206913808111479"

# Define meta-learner spec
# Use !!  to inject the value, prevents silent fit_resamples scope error
meta_spec <- linear_reg(
  penalty = !!best_penalty,
  mixture = 1 # LASSO
) %>%
  set_engine("glmnet", lower.limits = 0) # Non-negative 

# Define meta-learner recipe
meta_recipe <- recipe(motor_score_12m ~ ., data = l1_data) %>%
  update_role(center_id, new_role = "ID") %>%
  update_role(patient_id_global, new_role = "ID") %>%
  update_role(.row, new_role = "ID")

# Create workflow
eval_metrics <- metric_set(rmse, rsq, mae)
meta_wf <- workflow() %>%
  add_recipe(meta_recipe) %>%
  add_model(meta_spec)

# Evaluate with CV
set.seed(456)
tic("Meta-learner CV")
meta_cv_results <- fit_resamples(
  meta_wf,
  resamples = l1_folds,
  metrics = eval_metrics, # Pass the metrics
  control = control_resamples(save_pred = TRUE)
)
toc()

## Meta-learner CV: 23.01 sec elapsed

print("Meta-learner CV complete")

## [1] "Meta-learner CV complete"

# Extract CV predictions
cv_predictions <- tidyr::unnest(meta_cv_results, .predictions)

# Create final plotting data
plot_data_cv <- cv_predictions %>%
  left_join(
    modeling_data %>% 
      mutate(.row = row_number()) %>%
      select(.row, patient_id_global, baseline_ais),
    by = ".row"
  ) %>%
  select(patient_id_global, motor_score_12m, .pred, baseline_ais)

# Calculate final CV metrics
final_cv_metrics <- collect_metrics(meta_cv_results)

print(paste("Out-of-fold predictions extracted for", 
            nrow(plot_data_cv), "patients"))

## [1] "Out-of-fold predictions extracted for 1500 patients"

print("Performance metrics and predictions ready.")

## [1] "Performance metrics and predictions ready."

## [1] "DALEX explainer created successfully"

## Permutation importance: 1608.05 sec elapsed

## SHAP calculation (full sample): 4522.95 sec elapsed

## These are not valid conclusions. This is a programming exercise. For completeness, here are the results:

## Overall performance (5-fold CV):

##    - RMSE: 14.75 (13.76, 15.75)

##    - MAE: 11.28 (10.52, 12.05)

##    - R-squared: 0.739 (0.704, 0.774)

## Calibration:

##    - Calibration slope: 1.012 (ideal = 1)

##    - Calibration intercept: -0.326 (ideal = 0)

## Clinical utility:

##    - Predictions within MCID (±4 points): 23.3%

## Bland-Altman analysis:

##    - Mean bias: 0.02 points

##    - 95% Limits of Agreement: [-29.18, 29.23]

## Meta-Learner composition:

##    - Total members selected: 14 (out of 69 candidates)

##    - Dominant model type: MARS (48.2% of total weight)

## Subgroup performance:

##    - AIS A (n=617): RMSE=14.76, R-squared=0.718
##    - AIS B (n=200): RMSE=15.83, R-squared=0.722
##    - AIS C (n=284): RMSE=14.75, R-squared=0.753
##    - AIS D (n=399): RMSE=14.72, R-squared=0.754

## R version: R version 4.5.1 (2025-06-13 ucrt)

## Platform: x86_64-w64-mingw32

## Key ML packages:

## tidymodels 1.4.1 
## tune 2.0.1 
## stacks 1.1.1 
## ranger 0.17.0 
## xgboost 1.7.11.1 
## earth 5.3.4 
## DALEX 2.5.3 
## fastshap 0.1.1