Heart Disease Prediction Pipeline

Build a full ML classification pipeline to predict heart disease with scikit-learn.

Introduction

30 Second Summary

Every year, doctors assess patients using a combination of clinical measurements like blood pressure, cholesterol levels, and heart rate data. Machine learning can analyse these same measurements to predict whether a patient is at risk of heart disease, helping catch warning signs faster.

In this project, you will build a complete machine learning classification pipeline that predicts heart disease from real clinical patient data using scikit-learn, the standard Python ML library.

What You'll Build

A Python classification pipeline that loads the UCI Heart Disease dataset, explores and visualises the data, preprocesses features, trains three different models, evaluates them with proper metrics, and tunes the best performer with GridSearchCV.

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

  • A trained and tuned Random Forest model that predicts heart disease from patient features like age, cholesterol, and blood pressure.
  • A complete ML pipeline script (heart_disease_ml.py) covering data loading, EDA, preprocessing, model training, evaluation, and hyperparameter tuning.
  • A clear understanding of the full supervised learning workflow that maps directly to university ML courses.
  • Secret Mission: Build a feature importance visualisation and a reusable prediction function that diagnoses new patients.

Are there any prerequisites?

You should be comfortable with basic Python (variables, functions, loops) and have some familiarity with NumPy. No prior machine learning experience is required.

Before We Start

To predict heart disease with machine learning, you need a proper Python environment with the right packages installed.

In this step, you will set up a project folder, create an isolated virtual environment, and install the four libraries your pipeline needs: scikit-learn, pandas, matplotlib, and seaborn.

In this step, get ready to:
  • Create a project folder and virtual environment.
  • Install scikit-learn, pandas, matplotlib, and seaborn.
  • Verify the installation and create your main script file.
Check your Python version

This project requires Python 3.11 or higher because scikit-learn 1.8.0 and pandas 3.0.3 both depend on it.

  • Open your terminal by pressing Cmd+Space (macOS) or the Windows key (Windows) and typing Terminal (macOS) or PowerShell (Windows).
  • Check your Python version by running this command:

macOS

python3 --version

Windows

python --version

You should see output like Python 3.11.x or higher (e.g. Python 3.12.4).

✔️ I see version 3.11 or higher

Your Python version is compatible. Move on to the next section below.

ⓧ I see an older version

Your system has an older Python installed. You need 3.11 or higher for this project.

  • Download the latest Python installer from python.org/downloads.
  • Run the installer. On Windows, check the Add python.exe to PATH box before clicking Install Now.
  • Close and reopen your terminal, then re-run the version check.

macOS users with a shadowed version

macOS may still resolve python3 to an older system version even after installing a newer one. If the version still shows below 3.11 after installing, try running python3.11 --version or python3.12 --version to call the newer binary directly.

ⓧ Command not found

Python is not installed on your system yet.

  • Download the Python 3.11+ installer from python.org/downloads.
  • Run the installer. On Windows, check the Add python.exe to PATH box before clicking Install Now.
  • Close and reopen your terminal, then run the version check again.
Create a project folder and virtual environment
  • 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-ml
cd heart-disease-ml
  • Confirm the folder exists by running:
ls

The output should be empty since the folder is brand new.

  • Create a virtual environment by running this command:

macOS

python3 -m venv venv

Windows

python -m venv venv

What is a virtual environment?

A virtual environment is an isolated Python installation inside your project folder. Packages you install here stay separate from your system Python, so different projects can use different package versions without conflicts.

  • Activate the virtual environment by running:

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.

Don't see (venv) in your prompt?

On Windows, you may need to allow script execution first. Run Set-ExecutionPolicy -Scope CurrentUser RemoteSigned in PowerShell, then try the activate command again.

On macOS, make sure you typed source before the path. Without source, the script runs in a subprocess and does not affect your current terminal session.

Help me activate my virtual environment.

Install the required packages
  • Install scikit-learn, pandas, matplotlib, and seaborn by running:
pip install scikit-learn pandas matplotlib seaborn

What do these packages do?

scikit-learn provides the machine learning algorithms, preprocessing tools, and evaluation metrics you will use throughout this project.

pandas handles loading and manipulating tabular data. matplotlib and seaborn create the visualizations (heatmaps, distribution plots, confusion matrices).

The installation may take a minute. You should see a series of download and install messages ending with Successfully installed followed by a list of package names.

Seeing an installation error?

Make sure (venv) is visible in your prompt. If it is not, your virtual environment is not active and pip may be pointing to your system Python.

If you see a permissions error, do not use sudo. Reactivate your virtual environment instead.

Help me fix my pip install error.

Verify the installation and create the script file
  • Verify all four packages imported successfully by running:
python -c "import sklearn; import pandas; import matplotlib; import seaborn; print('All packages installed successfully!')"

You should see All packages installed successfully! printed in your terminal.

Seeing a ModuleNotFoundError?

This means python is not pointing to the virtual environment's interpreter. Confirm (venv) is in your prompt. If not, reactivate your virtual environment and try again.

Help me fix my ModuleNotFoundError.

  • Create the main script file for your project by running:

macOS

