Is This Water Safe to Drink?

# ---------------------------------------------------------------------------
# safe_water.py
#
# Capstone project: predicting, from a quick field-kit test, whether a
# drinking-water source will fail the laboratory's safety test - so unsafe
# sources can be flagged the same day instead of two days later.
#
# CBSE Class XII, Artificial Intelligence (843), Part C - Capstone Project.
# SDG 6 (Clean Water and Sanitation) and SDG 3 (Good Health and Well-being).
#
# Data Science Methodology, 10 steps. CLASSIFICATION (Safe / Unsafe). The
# field data has gaps, so the project compares four ways of handling
# missing values - the Unit 1 topic - with stratified 5-fold cross
# validation, before training the final Random Forest.
#
# Libraries: pandas, numpy, matplotlib, scikit-learn, joblib.
# ---------------------------------------------------------------------------

# %% Step 1 - Business Understanding
print("=" * 64)
print("IS THIS WATER SAFE TO DRINK?")
print("=" * 64)
print("Problem : the lab's verdict takes two days; people keep drinking meanwhile")
print("Goal    : predict the lab verdict from field-kit readings, the same day")
print()

# %% Step 2 - Analytic Approach
print("Analytic approach : predictive - classification (Safe or Unsafe)")
print()

# %% Step 3 - Data Requirements
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.impute import SimpleImputer
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import confusion_matrix, accuracy_score, precision_score, recall_score, f1_score
import joblib

pd.set_option("display.width", 140)
pd.set_option("display.max_columns", 20)

# %% Step 4 - Data Collection
data = pd.read_csv("water_samples.csv")
print("Samples collected :", len(data))
print(data.head(6))
print()

# %% Step 5 - Data Understanding
print("--- Step 5: Data Understanding ---")
missing = data.isnull().sum()
print("Missing values in each column:")
print(missing.to_string())
print()
print("Lab verdicts:")
print(data["verdict"].value_counts().to_string())
print("Share unsafe :", round((data["verdict"] == "Unsafe").mean() * 100, 1), "%")
print()
print("Share unsafe by kind of source (%):")
by_source = ((data["verdict"] == "Unsafe").groupby(data["source"]).mean() * 100).sort_values()
print(by_source.round(1).to_string())
print()

# Chart 1: how much is missing in each column.
m = (missing[missing > 0] / len(data) * 100).sort_values()
plt.figure(figsize=(7, 3.8))
plt.barh(m.index, m.values, color="#f28e2b")
plt.title("Missing values in the field data")
plt.xlabel("% of samples with no reading")
plt.tight_layout()
plt.savefig("chart1_missing.png")
plt.close()

# Chart 2: unsafe share by kind of source.
plt.figure(figsize=(7, 4))
plt.barh(by_source.index, by_source.values, color="#4e9fd1")
plt.title("Share of samples the lab found unsafe")
plt.xlabel("% unsafe")
plt.tight_layout()
plt.savefig("chart2_by_source.png")
plt.close()

# Chart 3: turbidity for safe and unsafe samples (log scale).
turb = pd.to_numeric(data["turbidity_ntu"])
plt.figure(figsize=(7, 4.2))
plt.boxplot([turb[data["verdict"] == "Safe"].dropna(), turb[data["verdict"] == "Unsafe"].dropna()],
            labels=["Safe", "Unsafe"])
plt.yscale("log")
plt.axhline(5, color="#e15759", linestyle="--", label="IS 10500 limit: 5 NTU")
plt.title("Turbidity of safe and unsafe samples")
plt.ylabel("Turbidity (NTU, log scale)")
plt.legend()
plt.tight_layout()
plt.savefig("chart3_turbidity.png")
plt.close()

# %% Step 6 - Data Preparation
print("--- Step 6: Data Preparation ---")
before = len(data)
data = data.drop_duplicates()
print("Duplicate samples removed :", before - len(data))
for col in ["fluoride_mg_l", "nitrate_mg_l", "turbidity_ntu"]:
    data[col] = pd.to_numeric(data[col])
