Heart Disease Prediction Pipeline

Build an end-to-end ML pipeline to predict heart disease with scikit-learn.

Introduction

30 Second Summary

Heart disease is the leading cause of death worldwide. Doctors assess risk using patient data like blood pressure, cholesterol, and heart rate, but a machine learning model can help flag high-risk patients automatically.

In this project, you will build an end-to-end machine learning pipeline that predicts heart disease from 13 clinical features. You'll train three different models, tune hyperparameters with cross-validation, and evaluate performance using the exact metrics ML Engineer interviewers test you on.

What You'll Build

A complete classification system using scikit-learn that handles preprocessing, model training, hyperparameter tuning, and evaluation on the UCI Heart Disease dataset.

By the end of this project, you'll have:

  • A leak-free preprocessing pipeline using ColumnTransformer and Pipeline that handles mixed numeric and categorical features without data leakage.
  • A model comparison report showing cross-validated F1 scores for Logistic Regression, Random Forest, and XGBoost.
  • Evaluation visualizations including a confusion matrix heatmap, ROC curve with AUC score, and feature importance chart showing which clinical factors predict heart disease.
  • Secret Mission: Save your best model and write a prediction function that takes raw patient data and returns a diagnosis with confidence score.

Are there any prerequisites?

You'll need basic Python knowledge and familiarity with pandas DataFrames. No prior machine learning experience is required.

Before We Start

Before We Start...

Set Up Your Python Environment

Before you can build a machine learning pipeline, you need the right tools in place. Python is the language powering your pipeline, and scikit-learn is the library that makes building, training, and evaluating models straightforward.

In this step, you will install Python, set up VS Code with a Jupyter notebook environment, and install the ML packages your project depends on.

In this step, get ready to:
  • Install Python 3.11+ and the Jupyter extension in VS Code.
  • Create a project folder, virtual environment, and install packages.
  • Create a Jupyter notebook and verify all imports work.
Install Python 3.11+ and the Jupyter extension

This project requires Python 3.11 or higher because scikit-learn 1.8.0 depends on it. Let's confirm you have the right version installed.

  • Check your Python version by running this command in a terminal:

macOS

python3 --version

Windows

python --version

✔️ I see version 3.11 or higher

Your Python version is compatible. Move on to opening VS Code below.

ⓧ I see an older version

You have Python installed, but the version is too old for this project. Download and install Python 3.11 or newer from python.org/downloads.

On macOS, your system may still default to an older version even after installing a newer one. After installing, try running python3.11 --version or python3.14 --version to confirm the new binary is accessible.

On Windows, the installer adds Python to your PATH automatically if you check the box during setup. Re-open your terminal after installing, then run python --version again.

ⓧ Command not found

Python is not installed yet. Download the latest installer from python.org/downloads.

  • On macOS, download the macOS installer and follow the prompts.
  • On Windows, download the Windows installer (64-bit). Check Add python.exe to PATH during installation.
  • After installing, close and reopen your terminal, then run the version check again.
  • Open VS Code by pressing Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Visual Studio Code and press Enter to open it.

Don't have VS Code?

Download it for free from code.visualstudio.com. Install the version for your operating system and follow the setup wizard.

  • Open the Extensions panel by pressing Cmd+Shift+X (macOS) or Ctrl+Shift+X (Windows).
  • Type Jupyter in the search box.
  • Click Install on the Jupyter extension by Microsoft.

What is the Jupyter extension?

The Jupyter extension lets you run interactive notebooks directly inside VS Code. Notebooks combine code, output, and explanations in one document, which is the standard format for ML projects.

Create your project folder and install packages

A virtual environment keeps your project's packages isolated from other Python projects on your machine. This prevents version conflicts between projects.

  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir heart-disease-pipeline
cd heart-disease-pipeline
  • Create a virtual environment by running:

macOS

python3 -m venv venv

Windows

python -m venv venv
  • Activate the virtual environment:

macOS

source venv/bin/activate

Windows

venv\Scripts\activate

You should see (venv) appear at the beginning of your terminal prompt. This confirms the virtual environment is active.

  • Install all the packages your pipeline needs by running:
pip install scikit-learn==1.8.0 xgboost==3.2.0 pandas matplotlib seaborn ucimlrepo

What are these packages?

  • scikit-learn provides the Pipeline, preprocessing, and model evaluation tools.
  • xgboost is a high-performance gradient boosting library used heavily in industry.
  • pandas handles your data as DataFrames for easy manipulation.
  • matplotlib and seaborn produce charts and visualizations.
  • ucimlrepo fetches the Heart Disease dataset directly from the UCI repository.

The install may take a minute or two. You should see a success message for each package once it finishes.

Seeing install errors?

  • Make sure your virtual environment is active. You should see (venv) in your prompt.
  • If you see a permissions error, do not use sudo. Deactivate and recreate the virtual environment instead.
  • If pip is not found, try pip3 instead.

help me fix my pip install error for scikit-learn or xgboost

Create your Jupyter notebook and verify imports
  • In VS Code, open your project folder by clicking File in the top menu bar.
  • Click Open Folder.
  • Select the heart-disease-pipeline folder on your Desktop, then click Open.
  • Open the Command Palette by pressing Cmd+Shift+P (macOS) or Ctrl+Shift+P (Windows).
  • Type Create: New Jupyter Notebook and select it from the dropdown.
  • Save the new notebook immediately by pressing Cmd+S (macOS) or Ctrl+S (Windows).
  • Name the file heart_disease_pipeline.ipynb and save it in the heart-disease-pipeline folder.
  • Select your virtual environment as the notebook kernel by clicking Select Kernel in the top right corner of the notebook.
  • Choose Python Environments, then select the venv environment you created.

What is a kernel?

A kernel is the engine that runs your notebook code. By selecting your virtual environment as the kernel, the notebook has access to all the packages you just installed.

Now add a title and the imports that your pipeline will use throughout the project.

  • The notebook starts with an empty code cell. Change it to a Markdown cell by clicking the cell type dropdown (it says Python or Code) and selecting Markdown.
  • Type this content into the Markdown cell:
# Heart Disease Prediction Pipeline
End-to-end ML classification using scikit-learn Pipeline, model comparison, and evaluation.
  • Run the cell by pressing Shift+Enter. This renders the Markdown and creates a new cell below.
  • Change the new cell to Markdown as well, then type:
## 1. Setup and Imports
  • Run this cell with Shift+Enter to render it and create the next cell.
  • The new cell should already be a Code cell. Paste the following imports into it:
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from ucimlrepo import fetch_ucirepo