touch heart_disease_ml.py

Windows

New-Item heart_disease_ml.py
  • Confirm the file was created by running:
ls

You should see heart_disease_ml.py and the venv folder in the output.

  • 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.
  • In VS Code, click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop, select the heart-disease-ml folder, and click Open.
  • In the Explorer sidebar on the left, you should see heart_disease_ml.py listed.

Your environment is ready. Next up, you will load a real heart disease dataset and explore its structure with pandas and seaborn.

Load and Explore the Dataset

Your environment is ready and your packages are installed. Now it's time to start building the actual machine learning pipeline.

Every ML project begins with understanding the data. In this step, you'll load a real-world medical dataset, inspect its structure, and create visualizations that reveal patterns your model can learn from.

In this step, get ready to:
  • Load the UCI Heart Disease dataset into a pandas DataFrame.
  • Inspect the dataset's shape, data types, and missing values.
  • Create visualizations to uncover correlations and feature distributions.
Load the UCI Heart Disease dataset

The UCI Heart Disease (Cleveland) dataset is a classic ML teaching dataset. It contains 303 patient records with 13 clinical features like age, cholesterol, and blood pressure.

Your model will learn to predict whether a patient has heart disease based on these features. First, you need to load the data and give each column a meaningful name.

  • Open heart_disease_ml.py in VS Code by clicking it in the Explorer sidebar.
  • Add the following imports at the top of the file by copying this code:
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay

What do these imports do?

  • pandas loads and manipulates tabular data as DataFrames.
  • numpy provides numerical computing utilities.
  • matplotlib and seaborn create the charts and heatmaps you'll save as PNG files.
  • scikit-learn provides the ML models, evaluation metrics, and preprocessing tools you'll use throughout the project.

You won't see any output yet. The next chunk uses these libraries to load the data.

  • Add the data loading code below your imports by copying this code:
# ============================================================
# STEP 2: Load and Explore the Dataset
# ============================================================

# Column names for the Cleveland Heart Disease dataset
column_names = [
    'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
    'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]

# Load the dataset from UCI repository
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
df = pd.read_csv(url, names=column_names, na_values='?')

What does this code do?

  • column_names defines the 14 attribute names (13 features + 1 target) for the dataset. The raw file has no header row, so you provide the names manually.
  • pd.read_csv() downloads the dataset directly from the UCI repository and loads it into a DataFrame called df.
  • na_values='?' tells pandas to treat ? characters as missing values instead of strings.
  • Save heart_disease_ml.py.

You won't see output yet. The next section adds print statements that reveal what the data looks like.

Inspect the data structure

Before building any model, you need to understand what you're working with. These print statements reveal the dataset's shape, check for missing values, and show how the target variable is distributed.

  • Add the data inspection code below your loading code:
# Inspect the data
print("=" * 60)
print("DATASET OVERVIEW")
print("=" * 60)
print(f"Shape: {df.shape}")
print(f"\nData types:\n{df.dtypes}")
print(f"\nMissing values:\n{df.isnull().sum()}")
print(f"\nTarget distribution:\n{df['target'].value_counts()}")
print(f"\nFirst 5 rows:\n{df.head()}")

What does this code do?

  • df.shape shows the number of rows (patients) and columns (features + target).
  • df.dtypes reveals the data type of each column, which helps you spot columns that might need conversion.
  • df.isnull().sum() counts missing values per column. Missing data needs to be handled before training.
  • df['target'].value_counts() shows how many patients have vs. don't have heart disease. An imbalanced target can bias your model.
  • Save heart_disease_ml.py.
  • Run the script in your terminal by executing:
python heart_disease_ml.py

You should see the DATASET OVERVIEW header followed by the shape (303, 14), a list of data types, missing value counts, and the target distribution.

Not seeing any output?

  • Make sure you saved the file before running the command.
  • Check that your virtual environment is still activated. You should see the environment name in parentheses in your terminal prompt.
  • If you see a network error, the script downloads data from the internet. Confirm you have an active connection.

Help me debug why my script isn't printing dataset output.

What is the target distribution telling you?

The original target values range from 0 to 4, where 0 means no disease and 1-4 indicate varying degrees of heart disease. You'll convert this to a simple binary (0 or 1) in the next step.

Create visualizations

Visualizations help you spot patterns that raw numbers might hide. A correlation heatmap shows how features relate to each other, and distribution plots reveal how feature values differ between patients with and without heart disease.

  • Add the correlation heatmap code below your inspection code:
# Correlation heatmap
plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, fmt='.2f', cmap='coolwarm', center=0)
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.close()
print("\nSaved: correlation_heatmap.png")

What does this code do?

  • df.corr(numeric_only=True) calculates the Pearson correlation coefficient between every pair of numerical columns.
  • sns.heatmap() renders those correlations as a color-coded grid. Values close to 1 or -1 indicate strong relationships.
  • plt.savefig() saves the chart as a PNG file instead of displaying it in a window. This makes the script non-interactive and reproducible.
  • Save heart_disease_ml.py.

You won't see the saved file yet. The next chunk adds distribution plots, then you'll run the full script.

  • Add the feature distribution code below the heatmap code:
# Distribution of key features by target
fig, axes = plt.subplots(2, 2, figsize=(12, 10))
features_to_plot = ['age', 'thalach', 'chol', 'oldpeak']
for ax, feature in zip(axes.flatten(), features_to_plot):
    sns.histplot(data=df, x=feature, hue='target', kde=True, ax=ax)
    ax.set_title(f'{feature} Distribution by Target')
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.close()
print("Saved: feature_distributions.png")

What does this code do?

  • plt.subplots(2, 2) creates a 2x2 grid of charts so you can compare four features side by side.
  • The loop plots a histogram for each feature, colored by the target variable. Overlapping distributions suggest a feature doesn't separate the classes well. Separated distributions suggest a useful predictor.
  • kde=True adds a smooth curve over the histogram, making the distribution shape easier to read.
  • Save heart_disease_ml.py.
  • Run the script again:
python heart_disease_ml.py

You should see the full DATASET OVERVIEW output followed by Saved: correlation_heatmap.png and Saved: feature_distributions.png.

  • Confirm the files were created by running:
ls *.png

You should see correlation_heatmap.png and feature_distributions.png in the output.

Don't see the PNG files?

  • Make sure you're running the command from inside your project folder. Use pwd to check your current directory.
  • If you see a MatplotlibDeprecationWarning, you can safely ignore it. The files should still be created.
  • Check for typos in plt.savefig(). The filename must be a string in quotes.

Help me figure out why my PNG files aren't being created.

✔️ Awesome, I've got everything!

Great. Double check you've saved heart_disease_ml.py and that both PNG files appear in your project folder.

ⓧ I'd like to double check the full code

Here's the complete heart_disease_ml.py so far:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay

# ============================================================
# STEP 2: Load and Explore the Dataset
# ============================================================

# Column names for the Cleveland Heart Disease dataset
column_names = [
    'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
    'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]

# Load the dataset from UCI repository
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
df = pd.read_csv(url, names=column_names, na_values='?')

# Inspect the data
print("=" * 60)
print("DATASET OVERVIEW")
print("=" * 60)
print(f"Shape: {df.shape}")
print(f"\nData types:\n{df.dtypes}")
print(f"\nMissing values:\n{df.isnull().sum()}")
print(f"\nTarget distribution:\n{df['target'].value_counts()}")
print(f"\nFirst 5 rows:\n{df.head()}")

# Correlation heatmap
plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, fmt='.2f', cmap='coolwarm', center=0)
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.close()
print("\nSaved: correlation_heatmap.png")

# Distribution of key features by target
fig, axes = plt.subplots(2, 2, figsize=(12, 10))
features_to_plot = ['age', 'thalach', 'chol', 'oldpeak']
for ax, feature in zip(axes.flatten(), features_to_plot):
    sns.histplot(data=df, x=feature, hue='target', kde=True, ax=ax)
    ax.set_title(f'{feature} Distribution by Target')
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.close()
print("Saved: feature_distributions.png")

Your data is loaded and visualized. Next up, you'll preprocess the features to get them ready for model training.

Preprocess the Data

In the last step, you explored the raw dataset and saw that it has a few missing values and a multi-class target column (values 0 through 4).

Machine learning models cannot handle missing data or ambiguous targets. In this step, you will clean the dataset, split it into training and test sets, and scale the features so each one contributes equally to the model.

In this step, get ready to:
  • Handle missing values and convert the target column to binary.
  • Split data into training and test sets with an 80/20 ratio.
  • Scale numerical features using StandardScaler.
Handle missing values and convert the target to binary

The Cleveland dataset uses ? to represent missing values in two columns (ca and thal). Only about 6 rows are affected, so dropping them is safe.

The original target column has values 0 through 4, where 0 means no disease and 1-4 represent different disease severities. For binary classification, you will collapse these into two categories: 0 (no disease) and 1 (disease present).

  • In heart_disease_ml.py, add the following code at the bottom of the file to clean the data and convert the target:
# ============================================================
# STEP 3: Preprocess the Data
# ============================================================

df = df.dropna()
print(f"\nShape after dropping missing values: {df.shape}")

df['target'] = (df['target'] > 0).astype(int)
print(f"Binary target distribution:\n{df['target'].value_counts()}")

What does this code do?

  • df.dropna() removes any row that contains a missing value. Since only ~6 rows out of 303 have gaps, this is a safe approach that avoids introducing bias from imputation.
  • (df['target'] > 0).astype(int) converts the multi-class target (0-4) into a simple binary column. Any value greater than 0 becomes 1 (disease present).
  • Save heart_disease_ml.py.
  • Run the script in your terminal by running:
python3 heart_disease_ml.py
  • In the terminal output, look for the new lines at the bottom. You should see a shape like (297, 14) and a binary target distribution showing counts for 0 and 1.

Seeing a KeyError or unexpected shape?

  • Make sure the new code is below your Step 2 visualization code, not inside it.
  • Check that you typed 'target' with matching quotes and exact spelling.

Why is my DataFrame shape wrong after dropna?