data["unsafe"] = (data["verdict"] == "Unsafe").astype(int)
data = pd.get_dummies(data, columns=["source"], dtype=int)
features = ["ph", "tds_mg_l", "turbidity_ntu", "nitrate_mg_l", "fluoride_mg_l", "chloride_mg_l",
            "residual_chlorine_mg_l", "source_Hand pump", "source_Open well", "source_Piped supply",
            "source_Pond", "source_Tube well"]
X = data[features]
y = data["unsafe"]

# Four ways of handling the missing values, each judged by stratified
# 5-fold cross validation. The filling happens INSIDE each pipeline, so it
# is learned from the training folds only and never sees the test fold.
folds = StratifiedKFold(n_splits=5, shuffle=True, random_state=2)
forest = RandomForestClassifier(n_estimators=300, min_samples_leaf=2, random_state=0)
complete = data.dropna(subset=features)
strategies = [
    ["A. Drop rows with a gap", None],
    ["B. Fill with the mean", make_pipeline(SimpleImputer(strategy="mean"), forest)],
    ["C. Fill with the median", make_pipeline(SimpleImputer(strategy="median"), forest)],
    ["D. Median + a 'was missing' column", make_pipeline(SimpleImputer(strategy="median", add_indicator=True), forest)],
]
print("Rows left if every row with a gap is dropped :", len(complete), "of", len(data))
recalls = []
for name, model in strategies:
    if model is None:
        r = cross_val_score(forest, complete[features], complete["unsafe"], cv=folds, scoring="recall")
        f = cross_val_score(forest, complete[features], complete["unsafe"], cv=folds, scoring="f1")
    else:
        r = cross_val_score(model, X, y, cv=folds, scoring="recall")
        f = cross_val_score(model, X, y, cv=folds, scoring="f1")
    recalls.append([name, r.mean(), f.mean()])
    print(name.ljust(36), "CV recall", round(r.mean(), 3), "  CV F1", round(f.mean(), 3))
print()

# Chart 4: the four strategies compared.
x = np.arange(4)
plt.figure(figsize=(8, 4.3))
plt.bar(x - 0.2, [r[1] for r in recalls], width=0.4, label="Recall (unsafe found)", color="#e15759")
plt.bar(x + 0.2, [r[2] for r in recalls], width=0.4, label="F1 score", color="#4e79a7")
plt.xticks(x, ["A. Drop", "B. Mean", "C. Median", "D. Median +\nindicator"])
plt.ylim(0.5, 1)
plt.ylabel("Mean over 5 folds")
plt.title("Four ways of handling missing values")
plt.legend()
plt.tight_layout()
plt.savefig("chart4_strategies.png")
plt.close()

# %% Step 7 - Modelling
print("--- Step 7: Modelling ---")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=5)
# Strategy A scored best, but only on the 797 complete rows, and a deployed
# model cannot refuse every sample with a missing test. B, C and D tied, and
# the indicator columns added nothing, so the simplest - C, the median - is used.
model = make_pipeline(SimpleImputer(strategy="median"),
                      RandomForestClassifier(n_estimators=300, min_samples_leaf=2, random_state=0))
model.fit(X_train, y_train)
print("Strategy used          : C (fill with the median)")
print("Training samples       :", len(X_train), " Testing samples :", len(X_test))
print()

# %% Step 8 - Evaluation
print("--- Step 8: Evaluation ---")
pred = model.predict(X_test)
tn, fp, fn, tp = confusion_matrix(y_test, pred).ravel()
print("TP, FN, FP, TN :", tp, fn, fp, tn)
print("Accuracy       :", round(accuracy_score(y_test, pred), 3))
print("Precision      :", round(precision_score(y_test, pred), 3))
print("Recall         :", round(recall_score(y_test, pred), 3))
print("F1 score       :", round(f1_score(y_test, pred), 3))
print()
cv_recall = cross_val_score(model, X, y, cv=folds, scoring="recall")
print("Recall in each of the 5 folds :", np.round(cv_recall, 3))
print()