from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold, GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, roc_auc_score
from xgboost import XGBClassifier
import joblib

What does this code do?

  • pandas and numpy handle data loading and numerical operations.
  • matplotlib and seaborn create visualizations like heatmaps and bar charts.
  • The sklearn imports cover the full ML workflow: splitting data, building pipelines, preprocessing, training models, and evaluating results.
  • XGBClassifier is an industry-standard gradient boosting model for tabular data.
  • joblib saves trained models to disk so you can reload them later.
  • Run the imports cell by pressing Shift+Enter.

If the cell runs without errors and a new empty cell appears below, all your packages are installed correctly.

Seeing a ModuleNotFoundError?

  • This means the notebook kernel is not using your virtual environment. Click Select Kernel in the top right and choose the venv Python interpreter.
  • If your venv does not appear in the kernel list, close and reopen VS Code with the heart-disease-pipeline folder open. VS Code needs to detect the environment.

help me fix ModuleNotFoundError in my Jupyter notebook

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook.

ⓧ I'd like to double check the full code

Your notebook should have two Markdown cells and one code cell. Here is the full content so far:

# %% [markdown]
# # Heart Disease Prediction Pipeline
# End-to-end ML classification using scikit-learn Pipeline, model comparison, and evaluation.

# %% [markdown]
# ## 1. Setup and Imports

# %%
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from ucimlrepo import fetch_ucirepo

from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold, GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, roc_auc_score
from xgboost import XGBClassifier
import joblib

Your environment is ready and every package is installed. Next up, you will load the Heart Disease dataset and explore its features with visualizations.

Load and Explore the Dataset

Before building any model, you need to understand the data you're working with. A dataset's shape, feature types, and quality issues determine every design decision in your preprocessing pipeline.

In this step, you'll load the UCI Heart Disease dataset, visualize how features correlate with each other, and identify which features are numeric vs. categorical. This analysis directly informs the preprocessing pipeline you'll build in the next step.

In this step, get ready to:
  • Load the UCI Heart Disease dataset and inspect its structure.
  • Visualize feature correlations and distributions by target class.
  • Identify numeric vs. categorical features to guide your preprocessing strategy.
Load the UCI Heart Disease dataset

The UCI Heart Disease dataset contains 13 clinical features from 303 patients. The target variable indicates whether the patient has heart disease (values 1-4) or not (value 0).

  • In your notebook, click below the imports cell to select it, then click the + Code button in the toolbar to add a new code cell.
  • Before the code cell, add a Markdown cell by clicking + Markdown and type:
## 2. Load and Explore the Dataset
  • Add a new code cell below the Markdown cell and paste this code to load and inspect the dataset:
# Load UCI Heart Disease dataset (Cleveland subset)
heart_disease = fetch_ucirepo(id=45)
X = heart_disease.data.features
y = heart_disease.data.targets

# Binary classification: 0 = no disease, 1-4 = disease present
y = (y.values.ravel() > 0).astype(int)

print(f"Dataset shape: {X.shape}")
print(f"Target distribution:\n{pd.Series(y).value_counts()}")
print(f"\nFeature types:\n{X.dtypes}")
print(f"\nMissing values:\n{X.isnull().sum()}")

What does this code do?

  • fetch_ucirepo(id=45) downloads the Cleveland Heart Disease dataset from the UCI Machine Learning Repository.
  • heart_disease.data.features gives you the 13 clinical features as a pandas DataFrame. heart_disease.data.targets gives you the diagnosis labels.
  • The original target has 5 values (0-4). You convert it to binary: 0 means no disease, 1 means disease present.
  • The print statements show the dataset shape, class balance, data types, and any missing values.
  • Run the cell by clicking the play button to its left (or press Shift+Enter).
  • You should see output showing the dataset has approximately 303 rows and 13 columns, with a roughly balanced target distribution.

What is the UCI Machine Learning Repository?

The UCI Machine Learning Repository is a collection of datasets maintained by the University of California, Irvine. It's one of the most-cited sources for benchmark datasets in ML research and practice.

Seeing a connection error?

  • fetch_ucirepo needs an internet connection to download the dataset. Check that you're connected to the internet.
  • If you see a ModuleNotFoundError, make sure your notebook kernel is set to the virtual environment where you installed ucimlrepo. Click the kernel selector in the top-right of the notebook to switch.

Help me fix my dataset loading error in Jupyter

Create visualizations

Visualizations help you spot patterns that inform model design. A correlation heatmap reveals which features move together, and distribution plots show how features differ between healthy and diseased patients.

  • Add a new code cell below and paste this code to create a correlation heatmap:
