XGBoost Machine Learning for Climate-Health Outcome Prediction

Trains an XGBoost gradient-boosted tree model to predict health outcomes (disease counts, mortality, hospitalizations) from climate and socioeconomic features. Accepts any aggregated data frame produced by the climasus4r pipeline (or a plain data.frame). Uses k-fold cross-validation to select the optimal number of trees (nrounds) and returns out-of-fold (OOF) predictions alongside the final model trained on full data.

Usage

sus_mod_ml(
  df,
  outcome_col,
  feature_cols = NULL,
  id_col = NULL,
  objective = "count:poisson",
  nrounds = 500L,
  max_depth = 6L,
  eta = 0.05,
  subsample = 0.8,
  colsample_bytree = 0.8,
  min_child_weight = 1,
  nfold = 5L,
  early_stopping = 50L,
  seed = 42L,
  lang = "pt",
  verbose = TRUE
)

Arguments

df: A data frame (plain or climasus_df) with one row per observation (e.g., city × day or city × month). Must contain outcome_col and at least one numeric feature column.
outcome_col: Character. Name of the column to predict (the target variable). Must be numeric.
feature_cols: Character vector or NULL. Names of columns to use as predictors. When NULL (default), all numeric columns except outcome_col and id_col are used automatically.
id_col: Character or NULL. City or group identifier column used for group-aware cross-validation: each group is kept entirely in one fold, preventing data leakage across cities. When NULL, standard random k-fold is used. Strongly recommended when data has a city dimension.
objective: Character. XGBoost objective function. One of "count:poisson" (default), "reg:squarederror", or "binary:logistic". See Details.
nrounds: Integer. Maximum number of boosting rounds. The optimal value is selected via early stopping during cross-validation. Default 500L.
max_depth: Integer. Maximum tree depth. Default 6L. Increase for more complex patterns; decrease to reduce overfitting.
eta: Numeric. Learning rate (step size shrinkage). Default 0.05. Smaller values require more rounds but often generalize better.
subsample: Numeric ∈ (0, 1]. Row subsampling ratio per tree. Default 0.8.
colsample_bytree: Numeric ∈ (0, 1]. Column subsampling ratio per tree. Default 0.8.
min_child_weight: Numeric. Minimum sum of instance weight in a leaf. Default 1. Increase to prevent learning rare patterns.
nfold: Integer. Number of cross-validation folds. Default 5L.
early_stopping: Integer. Stop training if the CV metric does not improve for this many consecutive rounds. Default 50L.
seed: Integer. Random seed for reproducibility. Default 42L.
lang: Character. Language for messages: "pt" (default), "en", "es".
verbose: Logical. Print progress messages. Default TRUE.

Value

A climasus_ml object (named list) with:

$predictions: Tibble with one row per (non-NA) observation: observed, fitted (train-set prediction), cv_predicted (out-of-fold CV prediction), residual (observed - cv_predicted), and the id_col and outcome_col columns if present.
$importance: Tibble sorted by Gain (descending): Feature, Gain (average loss reduction per split), Cover (average number of samples per split), Frequency (proportion of splits using this feature). These are the standard XGBoost importance metrics.
$performance: Named list: RMSE_train, MAE_train, R2_train, RMSE_cv, MAE_cv, R2_cv, Pearson_cv (correlation between observed and OOF predictions), and best_nrounds.
$model: The final XGBoost model (class xgb.Booster) trained on the full dataset using best_nrounds. Pass to predict() to generate forecasts for new data (see predict.climasus_ml()).
$cv_log: Data frame from xgb.cv() with per-round train and test metrics. Useful for plotting the learning curve.
$meta: Named list of all parameters used in this call.

Details

This function complements the DLNM epidemiological pipeline (sus_mod_dlnm()): DLNM models the exposure-response relationship and provides causal inference; sus_mod_ml() focuses on predictive accuracy across a wider feature space and can incorporate lagged climate variables, socioeconomic covariates, and spatial predictors simultaneously.

XGBoost objective

The objective controls the loss function and the scale of predictions:

"count:poisson": Poisson log-linear model for non-negative integer counts (deaths, hospitalizations). Predictions are in count scale (predicted lambda). Default — appropriate for most DATASUS outcomes.
"reg:squarederror": Squared-error regression for continuous outcomes (rates, indices). Predictions are in the same scale as the outcome.
"binary:logistic": Binary outcome (0/1). Predictions are probabilities ∈ (0, 1).

Examples

if (FALSE) { # \dontrun{
# From a climasus4r pipeline output (aggregated city x day data with
# climate columns from sus_climate_aggregate())
ml <- sus_mod_ml(
  df          = df_aggregated,
  outcome_col = "n_obitos",
  id_col      = "name_muni",        # group-aware CV by city
  objective   = "count:poisson",
  nfold       = 5L,
  lang        = "pt"
)
print(ml)
tidy(ml)    # predictions tibble

# Forecast for a new climate scenario
new_climate <- data.frame(tair_max_c = 36, rh_pct_mean = 60, ...)
predict(ml, newdata = new_climate)
} # }