library(pROC)
# Logistic regression
logit <- glm(SWP ~ Age + Sex + Fulltime + CTW + JobDemands +
JobControl + JobSupport + MentalWB, data=train, family=binomial)
p_logit <- predict(logit, newdata=test, type="response")
# Random forest
rf <- randomForest(factor(SWP) ~ Age + Sex + Fulltime + CTW + JobDemands +
JobControl + JobSupport + MentalWB,
data=train_f, ntree=300)
p_rf <- predict(rf, newdata=test_f, type="prob")[, "1"]
# XGBoost
xgb_fit <- xgb.train(
params = list(objective="binary:logistic", eta=0.1, max_depth=4,
eval_metric="auc"),
data = dtrain, nrounds=200, verbose=0)
p_xgb <- predict(xgb_fit, dtest)
# ROC curves
roc_logit <- roc(test$SWP, p_logit, quiet=TRUE)
roc_rf <- roc(test$SWP, p_rf, quiet=TRUE)
roc_xgb <- roc(test$SWP, p_xgb, quiet=TRUE)
# Plot
plot(roc_logit, col="#2980b9", lwd=2,
main="ROC curves β SWP classification",
xlab="1 β Specificity (FPR)", ylab="Sensitivity (TPR)")
lines(roc_rf, col="#27ae60", lwd=2)
lines(roc_xgb, col="#e74c3c", lwd=2)
abline(0, 1, lty=2, col="#888")
legend("bottomright",
legend=c(paste0("Logistic (AUC=", round(auc(roc_logit),3), ")"),
paste0("Random Forest (AUC=", round(auc(roc_rf),3), ")"),
paste0("XGBoost (AUC=", round(auc(roc_xgb),3), ")")),
col=c("#2980b9","#27ae60","#e74c3c"), lwd=2, bty="n")