Required libraries

library(readr)
library(dplyr)
library(tidyr)
library(ggplot2)
library(rsample)
library(parsnip)
library(yardstick)
library(ggplot2)
library(tune)
library(recipes)
library(workflows)
library(dials)
library(xgboost)

Random seed

set.seed(123)

Visualizing number of missing values.

First, I had to fill in many missing values. I modified the approach in the article Missing Value Visualization to create a visualization of the percent of missing values for three dataset columns.

water_total_rows = nrow(water)

water_missing <- water %>%
  pivot_longer(everything(), names_to = "name", values_to = "value") %>%
  mutate(is_missing = is.na(value)) %>%
  group_by(name, is_missing) %>%
  summarize(num_missing = n()) %>%
  filter(is_missing) %>%
  transmute(
    percent_missing = num_missing / water_total_rows * 100
  ) %>%
  rename(variable = name) %>%
  arrange(desc(percent_missing))

knitr::kable(water_missing)

ggplot(water_missing, aes(x = variable, y = percent_missing)) +
  geom_col() +
  labs(
    title = "Percentage of values missing",
    subtitle = "Higher is worse"
  )

Transform Potability into a factor

Second, I changed the potability column into a factor with two levels of potable and not_potable. I am training the models to predict when a water sample is drinkable, so the first level is the event of interest.

water_2 <- water %>%
  mutate(potable = factor(
      case_when(
        Potability == 0 ~ "not_potable",
        Potability == 1 ~ "potable"
      ),
      levels = c("potable", "not_potable")
    )
  ) %>%
  select(-Potability)

knitr::kable(head(water_2))

Train/test split

I split the data into 75% for training and 25% for testing. I use the same train/test split across all models in this comparison.

water_split <- initial_split(water_2, prop = 0.75, strata = potable)
water_train <- training(water_split)

Filling missing values with k nearest neighbors (kNN)

I use kNN imputation to fill the missing values in the data.

clean_water_recipe <- recipe(potable ~ ., data = water_train) %>%
  step_impute_knn(all_numeric(), neighbors = 10)

How many observations are in class?

More of the observations are of non-potable water samples. I suspect this made it harder to train the models on examples of potable water.

clean_water_class_count <- clean_water %>%
  group_by(potable) %>%
  summarize(class_count = n())

knitr::kable(clean_water_class_count)

ggplot(clean_water_class_count, aes(x = potable, y = class_count)) +
  geom_col() +
  labs(
      title = "Class count in entire dataset",
      subtitle = "Equal heights are better"
    )

Distributions variables depending on potability

After filling the missing values with kNN, I made violin plots lining up distributions for each variable depending on whether they were from potable samples.

ggplot(clean_water, aes(x = potable, y = ph)) +
  geom_violin() +
  ggtitle("pH")

ggplot(clean_water, aes(x = potable, y = Hardness)) +
  geom_violin() +
  ggtitle("Hardness")

ggplot(clean_water, aes(x = potable, y = Solids)) +
  geom_violin() +
  ggtitle("Solids")

ggplot(clean_water, aes(x = potable, y = Chloramines)) +
  geom_violin() +
  ggtitle("Chloramines")

ggplot(clean_water, aes(x = potable, y = Sulfate)) +
  geom_violin() +
  ggtitle("Sulfate")

ggplot(clean_water, aes(x = potable, y = Conductivity)) +
  geom_violin() +
  ggtitle("Conductivity")

ggplot(clean_water, aes(x = potable, y = Organic_carbon)) +
  geom_violin() +
  ggtitle("Organic Carbon")

ggplot(clean_water, aes(x = potable, y = Trihalomethanes)) +
  geom_violin() +
  ggtitle("Trihalomethanes")

ggplot(clean_water, aes(x = potable, y = Turbidity)) +
  geom_violin() +
  ggtitle("Turbidity")

k-fold cross validation

I use 20 folds stratified by the potable outcome variable for cross-validation.

water_folds <- vfold_cv(water_train, v = 20, strata = potable)