# A missed unsafe source is far worse than an extra lab test, so the
# threshold is lowered from 0.5 to 0.3.
threshold = 0.3
prob_test = model.predict_proba(X_test)[:, 1]
flag = (prob_test >= threshold).astype(int)
print("Threshold for 'unsafe'        :", threshold)
print("Recall at threshold 0.3       :", round(recall_score(y_test, flag), 3))
print("Precision at threshold 0.3    :", round(precision_score(y_test, flag), 3))
print("Unsafe sources missed at 0.3  :", int(((flag == 0) & (y_test == 1)).sum()), "of", int(y_test.sum()))
print()

# Chart 5: confusion matrix.
cm = confusion_matrix(y_test, pred, labels=[1, 0])
plt.figure(figsize=(5.5, 4.5))
plt.imshow(cm, cmap="Blues")
labels = [["TP", "FN"], ["FP", "TN"]]
for r in range(2):
    for c in range(2):
        plt.text(c, r, labels[r][c] + "\n" + str(cm[r][c]), ha="center", va="center", fontsize=13,
                 color="white" if cm[r][c] > cm.max() / 2 else "black")
plt.xticks([0, 1], ["Predicted unsafe", "Predicted safe"])
plt.yticks([0, 1], ["Unsafe", "Safe"])
plt.title("Confusion matrix (test samples)")
plt.tight_layout()
plt.savefig("chart5_confusion.png")
plt.close()

# Chart 6: which readings the forest relied on.
importance = pd.Series(model[-1].feature_importances_, index=features).sort_values()
print("Feature importance:")
print(importance.round(3).to_string())
print()
plt.figure(figsize=(7.5, 5.2))
plt.barh(importance.index, importance.values, color="#59a14f")
plt.title("What the Random Forest relied on")
plt.xlabel("Importance")
plt.tight_layout()
plt.savefig("chart6_importance.png")
plt.close()

# %% Step 9 - Deployment
print("--- Step 9: Deployment ---")
joblib.dump(model, "water_model.pkl")
model = joblib.load("water_model.pkl")
new = pd.DataFrame({
    "ph": [7.2, 7.8, 6.9], "tds_mg_l": [380, 1250, 520], "turbidity_ntu": [0.4, 2.1, 9.5],
    "nitrate_mg_l": [8.0, 30.0, np.nan], "fluoride_mg_l": [0.3, np.nan, 0.6], "chloride_mg_l": [110, 380, 150],
    "residual_chlorine_mg_l": [0.4, 0.0, 0.0], "source_Hand pump": [0, 1, 0], "source_Open well": [0, 0, 1],
    "source_Piped supply": [1, 0, 0], "source_Pond": [0, 0, 0], "source_Tube well": [0, 0, 0],
})
names = ["Piped supply, chlorinated", "Hand pump, fluoride not tested", "Open well, cloudy, nitrate not tested"]
prob = model.predict_proba(new)[:, 1]
for i in range(len(names)):
    if prob[i] >= threshold:
        verdict = "LIKELY UNSAFE - warn users, send to the lab first"
    else:
        verdict = "likely safe - confirm with the lab test"
    print(names[i].ljust(38), round(prob[i], 2), "->", verdict)
print()

# %% Step 10 - Feedback
print("--- Step 10: Feedback ---")
# After the monsoon, 250 more sources are tested. Rain washes dirt and
# sewage into wells and ponds, so contamination rises.
mon = pd.read_csv("monsoon_samples.csv")
mon["unsafe"] = (mon["verdict"] == "Unsafe").astype(int)
mon = pd.get_dummies(mon, columns=["source"], dtype=int)
mon_pred = (model.predict_proba(mon[features])[:, 1] >= threshold).astype(int)
print("Monsoon samples                    :", len(mon))
print("Share unsafe, before the monsoon   :", round(y.mean() * 100, 1), "%")
print("Share unsafe, after the monsoon    :", round(mon["unsafe"].mean() * 100, 1), "%")
print("Recall on the test samples (0.3)   :", round(recall_score(y_test, flag), 3))
print("Recall on the monsoon samples      :", round(recall_score(mon["unsafe"], mon_pred), 3))
print("Precision on the monsoon samples   :", round(precision_score(mon["unsafe"], mon_pred), 3))
if recall_score(mon["unsafe"], mon_pred) < recall_score(y_test, flag) - 0.05:
    print("Feedback: recall has fallen after the monsoon - add monsoon samples and retrain.")