# Correlation heatmap for numeric features
numeric_cols = X.select_dtypes(include=[np.number]).columns.tolist()
plt.figure(figsize=(10, 8))
sns.heatmap(X[numeric_cols].corr(), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()

What does this code do?

  • select_dtypes(include=[np.number]) filters to only numeric columns, since correlations don't apply to categorical data.
  • .corr() computes the Pearson correlation coefficient between every pair of numeric features. Values near +1 or -1 mean strong linear relationships.
  • The heatmap gives you a visual summary of which features are correlated. Highly correlated features can sometimes cause issues for linear models.
  • Run the cell.
  • You should see a color-coded matrix with correlation values annotated in each square. Look for strong positive (red) or negative (blue) correlations.

Heatmap not displaying?

  • Make sure you ran the data loading cell above first. The variable X must exist before this cell can reference it.
  • If you see a blank output, try adding %matplotlib inline at the top of the cell and rerunning.

Help me get my matplotlib plots to display in VS Code Jupyter

  • Add another new code cell and paste this code to plot feature distributions by target class:
# Distribution of key features by target class
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
features_to_plot = ['age', 'thalach', 'chol']
titles = ['Age', 'Max Heart Rate', 'Cholesterol']

for ax, feat, title in zip(axes, features_to_plot, titles):
    for label in [0, 1]:
        ax.hist(X[feat][y == label].dropna(), alpha=0.6, label=f'{"No Disease" if label == 0 else "Disease"}', bins=20)
    ax.set_xlabel(title)
    ax.set_ylabel('Count')
    ax.legend()

plt.suptitle('Feature Distributions by Heart Disease Status')
plt.tight_layout()
plt.show()

What does this code do?

  • This creates three side-by-side histograms for age, max heart rate (thalach), and cholesterol (chol).
  • Each histogram is split by target class (disease vs. no disease) so you can visually compare the distributions.
  • Features where the two classes have very different distributions are likely strong predictors. For example, patients with heart disease often have lower max heart rates.
  • Run the cell.
  • You should see three histograms with overlapping orange and blue distributions. Notice how the disease and no-disease groups separate differently across each feature.

Seeing a KeyError?

  • A KeyError means the column name doesn't match. Confirm your dataset loaded correctly by running X.columns.tolist() in a new cell to see the exact column names.
  • Make sure you're using the exact column names from the dataset: age, thalach, chol.

Help me fix a KeyError in my distribution plots

Identify numeric vs. categorical features

Your visualizations and data inspection revealed critical information about this dataset. The features fall into two distinct types that need different preprocessing strategies.

What are numeric vs. categorical features?

Numeric features are continuous values like age, blood pressure, and cholesterol. These need scaling so that features with large ranges don't dominate the model.

Categorical features represent discrete categories, even if they look like numbers. For example, cp (chest pain type) has values 1-4 representing different types of chest pain. These need encoding, not scaling.

Based on the dataset documentation and your exploration, here's how the 13 features split:

  • Numeric features (5): age, trestbps (resting blood pressure), chol (cholesterol), thalach (max heart rate), oldpeak (ST depression).
  • Categorical features (8): sex, cp (chest pain type), fbs (fasting blood sugar), restecg, exang (exercise angina), slope, ca (major vessels), thal.

Why do missing values matter?

You likely noticed missing values in the ca and thal columns. These are categorical features with a few missing entries.

If you ignore missing values, most models will crash or silently drop rows. In the next step, you'll handle this by building an imputer into your pipeline that fills missing values based only on training data.

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook before moving on.

ⓧ I'd like to double check the full code

Here's what your full notebook should look like up to this point:

# %% [markdown]
# # Heart Disease Prediction Pipeline
# End-to-end ML classification using scikit-learn Pipeline, model comparison, and evaluation.

# %% [markdown]
# ## 1. Setup and Imports

# %%
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from ucimlrepo import fetch_ucirepo

from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold, GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, roc_auc_score
from xgboost import XGBClassifier
import joblib

# %% [markdown]
# ## 2. Load and Explore the Dataset

# %%
# Load UCI Heart Disease dataset (Cleveland subset)
heart_disease = fetch_ucirepo(id=45)
X = heart_disease.data.features
y = heart_disease.data.targets

# Binary classification: 0 = no disease, 1-4 = disease present
y = (y.values.ravel() > 0).astype(int)

print(f"Dataset shape: {X.shape}")
print(f"Target distribution:\n{pd.Series(y).value_counts()}")
print(f"\nFeature types:\n{X.dtypes}")
print(f"\nMissing values:\n{X.isnull().sum()}")

# %%
# Correlation heatmap for numeric features
numeric_cols = X.select_dtypes(include=[np.number]).columns.tolist()
plt.figure(figsize=(10, 8))
sns.heatmap(X[numeric_cols].corr(), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()

# %%
# Distribution of key features by target class
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
features_to_plot = ['age', 'thalach', 'chol']
titles = ['Age', 'Max Heart Rate', 'Cholesterol']

for ax, feat, title in zip(axes, features_to_plot, titles):
    for label in [0, 1]:
        ax.hist(X[feat][y == label].dropna(), alpha=0.6, label=f'{"No Disease" if label == 0 else "Disease"}', bins=20)
    ax.set_xlabel(title)
    ax.set_ylabel('Count')
    ax.legend()

plt.suptitle('Feature Distributions by Heart Disease Status')
plt.tight_layout()
plt.show()

You now understand the shape and quality of your data. Next up, you'll use these insights to build a leak-free preprocessing pipeline that handles numeric scaling and categorical encoding in a single, reusable object.

Build a Leak-Free Preprocessing Pipeline

Your exploratory analysis revealed that the dataset has two types of features: numeric values like age and heart rate, and categorical codes like chest pain type and thalassemia status.

Each type needs different preprocessing before a model can use it. Numeric features need scaling so no single feature dominates. Categorical features need encoding into numbers a model can interpret. The challenge is doing this without accidentally leaking test data into your training process.

In this step, get ready to:
  • Create a ColumnTransformer that handles numeric and categorical features differently.
  • Split data into training and test sets with class balance preserved.
  • Document why this pipeline approach prevents data leakage.
Define feature lists and create the ColumnTransformer

A ColumnTransformer lets you apply different transformations to different columns of your data in a single step. Numeric features get imputed and scaled. Categorical features get imputed and one-hot encoded.

  • Add a new markdown cell below your last visualization cell. Type the following section heading:
## 3. Build a Leak-Free Preprocessing Pipeline
  • Add a new code cell below the markdown heading. Build the preprocessing pipeline by pasting this code:
# Identify feature types
numeric_features = ['age', 'trestbps', 'chol', 'thalach', 'oldpeak']
categorical_features = ['sex', 'cp', 'fbs', 'restecg', 'exang', 'slope', 'ca', 'thal']

# Create preprocessing pipelines for each feature type
numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median')),
    ('scaler', StandardScaler())
])

categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('encoder', OneHotEncoder(handle_unknown='ignore', sparse_output=False))
])

# Combine into a ColumnTransformer
preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, numeric_features),
        ('cat', categorical_transformer, categorical_features)
    ]
)

What does this code do?

This cell creates a two-track preprocessing system:

  • numeric_transformer fills missing values with the median, then scales each feature to have mean 0 and standard deviation 1.
  • categorical_transformer fills missing values with the most frequent category, then converts each category into a binary column.
  • preprocessor routes each feature to the correct transformer based on the lists you defined at the top.

Why median for numeric and most_frequent for categorical?

Median imputation is robust to outliers. If one patient has an extreme cholesterol reading, it won't skew the fill value for everyone else.

For categorical features, "most frequent" fills gaps with the most common category. You can't take the median of a category like "chest pain type 3."

  • Run the cell by pressing Shift+Enter.

The cell runs without output. That is expected. The preprocessor is now defined and ready to use in a pipeline.

Seeing a NameError?

  • Make sure you ran the imports cell at the top of your notebook first. Pipeline, ColumnTransformer, SimpleImputer, and StandardScaler all come from that first cell.
  • Check that your kernel is set to the virtual environment you created in Step 1.

Why am I getting a NameError for Pipeline or ColumnTransformer?

Split data into training and test sets

Before training any model, you need to set aside data it has never seen. This held-out test set is your reality check at the end.

  • Add a new code cell below the preprocessor cell. Split your data by pasting this code:
# Train-test split with stratification to preserve class balance
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"Training set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")
print(f"Train class balance: {np.bincount(y_train)}")
print(f"Test class balance: {np.bincount(y_test)}")

