Who Is at Risk of Heart Disease?

# ---------------------------------------------------------------------------
# heart_risk.py
#
# Capstone project: screening patients for heart disease from routine
# check-up measurements, so that those at risk are referred to a
# cardiologist early.
#
# CBSE Class XII, Artificial Intelligence (843), Part C - Capstone Project.
# SDG 3 (Good Health and Well-being).
#
# Data Science Methodology, 10 steps. A CLASSIFICATION problem: two models
# - Logistic Regression and a Decision Tree - compared with a train-test
# split and stratified 5-fold cross validation, using the confusion matrix,
# accuracy, precision, recall and F1. The decision threshold is then tuned,
# because in screening a missed patient costs far more than an extra test.
#
# Libraries: pandas, numpy, matplotlib, scikit-learn, joblib.
# ---------------------------------------------------------------------------

# %% Step 1 - Business Understanding
print("=" * 64)
print("WHO IS AT RISK OF HEART DISEASE?")
print("=" * 64)
print("Problem : heart disease is often found late; cardiologists are few")
print("Goal    : flag the patients at a check-up who should see a specialist")
print()

# %% Step 2 - Analytic Approach
print("Analytic approach : predictive - classification (disease: yes or no)")
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_validate
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
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("patients.csv")
print("Rows collected :", len(data))
print(data.head(6))
print()

# %% Step 5 - Data Understanding
print("--- Step 5: Data Understanding ---")
print("Missing values:")
print(data.isnull().sum().to_string())
print("Cholesterol recorded as 0 :", (data["cholesterol"] == 0).sum())
print()
print("Patients with and without heart disease:")
print(data["heart_disease"].value_counts().to_string())
print("Share with heart disease  :", round(data["heart_disease"].mean() * 100, 1), "%")
print()

sick = data[data["heart_disease"] == 1]
well = data[data["heart_disease"] == 0]

# Chart 1: age, with and without disease.
plt.figure(figsize=(7, 4.3))
plt.hist(well["age"], bins=20, alpha=0.6, color="#4e79a7", label="No heart disease")
plt.hist(sick["age"], bins=20, alpha=0.6, color="#e15759", label="Heart disease")
plt.title("Age of the patients")
plt.xlabel("Age (years)")
plt.ylabel("Patients")
plt.legend()
plt.tight_layout()
plt.savefig("chart1_age.png")
plt.close()

# Chart 2: maximum heart rate, as box plots.
plt.figure(figsize=(6, 4.3))
plt.boxplot([well["max_heart_rate"], sick["max_heart_rate"]], labels=["No heart disease", "Heart disease"])
plt.title("Maximum heart rate reached in an exercise test")
plt.ylabel("Beats per minute")
plt.grid(True, axis="y", linestyle="--", alpha=0.4)
plt.tight_layout()
plt.savefig("chart2_max_heart_rate.png")
plt.close()

# Chart 3: share with disease for each kind of chest pain, by sex.
share = data.groupby(["chest_pain", "sex"])["heart_disease"].mean().unstack() * 100
share = share.reindex(["typical", "atypical", "non-anginal", "none"])
print("Share with heart disease, by chest pain and sex (%):")
print(share.round(1))
print()
x = np.arange(4)
plt.figure(figsize=(8, 4.3))
plt.bar(x - 0.2, share["M"], width=0.4, label="Men", color="#4e79a7")
plt.bar(x + 0.2, share["F"], width=0.4, label="Women", color="#b07aa1")
plt.xticks(x, share.index)
plt.ylabel("% with heart disease")
plt.title("Heart disease by type of chest pain")
plt.legend()
plt.tight_layout()
plt.savefig("chart3_chest_pain.png")
plt.close()

# %% Step 6 - Data Preparation
print("--- Step 6: Data Preparation ---")
before = len(data)
data = data.drop_duplicates()
print("Duplicate rows removed               :", before - len(data))

# Cholesterol of 0 means "not measured", as in the UCI data it is modelled
# on. It is turned into a real missing value and then filled with the
# median - and a new column records that it was missing, in case being
# untested is itself informative.
data["cholesterol_missing"] = (data["cholesterol"] == 0).astype(int)
data["cholesterol"] = data["cholesterol"].replace(0, np.nan)
data["cholesterol"] = data["cholesterol"].fillna(data["cholesterol"].median())
data["resting_bp"] = pd.to_numeric(data["resting_bp"])
blank_bp = data["resting_bp"].isnull().sum()
data["resting_bp"] = data["resting_bp"].fillna(data["resting_bp"].median())
print("Cholesterol 0s replaced by the median:", data["cholesterol_missing"].sum())
print("Blank blood pressures filled         :", blank_bp)