Split features and target, then perform train/test split

Before training a model, you need to separate the input features from the target label. You also need to hold out a portion of data the model never sees during training, so you can fairly measure its performance later.

  • Add the following code below your binary target conversion to split the data:
X = df.drop('target', axis=1)
y = df['target']

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"\nTraining set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

What does this code do?

  • X holds all 13 feature columns (everything except the target). y holds the binary target label.
  • train_test_split divides the data into 80% training and 20% test. The stratify=y parameter ensures both sets have the same proportion of disease vs no-disease cases.
  • random_state=42 locks the split so you get reproducible results every time you run the script.

Why stratify the split?

If the split were purely random, you might accidentally put most disease cases in the test set and very few in training. Stratification guarantees both sets mirror the original class balance, giving the model a fair representation of each outcome.

  • Save heart_disease_ml.py and run the script again:
python3 heart_disease_ml.py
  • Scroll to the bottom of the terminal output. You should see the training set size (around 237) and the test set size (around 60).

Seeing a NameError for train_test_split?

  • This means the import at the top of the file is missing. Check that your imports include from sklearn.model_selection import train_test_split.
  • Make sure there are no typos in train_test_split (underscores, not hyphens).

Help me fix my train_test_split error.

Scale numerical features using StandardScaler

Features in this dataset have very different ranges. Age might be 29-77, while cholesterol ranges from 100-500.

StandardScaler transforms each feature to have a mean of 0 and a standard deviation of 1. This prevents features with large values from dominating the model's decisions.

  • Add the following code below your train/test split to scale the features:
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

Why fit on training data only?

fit_transform on the training set learns the mean and standard deviation from that data alone. You then call transform on the test set using those same values.

If you fitted the scaler on the entire dataset (including test data), the model would have indirect knowledge of test set statistics. This is called data leakage and produces unrealistically high accuracy scores.

  • Save heart_disease_ml.py.
  • Run the script one more time:
python3 heart_disease_ml.py
  • You won't see new printed output from the scaler itself. Confirm the script runs without errors and still shows the training/test set sizes from the previous chunk.

Seeing an error about X_train not being defined?

  • The scaler code must appear after the train_test_split call. Make sure all three chunks are in order at the bottom of your file.
  • Check that there are no extra blank lines or indentation issues breaking the flow.

Why is X_train not defined when I run my script?

✔️ Awesome, I've got everything!

Great. Double check you have saved heart_disease_ml.py and that the script runs without errors.

ⓧ I'd like to double check the full code

Here is the full heart_disease_ml.py file up to this point:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay

# ============================================================
# STEP 2: Load and Explore the Dataset
# ============================================================

column_names = [
    'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
    'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]

url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
df = pd.read_csv(url, names=column_names, na_values='?')

print("=" * 60)
print("DATASET OVERVIEW")
print("=" * 60)
print(f"Shape: {df.shape}")
print(f"\nData types:\n{df.dtypes}")
print(f"\nMissing values:\n{df.isnull().sum()}")
print(f"\nTarget distribution:\n{df['target'].value_counts()}")
print(f"\nFirst 5 rows:\n{df.head()}")

plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, fmt='.2f', cmap='coolwarm', center=0)
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.close()
print("\nSaved: correlation_heatmap.png")

fig, axes = plt.subplots(2, 2, figsize=(12, 10))
features_to_plot = ['age', 'thalach', 'chol', 'oldpeak']
for ax, feature in zip(axes.flatten(), features_to_plot):
    sns.histplot(data=df, x=feature, hue='target', kde=True, ax=ax)
    ax.set_title(f'{feature} Distribution by Target')
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.close()
print("Saved: feature_distributions.png")

# ============================================================
# STEP 3: Preprocess the Data
# ============================================================

df = df.dropna()
print(f"\nShape after dropping missing values: {df.shape}")

df['target'] = (df['target'] > 0).astype(int)
print(f"Binary target distribution:\n{df['target'].value_counts()}")

X = df.drop('target', axis=1)
y = df['target']

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"\nTraining set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

Your data is now clean, split, and scaled. Next up, you will train three different classification models and compare their performance to find the best predictor of heart disease.

Train and Compare Models

Your data is cleaned, preprocessed, and split into training and test sets. Now comes the core question of any ML project: which algorithm performs best on this particular dataset?

In this step, you'll train three different classifiers and compare them using proper evaluation metrics. This "model comparison" approach is standard practice in machine learning because no single algorithm dominates every dataset.

In this step, get ready to:
  • Train Logistic Regression, Decision Tree, and Random Forest classifiers on the scaled training data.
  • Evaluate each model with classification reports and cross-validation F1 scores.
  • Visualize the confusion matrix for the best-performing model.
Define and initialize three classifiers

Before training, you need to set up the models you want to compare. Each algorithm approaches classification differently.

A Logistic Regression draws a straight decision boundary. A Decision Tree splits data using a series of yes/no questions. A Random Forest combines many decision trees and takes their majority vote.

  • In heart_disease_ml.py, add the following code at the bottom of the file (after the StandardScaler lines from Step 3):