What does this code do?

  • train_test_split reserves 20% of data for testing and gives 80% for training.
  • stratify=y ensures both splits have the same ratio of disease to no-disease cases. Without this, random chance could put most disease cases in one split.
  • random_state=42 makes the split reproducible. Anyone running this notebook gets the same split.
  • Run the cell by pressing Shift+Enter.

You should see output showing approximately 242 training samples and 61 test samples, with similar class balances in both splits.

Class balance looks very uneven?

  • Check that your y variable was converted to binary in the earlier data loading cell. The line y = (y.values.ravel() > 0).astype(int) should produce roughly balanced classes (around 138 vs 165).

Why does my train-test split have unbalanced classes?

Explain why this approach prevents data leakage

Data leakage happens when information from your test set influences your training process. In interviews, this is one of the most common follow-up questions after you describe a pipeline.

  • Add a new markdown cell below the train-test split cell. Type the following explanation:
**Why this prevents data leakage:** The `ColumnTransformer` is inside the `Pipeline`.
During cross-validation, the scaler computes mean/std ONLY from the training fold,
and the imputer fills missing values ONLY based on training fold statistics.
The test fold is never seen during fitting. This mirrors real-world deployment
where you won't have access to future data when making predictions.

Why does data leakage matter in healthcare?

If you scale features using the mean of ALL data (including test), your model has indirectly "seen" the test set. Its accuracy score looks better than it actually is.

In a healthcare setting, an inflated score means you think your model catches 95% of heart disease cases when it really catches 85%. That 10% gap represents real patients who get missed.

Why Pipeline instead of manual preprocessing?

You could scale and encode data manually before splitting. But that approach is fragile. Any time you run cross-validation or tune hyperparameters, you'd need to remember to re-fit transformers on each fold separately.

scikit-learn's Pipeline automates this. It guarantees that transformers only see training data during fitting, even inside GridSearchCV or cross_val_score.

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook before moving on.

ⓧ I'd like to double check the full code

Here is the full notebook up to this point. Compare your cells against this:

# %% [markdown]
# # Heart Disease Prediction Pipeline
# End-to-end ML classification using scikit-learn Pipeline, model comparison, and evaluation.

# %% [markdown]
# ## 1. Setup and Imports

# %%
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from ucimlrepo import fetch_ucirepo

from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold, GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, roc_auc_score
from xgboost import XGBClassifier
import joblib

# %% [markdown]
# ## 2. Load and Explore the Dataset

# %%
heart_disease = fetch_ucirepo(id=45)
X = heart_disease.data.features
y = heart_disease.data.targets
y = (y.values.ravel() > 0).astype(int)

print(f"Dataset shape: {X.shape}")
print(f"Target distribution:\n{pd.Series(y).value_counts()}")
print(f"\nFeature types:\n{X.dtypes}")
print(f"\nMissing values:\n{X.isnull().sum()}")

# %%
numeric_cols = X.select_dtypes(include=[np.number]).columns.tolist()
plt.figure(figsize=(10, 8))
sns.heatmap(X[numeric_cols].corr(), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()

# %%
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
features_to_plot = ['age', 'thalach', 'chol']
titles = ['Age', 'Max Heart Rate', 'Cholesterol']

for ax, feat, title in zip(axes, features_to_plot, titles):
    for label in [0, 1]:
        ax.hist(X[feat][y == label].dropna(), alpha=0.6, label=f'{"No Disease" if label == 0 else "Disease"}', bins=20)
    ax.set_xlabel(title)
    ax.set_ylabel('Count')
    ax.legend()

plt.suptitle('Feature Distributions by Heart Disease Status')
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 3. Build a Leak-Free Preprocessing Pipeline

# %%
numeric_features = ['age', 'trestbps', 'chol', 'thalach', 'oldpeak']
categorical_features = ['sex', 'cp', 'fbs', 'restecg', 'exang', 'slope', 'ca', 'thal']

numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median')),
    ('scaler', StandardScaler())
])

categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('encoder', OneHotEncoder(handle_unknown='ignore', sparse_output=False))
])

preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, numeric_features),
        ('cat', categorical_transformer, categorical_features)
    ]
)

# %%
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"Training set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")
print(f"Train class balance: {np.bincount(y_train)}")
print(f"Test class balance: {np.bincount(y_test)}")

# %% [markdown]
# **Why this prevents data leakage:** The `ColumnTransformer` is inside the `Pipeline`.
# During cross-validation, the scaler computes mean/std ONLY from the training fold,
# and the imputer fills missing values ONLY based on training fold statistics.
# The test fold is never seen during fitting. This mirrors real-world deployment
# where you won't have access to future data when making predictions.

Your preprocessing pipeline is locked in and your data is cleanly split. Next up, you'll plug this preprocessor into three different models and see which one predicts heart disease most accurately.

Train and Compare Three Models

Your preprocessing pipeline is ready to transform raw patient data into model-ready features. Now comes the exciting part: training multiple models and finding out which one predicts heart disease most accurately.

You'll train three different classifiers and evaluate them using cross-validation. Cross-validation tests the model on multiple different data subsets, giving a much more reliable performance estimate than a single train-test split.

In this step, get ready to:
  • Create three full pipelines (preprocessing + classifier) for Logistic Regression, Random Forest, and XGBoost.
  • Run 5-fold cross-validation and compare F1 scores with a bar chart.
  • Tune the best model's hyperparameters using GridSearchCV.
Create three pipelines and run cross-validation

Each pipeline bundles the same preprocessor you already built with a different classification algorithm. This means every model sees identically transformed data, making the comparison fair.

  • Add a new markdown cell in your notebook with the section heading:
# %% [markdown]
# ## 4. Train and Compare Three Models
  • In a new code cell below the heading, define the three pipelines and run cross-validation by pasting this code:
# %%
# Define three pipelines with different classifiers
pipelines = {
    'Logistic Regression': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', LogisticRegression(max_iter=1000, random_state=42))
    ]),
    'Random Forest': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', RandomForestClassifier(n_estimators=100, random_state=42))
    ]),
    'XGBoost': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', XGBClassifier(
            n_estimators=100,
            eval_metric='logloss',
            random_state=42
        ))
    ])
}

# Run 5-fold stratified cross-validation
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
results = {}

for name, pipeline in pipelines.items():
    scores = cross_val_score(pipeline, X_train, y_train, cv=cv, scoring='f1')
    results[name] = scores
    print(f"{name}: F1 = {scores.mean():.4f} (+/- {scores.std():.4f})")