data["male"] = (data["sex"] == "M").astype(int)
data["exercise_angina"] = (data["exercise_angina"] == "Yes").astype(int)
data = pd.get_dummies(data, columns=["chest_pain"], dtype=int)
features = ["age", "male", "resting_bp", "cholesterol", "cholesterol_missing", "fasting_sugar_high",
            "max_heart_rate", "exercise_angina", "st_depression", "chest_pain_atypical",
            "chest_pain_non-anginal", "chest_pain_none", "chest_pain_typical"]
print("Features used                        :", len(features))
print()

# Chart 4: correlations between the number features, as a heat map.
cols = ["age", "resting_bp", "cholesterol", "max_heart_rate", "st_depression", "exercise_angina", "heart_disease"]
corr = data[cols].corr()
plt.figure(figsize=(7.5, 6))
plt.imshow(corr, cmap="RdBu_r", vmin=-1, vmax=1)
plt.colorbar(label="Correlation")
plt.xticks(range(len(cols)), cols, rotation=45, ha="right")
plt.yticks(range(len(cols)), cols)
for i in range(len(cols)):
    for j in range(len(cols)):
        plt.text(j, i, round(corr.iloc[i, j], 2), ha="center", va="center", fontsize=8)
plt.title("Correlation between the measurements")
plt.tight_layout()
plt.savefig("chart4_heatmap.png")
plt.close()
print("Correlation with heart disease:")
print(corr["heart_disease"].drop("heart_disease").round(2).sort_values().to_string())
print()

# %% Step 7 - Modelling
print("--- Step 7: Modelling ---")
X = data[features]
y = data["heart_disease"]
# stratify=y keeps the same share of disease in the training and test sets.
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=11)
print("Training rows :", len(X_train), " Testing rows :", len(X_test))

# Logistic Regression needs features on similar scales, so a scaler is
# placed in front of it in a pipeline: the scaler learns only from the
# training data, and the test data is scaled with what it learned.
logistic = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))
tree = DecisionTreeClassifier(max_depth=4, min_samples_leaf=10, random_state=0)
logistic.fit(X_train, y_train)
tree.fit(X_train, y_train)

coef = pd.Series(logistic[-1].coef_[0], index=features).sort_values()
print("Logistic Regression coefficients (scaled features):")
print(coef.round(2).to_string())
print()

# Chart 5: what Logistic Regression learned.
plt.figure(figsize=(7.5, 5))
plt.barh(coef.index, coef.values, color=["#59a14f" if v < 0 else "#e15759" for v in coef.values])
plt.axvline(0, color="#999999")
plt.title("Logistic Regression: red raises the risk, green lowers it")
plt.xlabel("Coefficient (on scaled features)")
plt.tight_layout()
plt.savefig("chart5_coefficients.png")
plt.close()

# %% Step 8 - Evaluation
print("--- Step 8: Evaluation ---")
for name, model in [["Logistic Regression", logistic], ["Decision Tree", tree]]:
    pred = model.predict(X_test)
    tn, fp, fn, tp = confusion_matrix(y_test, pred).ravel()
    print(name)
    print("  " + name + " TP, FN, FP, TN :", tp, fn, fp, tn)
    print("  " + name + " accuracy  :", round(accuracy_score(y_test, pred), 3))
    print("  " + name + " precision :", round(precision_score(y_test, pred), 3))
    print("  " + name + " recall    :", round(recall_score(y_test, pred), 3))
    print("  " + name + " F1        :", round(f1_score(y_test, pred), 3))
print()

# Stratified 5-fold cross validation: every fold keeps the same share of
# patients with disease as the whole dataset.
folds = StratifiedKFold(n_splits=5, shuffle=True, random_state=1)
cv = {}
for name, model in [["Logistic Regression", make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))],
                    ["Decision Tree", DecisionTreeClassifier(max_depth=4, min_samples_leaf=10, random_state=0)]]:
    scores = cross_validate(model, X, y, cv=folds, scoring=["accuracy", "recall"])
    cv[name] = scores
    print(name, "- recall in each fold   :", np.round(scores["test_recall"], 3))
    print(name, "- mean CV accuracy      :", round(scores["test_accuracy"].mean(), 3))
    print(name, "- mean CV recall        :", round(scores["test_recall"].mean(), 3))
print()

# Chart 6: both confusion matrices.
plt.figure(figsize=(10, 4.3))
for k, (name, model) in enumerate([["Logistic Regression", logistic], ["Decision Tree", tree]]):
    m = confusion_matrix(y_test, model.predict(X_test), labels=[1, 0])
    plt.subplot(1, 2, k + 1)
    plt.imshow(m, cmap="Reds")
    labels = [["TP", "FN"], ["FP", "TN"]]
    for r in range(2):
        for c in range(2):
            plt.text(c, r, labels[r][c] + "\n" + str(m[r][c]), ha="center", va="center", fontsize=12,
                     color="white" if m[r][c] > m.max() / 2 else "black")
    plt.xticks([0, 1], ["Predicted disease", "Predicted none"])
    plt.yticks([0, 1], ["Disease", "No disease"])
    plt.title(name)