# ============================================================
# STEP 4: Train and Compare Models
# ============================================================

print("\n" + "=" * 60)
print("MODEL COMPARISON")
print("=" * 60)

models = {
    'Logistic Regression': LogisticRegression(max_iter=1000, random_state=42),
    'Decision Tree': DecisionTreeClassifier(random_state=42),
    'Random Forest': RandomForestClassifier(n_estimators=100, random_state=42)
}

What does this code do?

This creates a dictionary called models that maps a human-readable name to each classifier instance.

  • LogisticRegression(max_iter=1000) allows up to 1000 iterations to converge on a solution.
  • DecisionTreeClassifier() builds a tree that splits data on the most informative features.
  • RandomForestClassifier(n_estimators=100) trains 100 separate trees and averages their predictions.
  • random_state=42 ensures reproducible results every time you run the script.

You won't see a change yet. The next chunk uses this dictionary to train and evaluate each model.

Train each model and compare performance

Now you'll loop through all three models, train each one on the scaled training data, and evaluate their predictions. You'll use two evaluation methods: a classification report (which shows precision, recall, and F1 per class) and cross-validation (which tests generalization by training on different data folds).

  • Add the following code directly below the models dictionary you just wrote:
results = {}
for name, model in models.items():
    model.fit(X_train_scaled, y_train)
    y_pred = model.predict(X_test_scaled)
    print(f"\n--- {name} ---")
    print(classification_report(y_test, y_pred))
    
    cv_scores = cross_val_score(model, X_train_scaled, y_train, cv=5, scoring='f1')
    results[name] = cv_scores.mean()
    print(f"5-Fold CV F1 Score: {cv_scores.mean():.4f} (+/- {cv_scores.std():.4f})")

best_model_name = max(results, key=results.get)
print(f"\nBest model (by CV F1): {best_model_name} ({results[best_model_name]:.4f})")

What does this code do?

This loop trains each model and measures how well it predicts heart disease.

  • model.fit() trains the model on the scaled training data.
  • model.predict() generates predictions on the test set.
  • classification_report prints precision (of all predicted positives, how many were correct), recall (of all actual positives, how many were found), and F1-score (the harmonic mean of both).
  • cross_val_score with cv=5 splits the training data into 5 folds and trains/tests 5 times. This gives a more reliable performance estimate than a single train/test split.
  • The final two lines find which model scored the highest average F1 across the 5 folds.
  • Save heart_disease_ml.py.
  • Run the script in your terminal:
python heart_disease_ml.py
  • In the terminal output, scroll to the MODEL COMPARISON section. You should see three classification reports (one per model) with precision, recall, and F1 columns, plus a 5-Fold CV F1 Score for each.

Not seeing three classification reports?

  • Make sure you saved the file before running it.
  • Check that the models dictionary and the for loop are at the same indentation level (no extra indentation).
  • If you see NameError: name 'X_train_scaled' is not defined, the Step 3 code (StandardScaler section) is missing or has an error.

Help me debug why my model training loop isn't printing results.

Visualize the confusion matrix for Random Forest

A confusion matrix shows exactly where a model gets confused. It breaks predictions into four categories: true positives, true negatives, false positives, and false negatives. This gives you a clearer picture than a single accuracy number.

  • Add this code below the best_model_name print statement:
rf_model = models['Random Forest']
y_pred_rf = rf_model.predict(X_test_scaled)
cm = confusion_matrix(y_test, y_pred_rf)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=['No Disease', 'Disease'])
disp.plot(cmap='Blues')
plt.title('Random Forest - Confusion Matrix')
plt.tight_layout()
plt.savefig('confusion_matrix.png', dpi=100)
plt.close()
print("Saved: confusion_matrix.png")

What does this code do?

  • rf_model = models['Random Forest'] pulls the already-trained Random Forest from the models dictionary.
  • confusion_matrix() computes the 2x2 matrix of correct vs incorrect predictions for each class.
  • ConfusionMatrixDisplay renders it as a labelled heatmap with 'No Disease' and 'Disease' as the axis labels.
  • The plot is saved to confusion_matrix.png so you can review it anytime.
  • Save heart_disease_ml.py and run the script again:
python heart_disease_ml.py
  • Confirm that your terminal prints Saved: confusion_matrix.png at the end of the output.
  • Open confusion_matrix.png from your project folder. You should see a blue heatmap with No Disease and Disease labels on both axes, with numbers in each cell showing correct and incorrect predictions.

Don't see confusion_matrix.png?

  • Check that plt.savefig('confusion_matrix.png', dpi=100) appears before plt.close(). Closing the figure first discards it.
  • Make sure you're looking in the same folder where you run the script. Run ls in your terminal to confirm the file exists.
  • If the script errors before reaching the confusion matrix code, fix the earlier error first. The script runs top to bottom.

Help me figure out why my confusion matrix isn't saving.

How do I read a confusion matrix?

Each row represents actual labels. Each column represents predicted labels. The top-left cell is true negatives (correctly predicted No Disease). The bottom-right cell is true positives (correctly predicted Disease).