What does this code do?

Each pipeline pairs the same preprocessor with a different model family:

  • Logistic Regression is a linear model. It's the simplest option and serves as a strong baseline.
  • Random Forest is an ensemble of decision trees (bagging). It captures non-linear relationships.
  • XGBoost is a gradient-boosted tree ensemble. It's the industry standard for tabular data.

StratifiedKFold splits training data into 5 folds while preserving the class balance in each fold. cross_val_score trains and evaluates on each fold, giving you 5 F1 scores per model.

  • Run the cell. You should see three printed lines showing the F1 score and standard deviation for each model.

What is an F1 score?

F1 is the harmonic mean of precision and recall. It balances catching as many disease cases as possible (recall) with avoiding false alarms (precision). A higher F1 means the model handles both goals well.

Seeing an error when the cell runs?

  • If you see a NameError for preprocessor, make sure you ran all the cells from Step 3 first. The preprocessor variable must exist in memory.
  • If XGBoost raises an import error, confirm you installed it with the correct package name. Return to your terminal and run the pip install command from Step 1 again.

Still stuck? Help me debug my cross-validation cell

Visualize the model comparison

Numbers are great, but a bar chart makes the comparison obvious at a glance. This is also a useful visualization for your portfolio.

  • Create a new code cell below and paste this visualization code:
# %%
# Visualize model comparison
model_names = list(results.keys())
means = [results[name].mean() for name in model_names]
stds = [results[name].std() for name in model_names]

plt.figure(figsize=(8, 5))
bars = plt.bar(model_names, means, yerr=stds, capsize=5, color=['#2196F3', '#4CAF50', '#FF9800'])
plt.ylabel('F1 Score (5-Fold CV)')
plt.title('Model Comparison: Cross-Validated F1 Scores')
plt.ylim(0.7, 1.0)

for bar, mean in zip(bars, means):
    plt.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.01,
             f'{mean:.3f}', ha='center', va='bottom', fontweight='bold')

plt.tight_layout()
plt.show()

What does this code do?

  • yerr=stds adds error bars showing the standard deviation across folds. Smaller bars mean more consistent performance.
  • The loop at the bottom labels each bar with its exact mean F1 score so you can compare without squinting at the y-axis.
  • Run the cell. You should see a bar chart with three colored bars showing F1 scores, each labeled with a numeric value and topped with error bars.

Chart is blank or not displaying?

  • Make sure you ran the cross-validation cell above first. The results dictionary must be populated before this cell can plot it.
  • If the chart renders but looks squished, check that the Jupyter output panel is wide enough in VS Code. You can drag the panel border to resize.

Need help? My bar chart isn't showing up

Tune the best model with GridSearchCV

Cross-validation showed you which model family works best. Now you'll squeeze more performance out of it by searching for the optimal hyperparameters. These are settings you choose before training (like how many trees to grow) rather than values the model learns from data.

  • Create a new code cell and paste this hyperparameter tuning code:
# %%
# Hyperparameter tuning for Random Forest
param_grid = {
    'classifier__n_estimators': [50, 100, 200],
    'classifier__max_depth': [3, 5, 10, None],
    'classifier__min_samples_split': [2, 5, 10]
}

grid_search = GridSearchCV(
    pipelines['Random Forest'],
    param_grid,
    cv=cv,
    scoring='f1',
    n_jobs=-1
)

grid_search.fit(X_train, y_train)
print(f"Best parameters: {grid_search.best_params_}")
print(f"Best F1 score: {grid_search.best_score_:.4f}")

What does this code do?

  • param_grid defines every combination of hyperparameters to try. The classifier__ prefix tells GridSearchCV to reach inside the pipeline and set parameters on the classifier step.
  • GridSearchCV tries all 36 combinations (3 x 4 x 3) using the same 5-fold cross-validation, then reports the best-performing set.
  • n_jobs=-1 uses all your CPU cores to run combinations in parallel, speeding up the search.
  • Run the cell. This may take 10-30 seconds depending on your machine. You should see the best parameter combination and its F1 score printed.

Why tune Random Forest specifically?

Random Forest typically performs well on this dataset and its hyperparameters are intuitive to explain in interviews. In your own projects, you'd tune whichever model scored highest in cross-validation.

Cell is taking a very long time?

  • GridSearchCV is training 36 model combinations across 5 folds. That's 180 training runs total. On older hardware this can take a minute or two.
  • If it hangs for more than 2 minutes, try changing n_jobs=-1 to n_jobs=1. Some systems have issues with parallel processing in Jupyter notebooks.

Still having issues? Help me debug GridSearchCV performance

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook file before moving on.

ⓧ I'd like to double check the full code

Here is the full notebook code up to this point. Compare it against your notebook to spot any differences.

# %% [markdown]
# # Heart Disease Prediction Pipeline
# End-to-end ML classification using scikit-learn Pipeline, model comparison, and evaluation.

# %% [markdown]
# ## 1. Setup and Imports

# %%
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from ucimlrepo import fetch_ucirepo

from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold, GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, roc_auc_score
from xgboost import XGBClassifier
import joblib

# %% [markdown]
# ## 2. Load and Explore the Dataset

# %%
heart_disease = fetch_ucirepo(id=45)
X = heart_disease.data.features
y = heart_disease.data.targets
y = (y.values.ravel() > 0).astype(int)

print(f"Dataset shape: {X.shape}")
print(f"Target distribution:\n{pd.Series(y).value_counts()}")
print(f"\nFeature types:\n{X.dtypes}")
print(f"\nMissing values:\n{X.isnull().sum()}")

# %%
numeric_cols = X.select_dtypes(include=[np.number]).columns.tolist()
plt.figure(figsize=(10, 8))
sns.heatmap(X[numeric_cols].corr(), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()

# %%
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
features_to_plot = ['age', 'thalach', 'chol']
titles = ['Age', 'Max Heart Rate', 'Cholesterol']

for ax, feat, title in zip(axes, features_to_plot, titles):
    for label in [0, 1]:
        ax.hist(X[feat][y == label].dropna(), alpha=0.6, label=f'{"No Disease" if label == 0 else "Disease"}', bins=20)
    ax.set_xlabel(title)
    ax.set_ylabel('Count')
    ax.legend()

plt.suptitle('Feature Distributions by Heart Disease Status')
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 3. Build a Leak-Free Preprocessing Pipeline

# %%
numeric_features = ['age', 'trestbps', 'chol', 'thalach', 'oldpeak']
categorical_features = ['sex', 'cp', 'fbs', 'restecg', 'exang', 'slope', 'ca', 'thal']

numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median')),
    ('scaler', StandardScaler())
])

categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('encoder', OneHotEncoder(handle_unknown='ignore', sparse_output=False))
])

preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, numeric_features),
        ('cat', categorical_transformer, categorical_features)
    ]
)

# %%
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"Training set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")
print(f"Train class balance: {np.bincount(y_train)}")
print(f"Test class balance: {np.bincount(y_test)}")

# %% [markdown]
# **Why this prevents data leakage:** The `ColumnTransformer` is inside the `Pipeline`.
# During cross-validation, the scaler computes mean/std ONLY from the training fold,
# and the imputer fills missing values ONLY based on training fold statistics.
# The test fold is never seen during fitting.

# %% [markdown]
# ## 4. Train and Compare Three Models

# %%
# Define three pipelines with different classifiers
pipelines = {
    'Logistic Regression': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', LogisticRegression(max_iter=1000, random_state=42))
    ]),
    'Random Forest': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', RandomForestClassifier(n_estimators=100, random_state=42))
    ]),
    'XGBoost': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', XGBClassifier(
            n_estimators=100,
            eval_metric='logloss',
            random_state=42
        ))
    ])
}

# Run 5-fold stratified cross-validation
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
results = {}

for name, pipeline in pipelines.items():
    scores = cross_val_score(pipeline, X_train, y_train, cv=cv, scoring='f1')
    results[name] = scores
    print(f"{name}: F1 = {scores.mean():.4f} (+/- {scores.std():.4f})")

# %%
# Visualize model comparison
model_names = list(results.keys())
means = [results[name].mean() for name in model_names]
stds = [results[name].std() for name in model_names]

plt.figure(figsize=(8, 5))
bars = plt.bar(model_names, means, yerr=stds, capsize=5, color=['#2196F3', '#4CAF50', '#FF9800'])
plt.ylabel('F1 Score (5-Fold CV)')
plt.title('Model Comparison: Cross-Validated F1 Scores')
plt.ylim(0.7, 1.0)

for bar, mean in zip(bars, means):
    plt.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.01,
             f'{mean:.3f}', ha='center', va='bottom', fontweight='bold')

plt.tight_layout()
plt.show()

# %%
# Hyperparameter tuning for Random Forest
param_grid = {
    'classifier__n_estimators': [50, 100, 200],
    'classifier__max_depth': [3, 5, 10, None],
    'classifier__min_samples_split': [2, 5, 10]
}

grid_search = GridSearchCV(
    pipelines['Random Forest'],
    param_grid,
    cv=cv,
    scoring='f1',
    n_jobs=-1
)

grid_search.fit(X_train, y_train)
print(f"Best parameters: {grid_search.best_params_}")
print(f"Best F1 score: {grid_search.best_score_:.4f}")

You now have a tuned model that's been rigorously validated. Next up, you'll evaluate this champion on the held-out test set and extract the feature importances that explain its predictions.

Evaluate Your Best Model

Your tuned Random Forest is ready. But a model is only useful if you can prove it works on data it has never seen before.

In this step, you'll evaluate your best model on the held-out test set using the metrics that interviewers actually ask about: precision, recall, F1, and ROC-AUC. You'll also uncover which patient features matter most for predicting heart disease.

In this step, get ready to:
  • Generate predictions and a classification report on the test set.
  • Visualize the confusion matrix and ROC curve.
  • Extract and interpret feature importances from your best model.
Generate predictions and the classification report

The classification report is the standard way to see precision, recall, and F1 for each class at a glance. These are the metrics interviewers expect you to know cold.

  • Add a new markdown cell at the bottom of your notebook with the section heading, then add a code cell below it with the following prediction code:
# %% [markdown]
# ## 5. Evaluate Your Best Model

# %%
# Use the tuned model for final evaluation
best_model = grid_search.best_estimator_
y_pred = best_model.predict(X_test)
y_prob = best_model.predict_proba(X_test)[:, 1]

# Classification report
print("Classification Report:")
print(classification_report(y_test, y_pred, target_names=['No Disease', 'Disease']))

What does this code do?

  • grid_search.best_estimator_ retrieves the full pipeline (preprocessor + classifier) with the best hyperparameters from your grid search.
  • y_pred stores the hard class predictions (0 or 1) on the test set.
  • y_prob stores the probability of disease (class 1) for each patient. You need these probabilities to plot the ROC curve later.
  • classification_report prints precision, recall, and F1 for both classes in a formatted table.
  • Run the cell by pressing Shift+Enter.
  • You should see a table with precision, recall, F1-score, and support columns for both No Disease and Disease classes.

Seeing a NameError?

If you see NameError: name 'grid_search' is not defined, make sure you ran all cells from Step 4 first. The kernel needs the previous cells to be executed in order.

Help me debug my classification report cell.

Plot the confusion matrix and ROC curve

Numbers in a table are useful, but visualizations make patterns obvious at a glance. A confusion matrix shows exactly where your model makes mistakes, and the ROC curve shows how well your model separates the two classes across all threshold values.

  • Add a new code cell and paste the following visualization code:
# %%
# Confusion Matrix
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Confusion matrix heatmap
cm = confusion_matrix(y_test, y_pred)
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=axes[0],
            xticklabels=['No Disease', 'Disease'],
            yticklabels=['No Disease', 'Disease'])
axes[0].set_xlabel('Predicted')
axes[0].set_ylabel('Actual')
axes[0].set_title('Confusion Matrix')

# ROC Curve
fpr, tpr, thresholds = roc_curve(y_test, y_prob)
auc_score = roc_auc_score(y_test, y_prob)
axes[1].plot(fpr, tpr, color='#2196F3', lw=2, label=f'ROC Curve (AUC = {auc_score:.3f})')
axes[1].plot([0, 1], [0, 1], color='gray', linestyle='--')
axes[1].set_xlabel('False Positive Rate')
axes[1].set_ylabel('True Positive Rate')
axes[1].set_title('ROC Curve')
axes[1].legend()

plt.tight_layout()
plt.show()

What does this code do?

  • The confusion matrix heatmap shows four quadrants: true positives, true negatives, false positives, and false negatives. The numbers in each box are raw patient counts.
  • roc_curve plots the trade-off between true positive rate and false positive rate at every probability threshold.
  • The AUC score (Area Under the Curve) summarizes overall performance in a single number. A score of 1.0 means perfect separation between classes. A score of 0.5 means the model is no better than random guessing.
  • Run the cell.
  • You should see a side-by-side figure with a blue heatmap on the left (confusion matrix) and a blue ROC curve on the right that arcs above the gray diagonal line.