else:
    print("Feedback: recall holds after the monsoon - the model can stay in use.")

================================================================
IS THIS WATER SAFE TO DRINK?
================================================================
Problem : the lab's verdict takes two days; people keep drinking meanwhile
Goal    : predict the lab verdict from field-kit readings, the same day

Analytic approach : predictive - classification (Safe or Unsafe)

Samples collected : 1203
  sample_id        source    ph  tds_mg_l  turbidity_ntu  nitrate_mg_l  fluoride_mg_l  chloride_mg_l  residual_chlorine_mg_l verdict
0    W70001     Open well  7.33     298.0            3.7          38.6           0.62           43.0                    0.00    Safe
1    W70002     Tube well  8.07     409.0            0.8           5.3           0.88           99.0                    0.00    Safe
2    W70003  Piped supply  7.23     500.0            0.3           9.5           0.96          140.0                    0.48    Safe
3    W70004     Tube well  7.78     638.0            0.6           NaN            NaN          202.0                    0.00    Safe
4    W70005     Hand pump  8.08     390.0            4.9          33.2            NaN          119.0                    0.00  Unsafe
5    W70006     Open well  7.95     319.0            1.7          34.1           0.41           93.0                    0.00    Safe

--- Step 5: Data Understanding ---
Missing values in each column:
sample_id                   0
source                      0
ph                          0
tds_mg_l                    0
turbidity_ntu              60
nitrate_mg_l              120
fluoride_mg_l             265
chloride_mg_l               0
residual_chlorine_mg_l      0
verdict                     0

Lab verdicts:
verdict
Safe      727
Unsafe    476
Share unsafe : 39.6 %

Share unsafe by kind of source (%):
source
Piped supply    10.8
Tube well       32.6
Hand pump       45.3
Open well       80.2
Pond            81.7

--- Step 6: Data Preparation ---
Duplicate samples removed : 3
Rows left if every row with a gap is dropped : 797 of 1200
A. Drop rows with a gap              CV recall 0.86   CV F1 0.91
B. Fill with the mean                CV recall 0.821   CV F1 0.877
C. Fill with the median              CV recall 0.819   CV F1 0.879
D. Median + a 'was missing' column   CV recall 0.821   CV F1 0.877

--- Step 7: Modelling ---
Strategy used          : C (fill with the median)
Training samples       : 960  Testing samples : 240

--- Step 8: Evaluation ---
TP, FN, FP, TN : 76 19 2 143
Accuracy       : 0.912
Precision      : 0.974
Recall         : 0.8
F1 score       : 0.879

Recall in each of the 5 folds : [0.811 0.842 0.832 0.811 0.8  ]

Threshold for 'unsafe'        : 0.3
Recall at threshold 0.3       : 0.874
Precision at threshold 0.3    : 0.847
Unsafe sources missed at 0.3  : 12 of 95

Feature importance:
source_Hand pump          0.007
source_Tube well          0.009
source_Pond               0.015
source_Open well          0.034
residual_chlorine_mg_l    0.044
chloride_mg_l             0.052
tds_mg_l                  0.054
source_Piped supply       0.065
fluoride_mg_l             0.148
ph                        0.167
nitrate_mg_l              0.183
turbidity_ntu             0.222

--- Step 9: Deployment ---
Piped supply, chlorinated              0.03 -> likely safe - confirm with the lab test
Hand pump, fluoride not tested         0.22 -> likely safe - confirm with the lab test
Open well, cloudy, nitrate not tested  0.98 -> LIKELY UNSAFE - warn users, send to the lab first

--- Step 10: Feedback ---
Monsoon samples                    : 250
Share unsafe, before the monsoon   : 39.6 %
Share unsafe, after the monsoon    : 52.4 %
Recall on the test samples (0.3)   : 0.874
Recall on the monsoon samples      : 0.931
Precision on the monsoon samples   : 0.961
Feedback: recall holds after the monsoon - the model can stay in use.