The off-diagonal cells are errors. Top-right is false positives (healthy patients flagged as sick). Bottom-left is false negatives (sick patients missed). In medical contexts, false negatives are usually more dangerous.

✔️ Awesome, I've got everything!

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

ⓧ I'd like to double check the full code

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay

# ============================================================
# STEP 2: Load and Explore the Dataset
# ============================================================

column_names = [
    'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
    'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]

url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
df = pd.read_csv(url, names=column_names, na_values='?')

print("=" * 60)
print("DATASET OVERVIEW")
print("=" * 60)
print(f"Shape: {df.shape}")
print(f"\nData types:\n{df.dtypes}")
print(f"\nMissing values:\n{df.isnull().sum()}")
print(f"\nTarget distribution:\n{df['target'].value_counts()}")
print(f"\nFirst 5 rows:\n{df.head()}")

plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, fmt='.2f', cmap='coolwarm', center=0)
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.close()
print("\nSaved: correlation_heatmap.png")

fig, axes = plt.subplots(2, 2, figsize=(12, 10))
features_to_plot = ['age', 'thalach', 'chol', 'oldpeak']
for ax, feature in zip(axes.flatten(), features_to_plot):
    sns.histplot(data=df, x=feature, hue='target', kde=True, ax=ax)
    ax.set_title(f'{feature} Distribution by Target')
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.close()
print("Saved: feature_distributions.png")

# ============================================================
# STEP 3: Preprocess the Data
# ============================================================

df = df.dropna()
print(f"\nShape after dropping missing values: {df.shape}")

df['target'] = (df['target'] > 0).astype(int)
print(f"Binary target distribution:\n{df['target'].value_counts()}")

X = df.drop('target', axis=1)
y = df['target']

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"\nTraining set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# ============================================================
# STEP 4: Train and Compare Models
# ============================================================

print("\n" + "=" * 60)
print("MODEL COMPARISON")
print("=" * 60)

models = {
    'Logistic Regression': LogisticRegression(max_iter=1000, random_state=42),
    'Decision Tree': DecisionTreeClassifier(random_state=42),
    'Random Forest': RandomForestClassifier(n_estimators=100, random_state=42)
}

results = {}
for name, model in models.items():
    model.fit(X_train_scaled, y_train)
    y_pred = model.predict(X_test_scaled)
    print(f"\n--- {name} ---")
    print(classification_report(y_test, y_pred))
    
    cv_scores = cross_val_score(model, X_train_scaled, y_train, cv=5, scoring='f1')
    results[name] = cv_scores.mean()
    print(f"5-Fold CV F1 Score: {cv_scores.mean():.4f} (+/- {cv_scores.std():.4f})")

best_model_name = max(results, key=results.get)
print(f"\nBest model (by CV F1): {best_model_name} ({results[best_model_name]:.4f})")

rf_model = models['Random Forest']
y_pred_rf = rf_model.predict(X_test_scaled)
cm = confusion_matrix(y_test, y_pred_rf)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=['No Disease', 'Disease'])
disp.plot(cmap='Blues')
plt.title('Random Forest - Confusion Matrix')
plt.tight_layout()
plt.savefig('confusion_matrix.png', dpi=100)
plt.close()
print("Saved: confusion_matrix.png")

You now have three trained models with clear performance metrics, and a confusion matrix showing exactly where Random Forest succeeds and fails. Next up, you'll squeeze even more performance out of the best model using hyperparameter tuning.

Tune the Best Model with GridSearchCV

In the previous step, you trained three models and compared their performance using cross-validation. You identified the strongest model based on F1 scores.

Now you'll take the Random Forest model and systematically search for the best combination of settings. This process is called hyperparameter tuning. Instead of guessing which values work best, you'll let GridSearchCV try every combination and report the winner.

In this step, get ready to:
  • Define a hyperparameter grid with multiple values for key Random Forest settings.
  • Run GridSearchCV with 5-fold cross-validation to find the best combination.
  • Evaluate the tuned model on the test set and compare it to the original.
Define the hyperparameter grid

A hyperparameter grid is a dictionary that maps each setting name to a list of values you want to test. GridSearchCV will try every possible combination of these values.

  • In heart_disease_ml.py, scroll to the bottom of your file (after the confusion matrix code) and add the hyperparameter tuning section:
# ============================================================
# STEP 5: Tune the Best Model with GridSearchCV
# ============================================================

print("\n" + "=" * 60)
print("HYPERPARAMETER TUNING (GridSearchCV)")
print("=" * 60)

# Define the combinations of settings to try
param_grid = {
    'n_estimators': [50, 100, 200],
    'max_depth': [None, 5, 10, 15],
    'min_samples_split': [2, 5, 10]
}

What does this code do?

Each key in param_grid is a hyperparameter name that Random Forest accepts. Each value is a list of options to try.

  • n_estimators: how many decision trees to build (50, 100, or 200).
  • max_depth: how deep each tree can grow (unlimited, 5, 10, or 15 levels).
  • min_samples_split: minimum samples required to split a node further (2, 5, or 10).

That gives 3 x 4 x 3 = 36 unique combinations to test.

You won't see a change yet. The next chunk uses this grid to run the actual search.