plt.tight_layout()
plt.savefig("chart6_confusion.png")
plt.close()

# Chart 7: recall in each fold.
x = np.arange(5)
plt.figure(figsize=(7.5, 4))
plt.bar(x - 0.2, cv["Logistic Regression"]["test_recall"], width=0.4, label="Logistic Regression", color="#e15759")
plt.bar(x + 0.2, cv["Decision Tree"]["test_recall"], width=0.4, label="Decision Tree", color="#9c9c9c")
plt.xticks(x, ["Fold 1", "Fold 2", "Fold 3", "Fold 4", "Fold 5"])
plt.ylim(0, 1)
plt.ylabel("Recall")
plt.title("Stratified 5-fold cross validation: recall")
plt.legend()
plt.tight_layout()
plt.savefig("chart7_kfold_recall.png")
plt.close()

# The threshold. Logistic Regression gives each patient a probability;
# by default a patient is flagged at 0.5. A lower threshold catches more
# patients with disease (recall up) at the cost of more referrals that
# turn out healthy (precision down).
prob = logistic.predict_proba(X_test)[:, 1]
thresholds = np.round(np.arange(0.2, 0.81, 0.05), 2)
prec_list = []
rec_list = []
for t in thresholds:
    flag = (prob >= t).astype(int)
    prec_list.append(precision_score(y_test, flag, zero_division=0))
    rec_list.append(recall_score(y_test, flag))
chosen = 0.35
flag = (prob >= chosen).astype(int)
print("Threshold chosen for screening   :", chosen)
print("Recall at threshold 0.35         :", round(recall_score(y_test, flag), 3))
print("Precision at threshold 0.35      :", round(precision_score(y_test, flag), 3))
print("Patients referred at 0.35        :", flag.sum(), "of", len(flag))
print()

# Chart 8: precision and recall at every threshold.
plt.figure(figsize=(7.5, 4.3))
plt.plot(thresholds, rec_list, marker="o", color="#e15759", label="Recall")
plt.plot(thresholds, prec_list, marker="s", color="#4e79a7", label="Precision")
plt.axvline(chosen, color="#999999", linestyle="--", label="Chosen threshold")
plt.title("Choosing the decision threshold")
plt.xlabel("Probability needed to refer a patient")
plt.ylim(0, 1.05)
plt.legend()
plt.grid(True, linestyle="--", alpha=0.4)
plt.tight_layout()
plt.savefig("chart8_threshold.png")
plt.close()

# %% Step 9 - Deployment
print("--- Step 9: Deployment ---")
joblib.dump(logistic, "heart_model.pkl")
model = joblib.load("heart_model.pkl")
new_patients = pd.DataFrame({
    "age": [62, 45, 51], "male": [1, 0, 1], "resting_bp": [150, 120, 135], "cholesterol": [290, 210, 245],
    "cholesterol_missing": [0, 0, 0], "fasting_sugar_high": [1, 0, 0], "max_heart_rate": [118, 172, 150],
    "exercise_angina": [1, 0, 0], "st_depression": [2.4, 0.2, 1.0], "chest_pain_atypical": [0, 1, 0],
    "chest_pain_non-anginal": [0, 0, 1], "chest_pain_none": [0, 0, 0], "chest_pain_typical": [1, 0, 0],
})
p_new = model.predict_proba(new_patients)[:, 1]
for i in range(len(new_patients)):
    if p_new[i] >= chosen:
        action = "REFER to a cardiologist"
    else:
        action = "routine follow-up"
    print("Patient", i + 1, "- probability", round(p_new[i], 3), "->", action)
print()

# %% Step 10 - Feedback
print("--- Step 10: Feedback ---")
# Months later, the diagnoses of 200 more patients are confirmed. Checking
# the model against them - for men and women separately - shows whether it
# works equally well for everyone.
fu = pd.read_csv("follow_up.csv")
fu["cholesterol_missing"] = 0
fu["male"] = (fu["sex"] == "M").astype(int)
fu["exercise_angina"] = (fu["exercise_angina"] == "Yes").astype(int)
fu = pd.get_dummies(fu, columns=["chest_pain"], dtype=int)
fu_flag = (model.predict_proba(fu[features])[:, 1] >= chosen).astype(int)
print("Follow-up patients                :", len(fu))
print("Recall on follow-up, everyone     :", round(recall_score(fu["heart_disease"], fu_flag), 3))
men = fu["male"] == 1
print("Recall on follow-up, men          :", round(recall_score(fu["heart_disease"][men], fu_flag[men]), 3))
print("Recall on follow-up, women        :", round(recall_score(fu["heart_disease"][~men], fu_flag[~men]), 3))
gap = recall_score(fu["heart_disease"][men], fu_flag[men]) - recall_score(fu["heart_disease"][~men], fu_flag[~men])
print("Recall gap, men minus women       :", round(gap * 100, 1), "points")
if gap > 0.05:
    print("Feedback: the model misses more women than men - collect more records")
    print("          from women and review the features before wider use.")