Evaluation metrics

I gather ROC AUC, sensitivity, and specificity metrics for each fold and hyperparameter combination during training.

common_metric_set <- metric_set(roc_auc, sens, spec)

Create and fit the logistic model

logistic_model <- logistic_reg() %>%
  set_engine("glm") %>%
  set_mode("classification")

logistic_last_fit <- logistic_model %>%
  last_fit(potable ~ ., split = water_split)

Evaluate logistic performance

The ROC AUC curve lines up mainly along the diagonal for the logistic classifier, so the classification performance is about the same as a coin toss. That’s pretty bad and won’t be hard to exceed using the other models.

logistic_last_fit_metrics <- logistic_last_fit %>%
  collect_metrics()

logistic_last_fit_roc_auc <- logistic_last_fit_metrics %>%
  filter(.metric == "roc_auc") %>%
  pull(.estimate)

knitr::kable(logistic_last_fit_metrics)

logistic_last_fit_results <- logistic_last_fit %>%
  collect_predictions()

logistic_last_fit_results %>%
  roc_curve(truth = potable, .pred_potable) %>%
  autoplot()

Decision tree workflow and grid

decision_tree_model <- decision_tree(
    cost_complexity = tune(),
    tree_depth = tune(),
    min_n = tune()
  ) %>%
  set_engine("rpart") %>%
  set_mode("classification")

decision_tree_workflow <- workflow() %>%
  add_model(decision_tree_model) %>%
  add_recipe(clean_water_recipe)

decision_tree_grid <- grid_random(parameters(decision_tree_model), size = 10)

Tune the decision tree

doParallel::registerDoParallel()

decision_tree_tuning <- decision_tree_workflow %>%
  tune_grid(
    resamples = water_folds,
    grid = decision_tree_grid,
    metrics = common_metric_set
  )

Fit the best decision tree

I select the top-performing decision tree based on the ROC AUC metric.

best_decision_tree_model <- decision_tree_tuning %>%
  select_best(metric = "roc_auc")

final_decision_tree_workflow <- decision_tree_workflow %>%
  finalize_workflow(best_decision_tree_model)

decision_tree_final_fit <- final_decision_tree_workflow %>%
  last_fit(split = water_split)

Evaluate decision tree performance

The decision tree’s ROC AUC is 0.62, which, while better than the logistic classifier, is still not very good.

decision_tree_final_fit_metrics <- decision_tree_final_fit %>%
  collect_metrics()

decision_tree_final_fit_roc_auc <- decision_tree_final_fit_metrics %>%
  filter(.metric == "roc_auc") %>%
  pull(.estimate)
  
knitr::kable(decision_tree_final_fit_metrics)

decision_tree_final_fit %>%
  collect_predictions() %>%
  roc_curve(truth = potable, .pred_potable) %>%
  autoplot()

Tune the xgboost

doParallel::registerDoParallel()

xgboost_tuning <- xgboost_workflow %>%
  tune_grid(
    resamples = water_folds,
    grid = xgboost_grid,
    metrics = common_metric_set
  )

Finalize the best xgboost

best_xgboost_model <- xgboost_tuning %>%
  select_best(metric = "roc_auc")

final_xgboost_workflow <- xgboost_workflow %>%
  finalize_workflow(best_xgboost_model)

xgboost_final_fit <- final_xgboost_workflow %>%
  last_fit(split = water_split)

Evaluate xgboost performance

With a ROC AUC of 0.65, the xgboost outperforms the decision tree, but not by much.

xgboost_final_fit_metrics <- xgboost_final_fit %>%
  collect_metrics()

xgboost_final_fit_roc_auc <- xgboost_final_fit_metrics %>%
  filter(.metric == "roc_auc") %>%
  pull(.estimate)

knitr::kable(xgboost_final_fit_metrics)

xgboost_final_fit %>%
  collect_predictions() %>%
  roc_curve(truth = potable, .pred_potable) %>%
  autoplot()