Run GridSearchCV with 5-fold cross-validation

GridSearchCV wraps a model and a parameter grid together. It trains the model once for every combination, using cross-validation each time to get a reliable score.

  • Add the following code directly below the param_grid dictionary:
# Set up GridSearchCV to try all combinations
grid_search = GridSearchCV(
    RandomForestClassifier(random_state=42),
    param_grid,
    cv=5,
    scoring='f1',
    n_jobs=-1,
    verbose=1
)

# Fit runs 36 combinations x 5 folds = 180 model trainings
grid_search.fit(X_train_scaled, y_train)

What does this code do?

  • GridSearchCV() creates the search object with a fresh RandomForestClassifier as the base model.
  • cv=5 means each combination is evaluated using 5-fold cross-validation (the data is split 5 ways, training on 4 and testing on 1, rotating each time).
  • scoring='f1' tells the search to rank combinations by their F1 score, which balances precision and recall.
  • n_jobs=-1 uses all available CPU cores to speed up the search.
  • The .fit() call trains 180 models (36 combinations times 5 folds) and stores the best one.

Why a fresh RandomForestClassifier instead of the trained one?

GridSearchCV needs an untrained estimator. It handles all the training internally so it can test each parameter combination from scratch with proper cross-validation.

You won't see a change yet. The next chunk prints the results and evaluates the tuned model.

Print the best parameters and evaluate the tuned model

Once GridSearchCV finishes, it stores the winning configuration and the fully trained best model. You'll extract both and run a final evaluation on the held-out test set.

  • Add the following code directly below the grid_search.fit() line:
print(f"\nBest parameters: {grid_search.best_params_}")
print(f"Best CV F1 score: {grid_search.best_score_:.4f}")

# Evaluate tuned model on test set
best_rf = grid_search.best_estimator_
y_pred_tuned = best_rf.predict(X_test_scaled)
print(f"\nTuned Random Forest - Test Set Results:")
print(classification_report(y_test, y_pred_tuned))

What does this code do?

  • grid_search.best_params_ is a dictionary of the winning hyperparameter values.
  • grid_search.best_score_ is the average F1 score of the best combination across all 5 folds.
  • grid_search.best_estimator_ is the actual trained model using those best parameters, ready to make predictions.
  • classification_report() prints the full precision, recall, and F1 breakdown so you can compare to the original Random Forest from Step 4.
  • Save heart_disease_ml.py.
  • Run the script by executing this command in your terminal:
python heart_disease_ml.py
  • Scroll to the bottom of the terminal output. You should see a HYPERPARAMETER TUNING (GridSearchCV) section showing the best parameters, the best CV F1 score, and a classification report for the tuned model.

The search takes a moment

GridSearchCV is training 180 models, so expect it to take 5-15 seconds depending on your machine. The verbose=1 setting prints progress so you know it's working.

Not seeing the tuning output?

  • Make sure you saved heart_disease_ml.py before running the script.
  • Check that the Step 5 code is placed after the confusion matrix block (not inside the Step 4 for loop).
  • If you see an indentation error, make sure the param_grid dictionary and grid_search variable are at the top level (no leading spaces).

My script runs but the tuning section doesn't appear. Help me debug.

✔️ Awesome, I've got everything!

Great. Double check you've saved your file.

ⓧ I'd like to double check the full code

Here is the complete heart_disease_ml.py file up to and including Step 5:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay

# ============================================================
# STEP 2: Load and Explore the Dataset
# ============================================================

column_names = [
    'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs',
    'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
]

url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"
df = pd.read_csv(url, names=column_names, na_values='?')

print("=" * 60)
print("DATASET OVERVIEW")
print("=" * 60)
print(f"Shape: {df.shape}")
print(f"\nData types:\n{df.dtypes}")
print(f"\nMissing values:\n{df.isnull().sum()}")
print(f"\nTarget distribution:\n{df['target'].value_counts()}")
print(f"\nFirst 5 rows:\n{df.head()}")

plt.figure(figsize=(12, 8))
sns.heatmap(df.corr(numeric_only=True), annot=True, fmt='.2f', cmap='coolwarm', center=0)
plt.title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.close()
print("\nSaved: correlation_heatmap.png")

fig, axes = plt.subplots(2, 2, figsize=(12, 10))
features_to_plot = ['age', 'thalach', 'chol', 'oldpeak']
for ax, feature in zip(axes.flatten(), features_to_plot):
    sns.histplot(data=df, x=feature, hue='target', kde=True, ax=ax)
    ax.set_title(f'{feature} Distribution by Target')
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.close()
print("Saved: feature_distributions.png")

# ============================================================
# STEP 3: Preprocess the Data
# ============================================================

df = df.dropna()
print(f"\nShape after dropping missing values: {df.shape}")

df['target'] = (df['target'] > 0).astype(int)
print(f"Binary target distribution:\n{df['target'].value_counts()}")

X = df.drop('target', axis=1)
y = df['target']

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"\nTraining set size: {X_train.shape[0]}")
print(f"Test set size: {X_test.shape[0]}")

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# ============================================================
# STEP 4: Train and Compare Models
# ============================================================