Plots not appearing?

If you see text output but no plots, make sure your Jupyter kernel is selected in VS Code (check the top-right corner of the notebook). Also confirm that matplotlib and seaborn imported successfully in your first cell.

Help me fix my confusion matrix and ROC curve plots.

Now add a markdown cell that explains why these metrics matter in a healthcare context. This kind of interpretation is what separates a junior candidate from someone who understands the domain.

  • Add a new markdown cell below the visualization cell with this explanation:
# %% [markdown]
# **Why F1 and Recall matter here:** In healthcare, missing a patient who actually
# has heart disease (False Negative) is more dangerous than a false alarm (False Positive).
# Recall measures how many actual disease cases we catch. F1 balances precision and recall.
# Accuracy alone can be misleading with imbalanced classes.

Why recall over accuracy?

Imagine a dataset where 90% of patients are healthy. A model that always predicts "No Disease" gets 90% accuracy but catches zero actual heart disease cases. That is a deadly model.

Recall specifically measures what percentage of actual disease cases you detect. In healthcare, you want high recall even if it means more false alarms. F1 balances recall with precision so you don't send every patient to the cardiologist.

Extract and visualize feature importances

Knowing which features drive predictions is just as important as the prediction itself. In an interview, explaining why your model makes certain predictions shows you understand the data, not just the algorithm.

  • Add a new code cell and paste the following feature importance code:
# %%
# Feature Importance
rf_model = best_model.named_steps['classifier']
feature_names = best_model.named_steps['preprocessor'].get_feature_names_out()

importances = rf_model.feature_importances_
importance_df = pd.DataFrame({
    'feature': feature_names,
    'importance': importances
}).sort_values('importance', ascending=False).head(10)

plt.figure(figsize=(10, 6))
sns.barplot(data=importance_df, x='importance', y='feature', palette='viridis')
plt.title('Top 10 Feature Importances')
plt.xlabel('Importance Score')
plt.ylabel('Feature')
plt.tight_layout()
plt.show()

print("\nTop 5 features for predicting heart disease:")
for i, row in importance_df.head(5).iterrows():
    print(f"  - {row['feature']}: {row['importance']:.4f}")

What does this code do?

  • best_model.named_steps['classifier'] reaches inside the pipeline to grab the trained Random Forest model directly.
  • get_feature_names_out() retrieves the actual feature names after preprocessing (including one-hot encoded categorical columns).
  • feature_importances_ is a Random Forest attribute that scores each feature by how much it reduces impurity across all trees. Higher scores mean the feature is more useful for splitting patients into disease vs. no-disease groups.
  • The bar chart visualizes the top 10 features, and the print loop lists the top 5 with their importance scores.
  • Run the cell.
  • You should see a horizontal bar chart showing the top 10 features ranked by importance, followed by a printed list of the top 5 features with their scores.

How to interpret your feature importances

You'll likely see features like chest pain type and max heart rate near the top. These are clinically meaningful. Chest pain type directly relates to cardiac symptoms, and lower max heart rate during exercise is a known indicator of cardiovascular problems.

If an interviewer asks "how do you explain your model's decisions?", feature importances are your first answer. They show that the model learned real medical patterns, not noise.

Seeing unfamiliar feature names?

After one-hot encoding, categorical features appear with prefixes like cat__. For example, cat__cp_4.0 means chest pain type 4. Numeric features have the num__ prefix. This is normal behavior from the ColumnTransformer.

Help me understand the feature names in my importance chart.

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook file.

ⓧ I'd like to double check the full code

Here is the complete notebook up to and including this step. Compare it against your file:

# %% [markdown]
# # Heart Disease Prediction Pipeline
# End-to-end ML classification using scikit-learn Pipeline, model comparison, and evaluation.

# %% [markdown]
# ## 1. Setup and Imports

# %%
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from ucimlrepo import fetch_ucirepo

from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold, GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, roc_auc_score
from xgboost import XGBClassifier
import joblib

# %% [markdown]
# ## 2. Load and Explore the Dataset

# %%
heart_disease = fetch_ucirepo(id=45)
X = heart_disease.data.features
y = heart_disease.data.targets
y = (y.values.ravel() > 0).astype(int)

print(f"Dataset shape: {X.shape}")
print(f"Target distribution:\n{pd.Series(y).value_counts()}")
print(f"\nFeature types:\n{X.dtypes}")
print(f"\nMissing values:\n{X.isnull().sum()}")

# %%
numeric_cols = X.select_dtypes(include=[np.number]).columns.tolist()
plt.figure(figsize=(10, 8))
sns.heatmap(X[numeric_cols].corr(), annot=True, cmap='coolwarm', center=0, fmt='.2f')
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.show()

# %%
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
features_to_plot = ['age', 'thalach', 'chol']
titles = ['Age', 'Max Heart Rate', 'Cholesterol']

for ax, feat, title in zip(axes, features_to_plot, titles):
    for label in [0, 1]:
        ax.hist(X[feat][y == label].dropna(), alpha=0.6, label=f'{"No Disease" if label == 0 else "Disease"}', bins=20)
    ax.set_xlabel(title)
    ax.set_ylabel('Count')
    ax.legend()

plt.suptitle('Feature Distributions by Heart Disease Status')
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 3. Build a Leak-Free Preprocessing Pipeline

# %%
numeric_features = ['age', 'trestbps', 'chol', 'thalach', 'oldpeak']
categorical_features = ['sex', 'cp', 'fbs', 'restecg', 'exang', 'slope', 'ca', 'thal']

numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median')),
    ('scaler', StandardScaler())
])

categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('encoder', OneHotEncoder(handle_unknown='ignore', sparse_output=False))
])

preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, numeric_features),
        ('cat', categorical_transformer, categorical_features)
    ]
)

# %%
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"Training set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")
print(f"Train class balance: {np.bincount(y_train)}")
print(f"Test class balance: {np.bincount(y_test)}")

# %% [markdown]
# **Why this prevents data leakage:** The `ColumnTransformer` is inside the `Pipeline`.
# During cross-validation, the scaler computes mean/std ONLY from the training fold,
# and the imputer fills missing values ONLY based on training fold statistics.
# The test fold is never seen during fitting.

# %% [markdown]
# ## 4. Train and Compare Three Models

# %%
pipelines = {
    'Logistic Regression': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', LogisticRegression(max_iter=1000, random_state=42))
    ]),
    'Random Forest': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', RandomForestClassifier(n_estimators=100, random_state=42))
    ]),
    'XGBoost': Pipeline([
        ('preprocessor', preprocessor),
        ('classifier', XGBClassifier(
            n_estimators=100,
            eval_metric='logloss',
            random_state=42
        ))
    ])
}

cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
results = {}

for name, pipeline in pipelines.items():
    scores = cross_val_score(pipeline, X_train, y_train, cv=cv, scoring='f1')
    results[name] = scores
    print(f"{name}: F1 = {scores.mean():.4f} (+/- {scores.std():.4f})")

# %%
model_names = list(results.keys())
means = [results[name].mean() for name in model_names]
stds = [results[name].std() for name in model_names]

plt.figure(figsize=(8, 5))
bars = plt.bar(model_names, means, yerr=stds, capsize=5, color=['#2196F3', '#4CAF50', '#FF9800'])
plt.ylabel('F1 Score (5-Fold CV)')
plt.title('Model Comparison: Cross-Validated F1 Scores')
plt.ylim(0.7, 1.0)

for bar, mean in zip(bars, means):
    plt.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.01,
             f'{mean:.3f}', ha='center', va='bottom', fontweight='bold')

plt.tight_layout()
plt.show()

# %%
param_grid = {
    'classifier__n_estimators': [50, 100, 200],
    'classifier__max_depth': [3, 5, 10, None],
    'classifier__min_samples_split': [2, 5, 10]
}

grid_search = GridSearchCV(
    pipelines['Random Forest'],
    param_grid,
    cv=cv,
    scoring='f1',
    n_jobs=-1
)

grid_search.fit(X_train, y_train)
print(f"Best parameters: {grid_search.best_params_}")
print(f"Best F1 score: {grid_search.best_score_:.4f}")

# %% [markdown]
# ## 5. Evaluate Your Best Model

# %%
best_model = grid_search.best_estimator_
y_pred = best_model.predict(X_test)
y_prob = best_model.predict_proba(X_test)[:, 1]

print("Classification Report:")
print(classification_report(y_test, y_pred, target_names=['No Disease', 'Disease']))

# %%
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

cm = confusion_matrix(y_test, y_pred)
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=axes[0],
            xticklabels=['No Disease', 'Disease'],
            yticklabels=['No Disease', 'Disease'])
axes[0].set_xlabel('Predicted')
axes[0].set_ylabel('Actual')
axes[0].set_title('Confusion Matrix')

fpr, tpr, thresholds = roc_curve(y_test, y_prob)
auc_score = roc_auc_score(y_test, y_prob)
axes[1].plot(fpr, tpr, color='#2196F3', lw=2, label=f'ROC Curve (AUC = {auc_score:.3f})')
axes[1].plot([0, 1], [0, 1], color='gray', linestyle='--')
axes[1].set_xlabel('False Positive Rate')
axes[1].set_ylabel('True Positive Rate')
axes[1].set_title('ROC Curve')
axes[1].legend()

plt.tight_layout()
plt.show()

# %% [markdown]
# **Why F1 and Recall matter here:** In healthcare, missing a patient who actually
# has heart disease (False Negative) is more dangerous than a false alarm (False Positive).
# Recall measures how many actual disease cases we catch. F1 balances precision and recall.
# Accuracy alone can be misleading with imbalanced classes.

# %%
rf_model = best_model.named_steps['classifier']
feature_names = best_model.named_steps['preprocessor'].get_feature_names_out()

importances = rf_model.feature_importances_
importance_df = pd.DataFrame({
    'feature': feature_names,
    'importance': importances
}).sort_values('importance', ascending=False).head(10)

plt.figure(figsize=(10, 6))
sns.barplot(data=importance_df, x='importance', y='feature', palette='viridis')
plt.title('Top 10 Feature Importances')
plt.xlabel('Importance Score')
plt.ylabel('Feature')
plt.tight_layout()
plt.show()

print("\nTop 5 features for predicting heart disease:")
for i, row in importance_df.head(5).iterrows():
    print(f"  - {row['feature']}: {row['importance']:.4f}")

Your evaluation is complete. You now have a full model report showing exactly how your pipeline performs, where it makes mistakes, and which features drive its decisions. This is the kind of analysis that makes an ML project interview-ready.

Secret mission

Deploy Your Model as a Prediction Function

Your model works great inside a notebook, but real-world ML needs models that can accept new patient data and return predictions on demand. In this secret mission, you'll save your trained pipeline to disk and build a reusable prediction function that mirrors how production ML systems work.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • Project folder on your Desktop (contains heart_disease_pipeline.ipynb, heart_disease_model.joblib, and your virtual environment).
  • Jupyter extension in VS Code.

Keep everything running

No action needed. Choose this if you plan to keep experimenting with the notebook, trying new models, or extending the pipeline.

  • Your project folder and virtual environment stay on disk ready for next time.
  • The Jupyter extension remains installed in VS Code for future notebooks.

Pause - I'll come back to this later

Shut down any running Jupyter kernels to free up memory. Your files remain safe on disk for when you return.

  • In VS Code, open the Command Palette with Cmd+Shift+P (macOS) or Ctrl+Shift+P (Windows).
  • Type Jupyter: Restart Kernel and select it to stop the running kernel.
  • If your virtual environment is still active in the terminal, deactivate it by running:
deactivate

Everything else stays in place. When you come back, reactivate your virtual environment and rerun the notebook cells.

Delete - I don't want to use this again

Remove all project resources and start fresh. This deletes your notebook, saved model, and virtual environment permanently.

  • If your virtual environment is still active, deactivate it first by running:
deactivate
  • Navigate to your Desktop and delete the project folder by running:
cd ~/Desktop
rm -rf heart-disease-pipeline

Windows users

On Windows, use this command instead:

cd ~\Desktop
rmdir /s /q heart-disease-pipeline
  • To remove the Jupyter extension from VS Code, open the Extensions sidebar with Cmd+Shift+X (macOS) or Ctrl+Shift+X (Windows).
  • Find the Jupyter extension, right-click it, and select Uninstall.

Nice Work!

Nice Work!

You just built a complete, interview-ready machine learning pipeline that predicts heart disease from real clinical data.

You've learned how to:

  • Build a leak-free preprocessing pipeline using ColumnTransformer and Pipeline to handle mixed feature types without data leakage.
  • Train, compare, and tune three models (Logistic Regression, Random Forest, and XGBoost) using cross-validation and GridSearchCV.
  • Evaluate model performance using precision, recall, F1, ROC-AUC, and feature importances to interpret results in a healthcare context.
  • Saved and deployed your best model as a reusable prediction function using joblib.

Ready to quiz yourself?