else:
    print("Feedback: recall is similar for men and women.")

================================================================
WHO IS AT RISK OF HEART DISEASE?
================================================================
Problem : heart disease is often found late; cardiologists are few
Goal    : flag the patients at a check-up who should see a specialist

Analytic approach : predictive - classification (disease: yes or no)

Rows collected : 1003
  patient_id  age sex   chest_pain  resting_bp  cholesterol  fasting_sugar_high  max_heart_rate exercise_angina  st_depression  \
0     H40001   49   F         none       126.0        230.0                   0           164.0              No            0.8   
1     H40002   39   F  non-anginal       144.0        250.0                   0           145.0              No            1.3   
2     H40003   71   M  non-anginal       117.0        217.0                   0           174.0              No            0.0   
3     H40004   49   M      typical       147.0        273.0                   1           131.0             Yes            2.1   
4     H40005   52   M      typical       163.0        278.0                   0           124.0             Yes            1.7   
5     H40006   68   M     atypical       127.0        260.0                   0           152.0              No            1.2   

   heart_disease  
0              0  
1              0  
2              0  
3              1  
4              1  
5              0  

--- Step 5: Data Understanding ---
Missing values:
patient_id            0
age                   0
sex                   0
chest_pain            0
resting_bp            6
cholesterol           0
fasting_sugar_high    0
max_heart_rate        0
exercise_angina       0
st_depression         0
heart_disease         0
Cholesterol recorded as 0 : 37

Patients with and without heart disease:
heart_disease
0    547
1    456
Share with heart disease  : 45.5 %

Share with heart disease, by chest pain and sex (%):
sex             F     M
chest_pain             
typical      73.2  83.1
atypical     28.3  26.8
non-anginal  21.5  17.9
none         38.0  30.8

--- Step 6: Data Preparation ---
Duplicate rows removed               : 3
Cholesterol 0s replaced by the median: 37
Blank blood pressures filled         : 6
Features used                        : 13

Correlation with heart disease:
max_heart_rate    -0.53
cholesterol        0.14
resting_bp         0.21
age                0.27
exercise_angina    0.47
st_depression      0.48

--- Step 7: Modelling ---
Training rows : 800  Testing rows : 200
Logistic Regression coefficients (scaled features):
max_heart_rate           -1.44
chest_pain_non-anginal   -0.45
chest_pain_none          -0.25
chest_pain_atypical      -0.22
male                     -0.16
cholesterol_missing       0.00
age                       0.02
fasting_sugar_high        0.24
cholesterol               0.25
resting_bp                0.48
chest_pain_typical        0.80
exercise_angina           0.98
st_depression             1.05

--- Step 8: Evaluation ---
Logistic Regression
  Logistic Regression TP, FN, FP, TN : 80 11 17 92
  Logistic Regression accuracy  : 0.86
  Logistic Regression precision : 0.825
  Logistic Regression recall    : 0.879
  Logistic Regression F1        : 0.851
Decision Tree
  Decision Tree TP, FN, FP, TN : 76 15 16 93
  Decision Tree accuracy  : 0.845
  Decision Tree precision : 0.826
  Decision Tree recall    : 0.835
  Decision Tree F1        : 0.831

Logistic Regression - recall in each fold   : [0.889 0.846 0.846 0.901 0.791]
Logistic Regression - mean CV accuracy      : 0.866
Logistic Regression - mean CV recall        : 0.855
Decision Tree - recall in each fold   : [0.767 0.813 0.736 0.868 0.736]
Decision Tree - mean CV accuracy      : 0.829
Decision Tree - mean CV recall        : 0.784

Threshold chosen for screening   : 0.35
Recall at threshold 0.35         : 0.923
Precision at threshold 0.35      : 0.808
Patients referred at 0.35        : 104 of 200

--- Step 9: Deployment ---
Patient 1 - probability 0.999 -> REFER to a cardiologist
Patient 2 - probability 0.012 -> routine follow-up
Patient 3 - probability 0.084 -> routine follow-up

--- Step 10: Feedback ---
Follow-up patients                : 200
Recall on follow-up, everyone     : 0.874
Recall on follow-up, men          : 0.889
Recall on follow-up, women        : 0.833
Recall gap, men minus women       : 5.6 points
Feedback: the model misses more women than men - collect more records
          from women and review the features before wider use.