print("\n" + "=" * 60)
print("MODEL COMPARISON")
print("=" * 60)

models = {
    'Logistic Regression': LogisticRegression(max_iter=1000, random_state=42),
    'Decision Tree': DecisionTreeClassifier(random_state=42),
    'Random Forest': RandomForestClassifier(n_estimators=100, random_state=42)
}

results = {}
for name, model in models.items():
    model.fit(X_train_scaled, y_train)
    y_pred = model.predict(X_test_scaled)
    print(f"\n--- {name} ---")
    print(classification_report(y_test, y_pred))
    
    cv_scores = cross_val_score(model, X_train_scaled, y_train, cv=5, scoring='f1')
    results[name] = cv_scores.mean()
    print(f"5-Fold CV F1 Score: {cv_scores.mean():.4f} (+/- {cv_scores.std():.4f})")

best_model_name = max(results, key=results.get)
print(f"\nBest model (by CV F1): {best_model_name} ({results[best_model_name]:.4f})")

rf_model = models['Random Forest']
y_pred_rf = rf_model.predict(X_test_scaled)
cm = confusion_matrix(y_test, y_pred_rf)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=['No Disease', 'Disease'])
disp.plot(cmap='Blues')
plt.title('Random Forest - Confusion Matrix')
plt.tight_layout()
plt.savefig('confusion_matrix.png', dpi=100)
plt.close()
print("Saved: confusion_matrix.png")

# ============================================================
# STEP 5: Tune the Best Model with GridSearchCV
# ============================================================

print("\n" + "=" * 60)
print("HYPERPARAMETER TUNING (GridSearchCV)")
print("=" * 60)

param_grid = {
    'n_estimators': [50, 100, 200],
    'max_depth': [None, 5, 10, 15],
    'min_samples_split': [2, 5, 10]
}

grid_search = GridSearchCV(
    RandomForestClassifier(random_state=42),
    param_grid,
    cv=5,
    scoring='f1',
    n_jobs=-1,
    verbose=1
)

grid_search.fit(X_train_scaled, y_train)

print(f"\nBest parameters: {grid_search.best_params_}")
print(f"Best CV F1 score: {grid_search.best_score_:.4f}")

# Evaluate tuned model on test set
best_rf = grid_search.best_estimator_
y_pred_tuned = best_rf.predict(X_test_scaled)
print(f"\nTuned Random Forest - Test Set Results:")
print(classification_report(y_test, y_pred_tuned))

Secret mission

Feature Importance and Prediction Function

Your model can predict heart disease, but which features matter most? In this bonus mission, you'll extract feature importances from your tuned Random Forest and build a reusable prediction function that takes new patient data and returns a diagnosis with confidence scores.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs or cloud resources to worry about. Decide whether to keep your files for future reference, or delete them to free up disk space.

Resources you used:

  • Project folder on your Desktop (contains heart_disease_ml.py, your virtual environment, and generated PNG files).
  • Python packages installed inside the virtual environment (scikit-learn, pandas, matplotlib, seaborn).

Keep everything running

No action needed. Choose this if you want to keep experimenting with your heart disease prediction pipeline or revisit the code later.

  • Your project folder stays on your Desktop with all files intact.
  • The virtual environment remains ready to activate whenever you want to run the script again.
  • The generated PNG visualizations (correlation_heatmap.png, feature_distributions.png, confusion_matrix.png, feature_importance.png) stay in the folder for reference.

Pause - I'll come back to this later

Deactivate your virtual environment to free up your terminal session. Your files stay on disk for when you return.

  • If your virtual environment is still active (you see (venv) in your terminal prompt), deactivate it by running:
deactivate
  • Close VS Code.
  • When you come back, reactivate the virtual environment and run heart_disease_ml.py again.

Delete - I don't want to use this again

Remove all project files, the virtual environment, and generated images from your machine.

  • If your virtual environment is still active, deactivate it first by running:
deactivate
  • Navigate to your Desktop and delete the project folder. On macOS or Linux, run:
cd ~/Desktop
rm -rf [[PROJECT_FOLDER="your project folder name"]]
  • On Windows, run:
cd ~\Desktop
rmdir /s /q [[PROJECT_FOLDER="your project folder name"]]
  • This removes heart_disease_ml.py, all generated PNG files, and the virtual environment (including all installed packages).

Nice Work!

Nice Work!

You just built a complete machine learning classification pipeline from scratch using scikit-learn. You went from raw patient data all the way to a tuned model that predicts heart disease.

You've learned how to:

  • Loaded and explored a real-world medical dataset with pandas and seaborn, including correlation heatmaps and feature distribution plots.
  • Preprocessed data by handling missing values, encoding the target variable to binary, performing a stratified train/test split, and scaling features with StandardScaler.
  • Trained and compared three classifiers using cross-validation and proper evaluation metrics (precision, recall, F1, confusion matrix), then tuned hyperparameters with GridSearchCV.
  • Built a feature importance visualization and a reusable predict_patient() function that takes new patient data and returns a diagnosis with confidence scores.

Ready to quiz yourself?