Predict Machine Failures with ML

Train classifiers on sensor data to predict industrial equipment failures.

Introduction

30 Second Summary

Every piece of industrial equipment eventually breaks down. The difference between a planned repair and a catastrophic failure on an offshore platform can cost over $500,000 per day in lost production.

In this project, you will build a machine learning model that predicts equipment failure from real sensor data. You will train and compare classifiers, handle class imbalance, and interpret results through a mechanical engineering lens.

What You'll Build

A predictive maintenance classifier that predicts whether industrial equipment will fail based on operational sensor readings like temperature, torque, rotational speed, and tool wear.

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

  • A complete Jupyter notebook with exploratory data analysis, preprocessing, model training, evaluation, and interpretation of a real industrial dataset.
  • A trained Random Forest model with feature importance analysis linking predictions back to mechanical engineering principles like heat dissipation, overstrain, and power failure.
  • A comparison of three classification models evaluated using precision, recall, and F1-score metrics appropriate for imbalanced data.
  • Secret Mission: Build a multi-class failure mode classifier that predicts not just IF a machine will fail, but WHAT TYPE of failure will occur.

Are there any prerequisites?

This project assumes basic Python knowledge and familiarity with classification algorithms (Logistic Regression, Decision Trees, Random Forest) from a course like Udemy's Machine Learning A-Z. You should also have VS Code, Python, and libraries like pandas, NumPy, scikit-learn, matplotlib, and seaborn already installed.

Set Up and Load the Data

Before you can predict equipment failures, you need real sensor data from industrial machines. The AI4I 2020 Predictive Maintenance Dataset from the UCI Machine Learning Repository contains 10,000 data points of manufacturing equipment readings.

In this step, you'll import your Python libraries, load the dataset, and check how many machines actually failed. Failures are rare in industrial equipment, which is realistic but creates a challenge for machine learning.

In this step, get ready to:
  • Import all required Python libraries for data analysis and machine learning.
  • Load the AI4I 2020 Predictive Maintenance Dataset from UCI.
  • Explore the dataset and verify the class imbalance.
Import the libraries and load the dataset

Your project needs several libraries working together. pandas and NumPy handle data manipulation. matplotlib and seaborn produce visualizations. scikit-learn provides the machine learning models.

  • In predictive_maintenance.py, add the following imports at the top of the file:
# Predictive Maintenance: Machine Failure Prediction
# Using the AI4I 2020 Predictive Maintenance Dataset (UCI ML Repository)

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
from sklearn.preprocessing import LabelEncoder
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 and numpy handle loading, cleaning, and manipulating the sensor data as DataFrames and arrays.
  • matplotlib and seaborn produce charts and heatmaps to visualize patterns in the data.
  • ucimlrepo downloads the dataset directly from the UCI Machine Learning Repository without needing a separate file download.
  • The sklearn imports bring in the classification models (Logistic Regression, Decision Tree, Random Forest) and evaluation tools you'll use later.

You won't see output yet. The next chunk loads the actual data.

  • Add this code below your imports to fetch the dataset and combine features with targets into a single DataFrame:
# ============================================================
# STEP 1: Load the Dataset
# ============================================================

dataset = fetch_ucirepo(id=601)
X = dataset.data.features
y = dataset.data.targets

# Combine features and targets for exploration
df = pd.concat([X, y], axis=1)

What does this code do?

  • fetch_ucirepo(id=601) downloads the AI4I 2020 Predictive Maintenance Dataset directly from the UCI repository. ID 601 is the unique identifier for this specific dataset.
  • dataset.data.features returns the input features (sensor readings) as a pandas DataFrame, stored in X.
  • dataset.data.targets returns the target variable (whether the machine failed) as a DataFrame, stored in y.
  • pd.concat([X, y], axis=1) merges features and targets side-by-side into one DataFrame called df so you can explore everything together.

You won't see output yet. The next section adds print statements to explore the data.

Explore the dataset structure

Before building any model, you need to understand what the data looks like. Each row represents one machine operation, and each column is a sensor reading or identifier.

  • Add these print statements below your data loading code to inspect the dataset:
print("Dataset Shape:", df.shape)
print("\nColumn Names:")
print(df.columns.tolist())
print("\nFirst 5 rows:")
print(df.head())
print("\nData Types:")
print(df.dtypes)
print("\nBasic Statistics:")
print(df.describe())

What does each print reveal?

  • df.shape shows 10,000 rows and the number of columns, confirming the full dataset loaded.
  • df.columns.tolist() lists all column names including sensor features: Air temperature [K], Process temperature [K], Rotational speed [rpm], Torque [Nm], and Tool wear [min].
  • df.head() displays the first 5 rows so you can see actual values.
  • df.dtypes confirms which columns are numeric vs. categorical (the Type column is L/M/H for product quality).
  • df.describe() gives min, max, mean, and standard deviation for each numeric feature.
  • Save predictive_maintenance.py.
  • Run the script in your terminal by running:
python predictive_maintenance.py

You should see the dataset shape printed as (10000, 14), followed by column names, sample rows, data types, and summary statistics.

Seeing a ModuleNotFoundError?

  • If the error mentions ucimlrepo, confirm the package is installed by running pip install ucimlrepo in your terminal.
  • If the error mentions sklearn, install scikit-learn by running pip install scikit-learn.
  • Make sure you're running the script from the same environment where you installed the packages.

Help me troubleshoot this error

Verify the class imbalance

The target variable is Machine failure. It's binary: 0 means no failure, 1 means failure. In real industrial equipment, failures are rare events. This creates class imbalance, where one class vastly outnumbers the other.

  • Add this code below your print statements to check and visualize the failure rate:
# Check class distribution
print("\n--- Class Distribution ---")
print(df['Machine failure'].value_counts())
print(f"\nFailure Rate: {df['Machine failure'].mean() * 100:.2f}%")

# Visualize class imbalance
fig, ax = plt.subplots(figsize=(6, 4))
df['Machine failure'].value_counts().plot(kind='bar', ax=ax, color=['steelblue', 'coral'])
ax.set_title('Machine Failure Distribution')
ax.set_xlabel('Machine Failure (0 = No, 1 = Yes)')
ax.set_ylabel('Count')
ax.set_xticklabels(['No Failure', 'Failure'], rotation=0)
plt.tight_layout()
plt.savefig('class_distribution.png', dpi=100)
plt.show()

Why does class imbalance matter?

You'll see approximately 339 failures out of 10,000 data points. That's a 3.4% failure rate. If a model simply predicted "no failure" every time, it would be 96.6% accurate while being completely useless.

This is why you'll use specialized metrics like precision, recall, and F1-score later instead of plain accuracy. You'll also use class_weight='balanced' when training models to compensate for this imbalance.

  • Save predictive_maintenance.py.
  • Run the script again in your terminal:
python predictive_maintenance.py

You should see the class distribution counts, the failure rate percentage (approximately 3.39%), and a bar chart showing the extreme imbalance between the two classes.

The chart is also saved as class_distribution.png in your project folder.

Bar chart not appearing?

  • If the chart window doesn't pop up, check that your matplotlib backend supports display. The chart is still saved as class_distribution.png regardless. Open it from your file explorer to verify.
  • If you see a blank chart, make sure you added all the code including plt.show() at the end.

Help me troubleshoot my matplotlib chart

✔️ Awesome, I've got everything!

Great. Double check you've saved predictive_maintenance.py and move on to the next step.

ⓧ I'd like to double check the full code

# Predictive Maintenance: Machine Failure Prediction
# Using the AI4I 2020 Predictive Maintenance Dataset (UCI ML Repository)

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
from sklearn.preprocessing import LabelEncoder
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 1: Load the Dataset
# ============================================================

dataset = fetch_ucirepo(id=601)
X = dataset.data.features
y = dataset.data.targets

# Combine features and targets for exploration
df = pd.concat([X, y], axis=1)

print("Dataset Shape:", df.shape)
print("\nColumn Names:")
print(df.columns.tolist())
print("\nFirst 5 rows:")
print(df.head())
print("\nData Types:")
print(df.dtypes)
print("\nBasic Statistics:")
print(df.describe())

# Check class distribution
print("\n--- Class Distribution ---")
print(df['Machine failure'].value_counts())
print(f"\nFailure Rate: {df['Machine failure'].mean() * 100:.2f}%")

# Visualize class imbalance
fig, ax = plt.subplots(figsize=(6, 4))
df['Machine failure'].value_counts().plot(kind='bar', ax=ax, color=['steelblue', 'coral'])
ax.set_title('Machine Failure Distribution')
ax.set_xlabel('Machine Failure (0 = No, 1 = Yes)')
ax.set_ylabel('Count')
ax.set_xticklabels(['No Failure', 'Failure'], rotation=0)
plt.tight_layout()
plt.savefig('class_distribution.png', dpi=100)
plt.show()

Your data is loaded and you've confirmed the class imbalance. Next up, you'll dig deeper with exploratory data analysis and engineer new features using your mechanical engineering knowledge.

Explore the Data and Engineer Features

You have your dataset loaded and confirmed the class imbalance. Only about 3.4% of the machines actually fail.

Now you need to understand what the sensor data looks like before and during failures. You will visualize the features, identify correlations, and then create new features based on mechanical engineering principles that explain how these failures physically happen.

In this step, get ready to:
  • Visualize how sensor readings differ between failed and non-failed machines.
  • Engineer new features based on physical failure mechanisms.
  • Prepare the dataset for model training with a stratified train-test split.
Visualize feature distributions and correlations

Before training a model, you need to see how each sensor reading behaves when a machine is about to fail versus when it runs normally. Overlapping histograms reveal which features have the strongest signal for predicting failure.

  • In predictive_maintenance.py, add the following code below your existing plt.show() line to create feature distribution histograms colored by failure status:
# ============================================================
# STEP 2: Exploratory Data Analysis and Feature Engineering
# ============================================================

# Feature distributions by failure status
numerical_features = ['Air temperature [K]', 'Process temperature [K]',
                      'Rotational speed [rpm]', 'Torque [Nm]', 'Tool wear [min]']

fig, axes = plt.subplots(2, 3, figsize=(14, 8))
axes = axes.flatten()
for i, feature in enumerate(numerical_features):
    ax = axes[i]
    df[df['Machine failure'] == 0][feature].hist(ax=ax, bins=30, alpha=0.5, label='No Failure', color='steelblue')
    df[df['Machine failure'] == 1][feature].hist(ax=ax, bins=30, alpha=0.5, label='Failure', color='coral')
    ax.set_title(feature)
    ax.legend()
axes[5].set_visible(False)
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.show()

What does this code do?

This creates a 2x3 grid of histograms, one for each numerical sensor feature. Each histogram overlays two distributions: blue for non-failed machines and coral for failed machines.

Where the two colors separate, the feature carries a strong failure signal. Where they overlap heavily, the feature alone does not distinguish failure from normal operation.

  • Save predictive_maintenance.py and run it.
  • You should see a grid of 5 histograms pop up. Look for features where the coral (failure) bars separate from the blue (no failure) bars. Torque and rotational speed typically show the clearest separation.

Not seeing histograms?

Make sure you saved the file before running. Check that the numerical_features list uses the exact column names with brackets and units, for example Air temperature [K]. A typo in any column name will cause a KeyError.

Help me debug my histogram code

  • Next, add the correlation heatmap below your histogram code to see how features relate to each other and to the failure target:
# Correlation heatmap
fig, ax = plt.subplots(figsize=(8, 6))
correlation_matrix = df[numerical_features + ['Machine failure']].corr()
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0, ax=ax)
ax.set_title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.show()

What does the correlation heatmap show?

The heatmap shows how strongly each pair of features correlates (values range from -1 to 1). Red means strong positive correlation. Blue means strong negative correlation.

Look at the Machine failure row. Torque typically shows the highest positive correlation with failure, while rotational speed shows a negative correlation. This matches engineering intuition: high torque at low speeds strains mechanical components.

  • Save and run the script.
  • You should see a color-coded heatmap with numbers in each cell. Both feature_distributions.png and correlation_heatmap.png are now saved in your project folder.

Heatmap looks blank or garbled?

Check that you imported seaborn as sns at the top of your file. Also confirm that numerical_features was defined in the histogram block above. If you skipped the histograms, this code cannot find the variable.

Help me fix my seaborn heatmap

Engineer features using mechanical engineering knowledge

Raw sensor readings alone do not capture how failures physically happen. In mechanical systems, failures emerge from combinations of conditions, not single readings in isolation.

Why engineer new features?

The AI4I 2020 dataset documents specific physical thresholds for each failure mode. Heat dissipation failure (HDF) occurs when the temperature difference between process and air is too low. Power failure (PWF) happens when power falls outside the 3500-9000W range. Overstrain failure (OSF) depends on the product of tool wear and torque exceeding a threshold.

By creating features that directly compute these physical quantities, you give the model the same information an engineer uses to diagnose failures.

  • Add the following code below your heatmap to create three engineered features:
# Feature Engineering using mechanical engineering knowledge
df['Temp_difference'] = df['Process temperature [K]'] - df['Air temperature [K]']
df['Power_W'] = df['Torque [Nm]'] * df['Rotational speed [rpm]'] * (2 * np.pi / 60)
df['Wear_torque_interaction'] = df['Tool wear [min]'] * df['Torque [Nm]']

print("\n--- Engineered Features (first 5 rows) ---")
print(df[['Temp_difference', 'Power_W', 'Wear_torque_interaction']].head())

What do these engineered features represent?

  • Temp_difference (process temp minus air temp): heat dissipation failure occurs when this difference drops below 8.6K at low speeds. A smaller gap means the machine cannot shed heat fast enough.
  • Power_W (torque multiplied by rotational speed in rad/s): power failure triggers when this value falls below 3500W or exceeds 9000W. The 2 * np.pi / 60 converts RPM to radians per second.
  • Wear_torque_interaction (tool wear multiplied by torque): overstrain failure occurs when this product exceeds specific thresholds per quality type (11000/12000/13000 minNm for L/M/H).
  • Save and run the script.
  • You should see a table of the first 5 rows of your three new columns printed in the terminal. Temp_difference values will be around 8-12K. Power_W values will be in the thousands.

Seeing NaN values in engineered features?

Check that the column names in your code match the exact names from the dataset, including spaces and brackets. For example, Process temperature [K] must include the space before the bracket. Also confirm numpy is imported as np at the top of your file.

Help me debug my feature engineering code

Prepare the data for modeling

Your DataFrame now has both raw sensor columns and engineered features. Before training a model, you need to encode the categorical Type column (L/M/H product quality) as a number, build the final feature matrix, and split into training and test sets.

  • Add the data preparation code below your engineered features:
# Prepare features for modeling
le = LabelEncoder()
df['Type_encoded'] = le.fit_transform(df['Type'])

feature_columns = numerical_features + ['Type_encoded', 'Temp_difference', 'Power_W', 'Wear_torque_interaction']
X_model = df[feature_columns]
y_model = df['Machine failure']

# Train-test split with stratification
X_train, X_test, y_train, y_test = train_test_split(
    X_model, y_model, test_size=0.2, random_state=42, stratify=y_model
)

print(f"\nTraining set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")
print(f"Training failure rate: {y_train.mean() * 100:.2f}%")
print(f"Test failure rate: {y_test.mean() * 100:.2f}%")

What does this code do?

  • LabelEncoder() converts the categorical Type column (L, M, H) into numeric values (0, 1, 2) so the model can process it.
  • feature_columns combines the 5 original numerical features, the encoded type, and your 3 engineered features into a single list of 9 columns.
  • stratify=y_model ensures both the training and test sets preserve the original 3.4% failure rate. Without stratification, the test set could randomly end up with very few or no failures, making evaluation unreliable.

Why stratified splitting matters for imbalanced data

With only 339 failures in 10,000 rows, a random split could put most failures in training and almost none in testing. Stratified splitting guarantees the class ratio is preserved in both sets, so your evaluation metrics reflect real-world performance.

  • Save and run the script.
  • You should see the training set has 8000 samples, the test set has 2000 samples, and both show a failure rate of approximately 3.39%.

Failure rates don't match between training and test?

Double check that you included stratify=y_model in your train_test_split call. Without it, the split is purely random and the rates will differ. Also confirm y_model is set to df['Machine failure'] (the binary target column).

Help me fix my train-test split

✔️ Awesome, I've got everything!

Great. Double check you have saved your file before moving on.

ⓧ I'd like to double check the full code

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

# Predictive Maintenance: Machine Failure Prediction
# Using the AI4I 2020 Predictive Maintenance Dataset (UCI ML Repository)

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
from sklearn.preprocessing import LabelEncoder
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 1: Load the Dataset
# ============================================================

dataset = fetch_ucirepo(id=601)
X = dataset.data.features
y = dataset.data.targets

# Combine features and targets for exploration
df = pd.concat([X, y], axis=1)

print("Dataset Shape:", df.shape)
print("\nColumn Names:")
print(df.columns.tolist())
print("\nFirst 5 rows:")
print(df.head())
print("\nData Types:")
print(df.dtypes)
print("\nBasic Statistics:")
print(df.describe())

# Check class distribution
print("\n--- Class Distribution ---")
print(df['Machine failure'].value_counts())
print(f"\nFailure Rate: {df['Machine failure'].mean() * 100:.2f}%")

# Visualize class imbalance
fig, ax = plt.subplots(figsize=(6, 4))
df['Machine failure'].value_counts().plot(kind='bar', ax=ax, color=['steelblue', 'coral'])
ax.set_title('Machine Failure Distribution')
ax.set_xlabel('Machine Failure (0 = No, 1 = Yes)')
ax.set_ylabel('Count')
ax.set_xticklabels(['No Failure', 'Failure'], rotation=0)
plt.tight_layout()
plt.savefig('class_distribution.png', dpi=100)
plt.show()

# ============================================================
# STEP 2: Exploratory Data Analysis and Feature Engineering
# ============================================================

# Feature distributions by failure status
numerical_features = ['Air temperature [K]', 'Process temperature [K]',
                      'Rotational speed [rpm]', 'Torque [Nm]', 'Tool wear [min]']

fig, axes = plt.subplots(2, 3, figsize=(14, 8))
axes = axes.flatten()
for i, feature in enumerate(numerical_features):
    ax = axes[i]
    df[df['Machine failure'] == 0][feature].hist(ax=ax, bins=30, alpha=0.5, label='No Failure', color='steelblue')
    df[df['Machine failure'] == 1][feature].hist(ax=ax, bins=30, alpha=0.5, label='Failure', color='coral')
    ax.set_title(feature)
    ax.legend()
axes[5].set_visible(False)
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.show()

# Correlation heatmap
fig, ax = plt.subplots(figsize=(8, 6))
correlation_matrix = df[numerical_features + ['Machine failure']].corr()
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0, ax=ax)
ax.set_title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.show()

# Feature Engineering using mechanical engineering knowledge
df['Temp_difference'] = df['Process temperature [K]'] - df['Air temperature [K]']
df['Power_W'] = df['Torque [Nm]'] * df['Rotational speed [rpm]'] * (2 * np.pi / 60)
df['Wear_torque_interaction'] = df['Tool wear [min]'] * df['Torque [Nm]']

print("\n--- Engineered Features (first 5 rows) ---")
print(df[['Temp_difference', 'Power_W', 'Wear_torque_interaction']].head())

# Prepare features for modeling
le = LabelEncoder()
df['Type_encoded'] = le.fit_transform(df['Type'])

feature_columns = numerical_features + ['Type_encoded', 'Temp_difference', 'Power_W', 'Wear_torque_interaction']
X_model = df[feature_columns]
y_model = df['Machine failure']

# Train-test split with stratification
X_train, X_test, y_train, y_test = train_test_split(
    X_model, y_model, test_size=0.2, random_state=42, stratify=y_model
)

print(f"\nTraining set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")
print(f"Training failure rate: {y_train.mean() * 100:.2f}%")
print(f"Test failure rate: {y_test.mean() * 100:.2f}%")

Your data is explored, enriched with engineering-driven features, and split into training and test sets. Next up, you will train three classification models and see which one catches the most failures.

Train and Compare Classification Models

Your data is prepped with engineered features that capture real mechanical failure physics. Now it's time to see which algorithm does the best job at catching equipment failures before they happen.

In predictive maintenance, a missed failure is far more expensive than a false alarm. One unplanned shutdown on a rotating machine can cost hundreds of thousands of dollars. Your models need to maximize recall for the failure class, meaning they catch as many real failures as possible.

In this step, get ready to:
  • Train Logistic Regression, Decision Tree, and Random Forest classifiers with class imbalance handling.
  • Evaluate each model using precision, recall, and F1-score for the failure class.
  • Visualize confusion matrices to compare model performance side by side.
Train and evaluate three classifiers

With only 3.4% of your data representing failures, a model could predict "no failure" every single time and still achieve 96.6% accuracy. That number looks great but is completely useless. You need metrics that focus specifically on how well the model detects the rare failure cases.

Why class_weight='balanced'?

The class_weight='balanced' parameter tells each model to pay more attention to the rare failure cases. It automatically adjusts weights inversely proportional to class frequencies.

Without this, models would learn to always predict "no failure" because that's correct 96.6% of the time. With balanced weights, the model treats each failure case as roughly 29x more important than a non-failure case.

  • In predictive_maintenance.py, scroll to the bottom of your file (after the train-test split output).
  • Add the following code to define your three models, train them, and print classification reports:
# ============================================================
# STEP 3: Train and Compare Classification Models
# ============================================================

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

results = {}

for name, model in models.items():
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    results[name] = {
        'model': model,
        'predictions': y_pred
    }
    print(f"\n{'='*50}")
    print(f"Model: {name}")
    print('='*50)
    print(classification_report(y_test, y_pred, target_names=['No Failure', 'Failure']))

What does this code do?

  • The models dictionary stores three classifiers, each configured with class_weight='balanced' to handle the 96.6%/3.4% imbalance.
  • The for loop trains each model on X_train, generates predictions on X_test, and stores both the model object and its predictions in the results dictionary.
  • The classification_report prints precision, recall, and F1-score for both classes. Focus on the Failure row to see how well each model catches real failures.
  • Save predictive_maintenance.py and run the file.
  • You should see three classification reports printed in your terminal, one for each model. Each report shows precision, recall, and F1-score for both the No Failure and Failure classes.

How do I read the classification report?

  • Precision for Failure: of all the times the model predicted failure, what percentage were actually failures? Higher precision means fewer false alarms.
  • Recall for Failure: of all the actual failures, what percentage did the model catch? Higher recall means fewer missed breakdowns.
  • F1-score: the harmonic mean of precision and recall. A balanced measure when you care about both.

Seeing a ConvergenceWarning?

If you see a warning from Logistic Regression about convergence, double check that you have max_iter=1000 in the LogisticRegression constructor. The default iteration limit is too low for this dataset.

If the reports all show 0.00 for the Failure class, make sure class_weight='balanced' is set on each model. Without it, models ignore the rare failure cases entirely.

Help me troubleshoot my classification reports.

Visualize model performance with confusion matrices

A confusion matrix shows exactly where each model gets things right and wrong. You can see at a glance how many failures were caught versus missed, and how many false alarms were raised.

  • Add the following code below your training loop to create side-by-side confusion matrices for all three models:
# Confusion matrices side by side
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
for i, (name, result) in enumerate(results.items()):
    ConfusionMatrixDisplay.from_predictions(
        y_test, result['predictions'],
        display_labels=['No Failure', 'Failure'],
        cmap='Blues',
        ax=axes[i]
    )
    axes[i].set_title(name)
plt.tight_layout()
plt.savefig('confusion_matrices.png', dpi=100)
plt.show()

What does this code do?

  • Creates a 1x3 grid of subplots so all three confusion matrices appear side by side.
  • ConfusionMatrixDisplay.from_predictions takes the true labels and predicted labels and plots the matrix directly.
  • The top-left cell shows correct "No Failure" predictions. The bottom-right cell shows correct "Failure" predictions. The bottom-left cell shows missed failures (the most dangerous in practice).
  • Save predictive_maintenance.py and run the file again.
  • You should see three confusion matrices displayed side by side, and a confusion_matrices.png file saved in your project folder.
  • Compare the bottom-right cell across all three matrices. The model with the highest number there catches the most real failures.

Not seeing the confusion matrices?

Make sure this code comes after the training loop, not inside it. The results dictionary must be fully populated before you loop through it to plot.

If the plot window appears blank, check that plt.tight_layout() and plt.show() are both present at the end of the block.

Help me debug my confusion matrix visualization.

✔️ 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

Here is the complete predictive_maintenance.py file up to this point:

# Predictive Maintenance: Machine Failure Prediction
# Using the AI4I 2020 Predictive Maintenance Dataset (UCI ML Repository)

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
from sklearn.preprocessing import LabelEncoder
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 1: Load the Dataset
# ============================================================

dataset = fetch_ucirepo(id=601)
X = dataset.data.features
y = dataset.data.targets

# Combine features and targets for exploration
df = pd.concat([X, y], axis=1)

print("Dataset Shape:", df.shape)
print("\nColumn Names:")
print(df.columns.tolist())
print("\nFirst 5 rows:")
print(df.head())
print("\nData Types:")
print(df.dtypes)
print("\nBasic Statistics:")
print(df.describe())

# Check class distribution
print("\n--- Class Distribution ---")
print(df['Machine failure'].value_counts())
print(f"\nFailure Rate: {df['Machine failure'].mean() * 100:.2f}%")

# Visualize class imbalance
fig, ax = plt.subplots(figsize=(6, 4))
df['Machine failure'].value_counts().plot(kind='bar', ax=ax, color=['steelblue', 'coral'])
ax.set_title('Machine Failure Distribution')
ax.set_xlabel('Machine Failure (0 = No, 1 = Yes)')
ax.set_ylabel('Count')
ax.set_xticklabels(['No Failure', 'Failure'], rotation=0)
plt.tight_layout()
plt.savefig('class_distribution.png', dpi=100)
plt.show()

# ============================================================
# STEP 2: Exploratory Data Analysis and Feature Engineering
# ============================================================

# Feature distributions by failure status
numerical_features = ['Air temperature [K]', 'Process temperature [K]',
                      'Rotational speed [rpm]', 'Torque [Nm]', 'Tool wear [min]']

fig, axes = plt.subplots(2, 3, figsize=(14, 8))
axes = axes.flatten()
for i, feature in enumerate(numerical_features):
    ax = axes[i]
    df[df['Machine failure'] == 0][feature].hist(ax=ax, bins=30, alpha=0.5, label='No Failure', color='steelblue')
    df[df['Machine failure'] == 1][feature].hist(ax=ax, bins=30, alpha=0.5, label='Failure', color='coral')
    ax.set_title(feature)
    ax.legend()
axes[5].set_visible(False)
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.show()

# Correlation heatmap
fig, ax = plt.subplots(figsize=(8, 6))
correlation_matrix = df[numerical_features + ['Machine failure']].corr()
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0, ax=ax)
ax.set_title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.show()

# Feature Engineering using mechanical engineering knowledge
df['Temp_difference'] = df['Process temperature [K]'] - df['Air temperature [K]']
df['Power_W'] = df['Torque [Nm]'] * df['Rotational speed [rpm]'] * (2 * np.pi / 60)
df['Wear_torque_interaction'] = df['Tool wear [min]'] * df['Torque [Nm]']

print("\n--- Engineered Features (first 5 rows) ---")
print(df[['Temp_difference', 'Power_W', 'Wear_torque_interaction']].head())

# Prepare features for modeling
le = LabelEncoder()
df['Type_encoded'] = le.fit_transform(df['Type'])

feature_columns = numerical_features + ['Type_encoded', 'Temp_difference', 'Power_W', 'Wear_torque_interaction']
X_model = df[feature_columns]
y_model = df['Machine failure']

# Train-test split with stratification
X_train, X_test, y_train, y_test = train_test_split(
    X_model, y_model, test_size=0.2, random_state=42, stratify=y_model
)

print(f"\nTraining set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")
print(f"Training failure rate: {y_train.mean() * 100:.2f}%")
print(f"Test failure rate: {y_test.mean() * 100:.2f}%")

# ============================================================
# STEP 3: Train and Compare Classification Models
# ============================================================

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

results = {}

for name, model in models.items():
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    results[name] = {
        'model': model,
        'predictions': y_pred
    }
    print(f"\n{'='*50}")
    print(f"Model: {name}")
    print('='*50)
    print(classification_report(y_test, y_pred, target_names=['No Failure', 'Failure']))

# Confusion matrices side by side
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
for i, (name, result) in enumerate(results.items()):
    ConfusionMatrixDisplay.from_predictions(
        y_test, result['predictions'],
        display_labels=['No Failure', 'Failure'],
        cmap='Blues',
        ax=axes[i]
    )
    axes[i].set_title(name)
plt.tight_layout()
plt.savefig('confusion_matrices.png', dpi=100)
plt.show()

You now have three trained models and a clear picture of how each one performs on failure detection. Next up, you'll dig into feature importance to understand which physical parameters drive the Random Forest's predictions.

Analyze Feature Importance and Interpret Results

In the previous step, you trained three classifiers and compared their performance using classification reports and confusion matrices. You now know which model catches the most failures.

Knowing which model performs best is only half the story. The real value for a mechanical engineer is understanding why certain sensor readings predict failure. In this step, you will extract feature importance from the Random Forest model and connect the results back to the physical failure mechanisms you already understand from engineering.

In this step, get ready to:
  • Extract and visualize feature importance from the Random Forest model.
  • Interpret the top features through a mechanical engineering lens.
  • Summarize your findings with a final results cell that quantifies business impact.
Extract and visualize feature importance

Random Forest models track how much each feature contributes to correct predictions. The .feature_importances_ attribute gives you a score for each input feature, where higher scores mean the feature played a larger role in splitting decisions across all trees.

  • In predictive_maintenance.py, scroll to the bottom of the file (below the confusion matrices code).
  • Add the Step 4 section header and feature importance extraction code by copying the following:
# ============================================================
# STEP 4: Feature Importance and Engineering Interpretation
# ============================================================

rf_model = results['Random Forest']['model']
importances = rf_model.feature_importances_
feature_importance_df = pd.DataFrame({
    'Feature': feature_columns,
    'Importance': importances
}).sort_values('Importance', ascending=True)

fig, ax = plt.subplots(figsize=(8, 6))
ax.barh(feature_importance_df['Feature'], feature_importance_df['Importance'], color='steelblue')
ax.set_xlabel('Feature Importance')
ax.set_title('Random Forest Feature Importance for Failure Prediction')
plt.tight_layout()
plt.savefig('feature_importance.png', dpi=100)
plt.show()

print("\n--- Feature Importance Ranking ---")
for _, row in feature_importance_df.sort_values('Importance', ascending=False).iterrows():
    print(f"  {row['Feature']}: {row['Importance']:.4f}")

What does this code do?

rf_model = results['Random Forest']['model'] pulls the trained Random Forest out of the results dictionary you built in the previous step.

.feature_importances_ returns an array of importance scores, one per feature. The code wraps these into a DataFrame sorted in ascending order so the horizontal bar chart displays the most important features at the top.

The final loop prints each feature and its numeric score in descending order, giving you a clean text summary alongside the chart.

  • Save predictive_maintenance.py and run it.
  • You should see a horizontal bar chart showing each feature ranked by importance, saved as feature_importance.png. Below it, the terminal prints the numeric importance score for every feature.

Not seeing the chart?

Make sure you saved the file before running. If you see a KeyError on 'Random Forest', confirm that the Step 3 training loop ran successfully and populated the results dictionary.

help me debug a KeyError when accessing results['Random Forest']

Interpret results through a mechanical engineering lens

The chart shows you which features the model relies on most. Your engineering background lets you explain why these features matter physically.

  • Look at the feature importance ranking in your terminal output.
  • Identify which engineered features (like Wear_torque_interaction, Power_W, and Temp_difference) appear in the top positions compared to the raw sensor readings.

Why do these features matter physically?

Each top feature maps directly to a real failure mechanism in rotating equipment:

  • Wear_torque_interaction captures overstrain failure. When a worn tool is forced to cut under high torque, the combined stress exceeds the material threshold.
  • Power_W captures power failure. Equipment fails when power falls outside the 3500-9000W operating range.
  • Temp_difference captures heat dissipation failure. When the gap between process and air temperature drops below 8.6K with low rotational speed, cooling is insufficient.

In practice, an engineer would set alerts on these thresholds so maintenance teams can intervene before a failure occurs.

Why this matters for your portfolio

Hiring managers look for candidates who can bridge the gap between ML outputs and domain expertise. Anyone can run a Random Forest. Explaining that high Wear_torque_interaction importance aligns with known overstrain physics shows you understand the system you are modeling.

Summarize your findings with a final results cell

The last piece of any ML project is a clear summary that quantifies the model's real-world impact. You will calculate how many equipment failures the model catches versus misses in the test set.

  • Add the final summary code below the feature importance ranking loop in predictive_maintenance.py:
# Final summary
print("\n" + "="*50)
print("FINAL RESULTS SUMMARY")
print("="*50)
best_model_name = 'Random Forest'
best_predictions = results[best_model_name]['predictions']
from sklearn.metrics import recall_score, precision_score, f1_score
print(f"\nBest Model: {best_model_name}")
print(f"Failure Recall: {recall_score(y_test, best_predictions):.3f}")
print(f"Failure Precision: {precision_score(y_test, best_predictions):.3f}")
print(f"Failure F1-Score: {f1_score(y_test, best_predictions):.3f}")
print(f"\nBusiness Impact:")
total_failures = y_test.sum()
caught_failures = ((y_test == 1) & (best_predictions == 1)).sum()
print(f"  Total failures in test set: {total_failures}")
print(f"  Failures caught by model: {caught_failures}")
print(f"  Failures missed: {total_failures - caught_failures}")

What does this code do?

recall_score measures what proportion of actual failures the model correctly identified. precision_score measures what proportion of the model's failure predictions were correct. f1_score is the harmonic mean of both.

The business impact section counts the raw numbers. total_failures is the true count of failures in the test set. caught_failures counts how many of those the model correctly flagged. The difference tells you how many failures would slip through undetected.

  • Save predictive_maintenance.py and run it again.
  • Scroll to the bottom of the terminal output. You should see a FINAL RESULTS SUMMARY block showing the best model name, recall, precision, and F1 scores, followed by the count of failures caught and missed.

Seeing an ImportError for recall_score?

The from sklearn.metrics import recall_score, precision_score, f1_score line imports these metrics inline. If you see an import error, make sure scikit-learn is installed in your environment. Run pip install scikit-learn if needed.

help me fix an ImportError for sklearn.metrics

Why recall matters more than accuracy here

With only ~3.4% of data points being failures, a model that predicts "no failure" every single time would score 96.6% accuracy. That metric is misleading for imbalanced data.

Recall for the failure class tells you what percentage of real failures the model catches. In predictive maintenance, a missed failure means unplanned downtime. A false alarm just means an extra inspection. The cost of missing failures is far higher than the cost of extra inspections.

✔️ Awesome, I've got everything!

Great. Double check you have saved predictive_maintenance.py with all your Step 4 code included.

ⓧ I'd like to double check the full code

Here is the complete predictive_maintenance.py file from start to finish. Compare yours against this:

# Predictive Maintenance: Machine Failure Prediction
# Using the AI4I 2020 Predictive Maintenance Dataset (UCI ML Repository)

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
from sklearn.preprocessing import LabelEncoder
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 1: Load the Dataset
# ============================================================

dataset = fetch_ucirepo(id=601)
X = dataset.data.features
y = dataset.data.targets

# Combine features and targets for exploration
df = pd.concat([X, y], axis=1)

print("Dataset Shape:", df.shape)
print("\nColumn Names:")
print(df.columns.tolist())
print("\nFirst 5 rows:")
print(df.head())
print("\nData Types:")
print(df.dtypes)
print("\nBasic Statistics:")
print(df.describe())

# Check class distribution
print("\n--- Class Distribution ---")
print(df['Machine failure'].value_counts())
print(f"\nFailure Rate: {df['Machine failure'].mean() * 100:.2f}%")

# Visualize class imbalance
fig, ax = plt.subplots(figsize=(6, 4))
df['Machine failure'].value_counts().plot(kind='bar', ax=ax, color=['steelblue', 'coral'])
ax.set_title('Machine Failure Distribution')
ax.set_xlabel('Machine Failure (0 = No, 1 = Yes)')
ax.set_ylabel('Count')
ax.set_xticklabels(['No Failure', 'Failure'], rotation=0)
plt.tight_layout()
plt.savefig('class_distribution.png', dpi=100)
plt.show()

# ============================================================
# STEP 2: Exploratory Data Analysis and Feature Engineering
# ============================================================

# Feature distributions by failure status
numerical_features = ['Air temperature [K]', 'Process temperature [K]',
                      'Rotational speed [rpm]', 'Torque [Nm]', 'Tool wear [min]']

fig, axes = plt.subplots(2, 3, figsize=(14, 8))
axes = axes.flatten()
for i, feature in enumerate(numerical_features):
    ax = axes[i]
    df[df['Machine failure'] == 0][feature].hist(ax=ax, bins=30, alpha=0.5, label='No Failure', color='steelblue')
    df[df['Machine failure'] == 1][feature].hist(ax=ax, bins=30, alpha=0.5, label='Failure', color='coral')
    ax.set_title(feature)
    ax.legend()
axes[5].set_visible(False)
plt.tight_layout()
plt.savefig('feature_distributions.png', dpi=100)
plt.show()

# Correlation heatmap
fig, ax = plt.subplots(figsize=(8, 6))
correlation_matrix = df[numerical_features + ['Machine failure']].corr()
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0, ax=ax)
ax.set_title('Feature Correlation Heatmap')
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=100)
plt.show()

# Feature Engineering using mechanical engineering knowledge
df['Temp_difference'] = df['Process temperature [K]'] - df['Air temperature [K]']
df['Power_W'] = df['Torque [Nm]'] * df['Rotational speed [rpm]'] * (2 * np.pi / 60)
df['Wear_torque_interaction'] = df['Tool wear [min]'] * df['Torque [Nm]']

print("\n--- Engineered Features (first 5 rows) ---")
print(df[['Temp_difference', 'Power_W', 'Wear_torque_interaction']].head())

# Prepare features for modeling
le = LabelEncoder()
df['Type_encoded'] = le.fit_transform(df['Type'])

feature_columns = numerical_features + ['Type_encoded', 'Temp_difference', 'Power_W', 'Wear_torque_interaction']
X_model = df[feature_columns]
y_model = df['Machine failure']

# Train-test split with stratification
X_train, X_test, y_train, y_test = train_test_split(
    X_model, y_model, test_size=0.2, random_state=42, stratify=y_model
)

print(f"\nTraining set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")
print(f"Training failure rate: {y_train.mean() * 100:.2f}%")
print(f"Test failure rate: {y_test.mean() * 100:.2f}%")

# ============================================================
# STEP 3: Train and Compare Classification Models
# ============================================================

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

results = {}

for name, model in models.items():
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    results[name] = {
        'model': model,
        'predictions': y_pred
    }
    print(f"\n{'='*50}")
    print(f"Model: {name}")
    print('='*50)
    print(classification_report(y_test, y_pred, target_names=['No Failure', 'Failure']))

# Confusion matrices side by side
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
for i, (name, result) in enumerate(results.items()):
    ConfusionMatrixDisplay.from_predictions(
        y_test, result['predictions'],
        display_labels=['No Failure', 'Failure'],
        cmap='Blues',
        ax=axes[i]
    )
    axes[i].set_title(name)
plt.tight_layout()
plt.savefig('confusion_matrices.png', dpi=100)
plt.show()

# ============================================================
# STEP 4: Feature Importance and Engineering Interpretation
# ============================================================

rf_model = results['Random Forest']['model']
importances = rf_model.feature_importances_
feature_importance_df = pd.DataFrame({
    'Feature': feature_columns,
    'Importance': importances
}).sort_values('Importance', ascending=True)

fig, ax = plt.subplots(figsize=(8, 6))
ax.barh(feature_importance_df['Feature'], feature_importance_df['Importance'], color='steelblue')
ax.set_xlabel('Feature Importance')
ax.set_title('Random Forest Feature Importance for Failure Prediction')
plt.tight_layout()
plt.savefig('feature_importance.png', dpi=100)
plt.show()

print("\n--- Feature Importance Ranking ---")
for _, row in feature_importance_df.sort_values('Importance', ascending=False).iterrows():
    print(f"  {row['Feature']}: {row['Importance']:.4f}")

# Final summary
print("\n" + "="*50)
print("FINAL RESULTS SUMMARY")
print("="*50)
best_model_name = 'Random Forest'
best_predictions = results[best_model_name]['predictions']
from sklearn.metrics import recall_score, precision_score, f1_score
print(f"\nBest Model: {best_model_name}")
print(f"Failure Recall: {recall_score(y_test, best_predictions):.3f}")
print(f"Failure Precision: {precision_score(y_test, best_predictions):.3f}")
print(f"Failure F1-Score: {f1_score(y_test, best_predictions):.3f}")
print(f"\nBusiness Impact:")
total_failures = y_test.sum()
caught_failures = ((y_test == 1) & (best_predictions == 1)).sum()
print(f"  Total failures in test set: {total_failures}")
print(f"  Failures caught by model: {caught_failures}")
print(f"  Failures missed: {total_failures - caught_failures}")

Secret mission

Multi-Class Failure Mode Classification

Your binary classifier predicts IF a machine will fail. But in a real plant, the maintenance crew needs to know WHAT TYPE of failure is coming so they can bring the right tools and parts. Can you build a classifier that distinguishes between heat dissipation, power, overstrain, tool wear, and random failures?

Clean Up Your Resources

Clean Up Your Resources

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

Resources you used:

  • predictive_maintenance.py (main project script).
  • predictive_maintenance_multiclass.py (secret mission script).
  • Generated PNG images (class_distribution.png, feature_distributions.png, correlation_heatmap.png, confusion_matrices.png, feature_importance.png, multiclass_confusion_matrix.png, multiclass_feature_importance.png).
  • ucimlrepo Python package (installed via pip).

Keep everything running

No action needed. Choose this if you plan to revisit your models, tweak hyperparameters, or extend the analysis further.

  • Your Python files and generated images stay in place.
  • The ucimlrepo package remains installed in your Python environment.

Pause - I'll come back to this later

Close VS Code to free up memory. Your files and installed packages stay on disk so you can pick up where you left off.

  • Close VS Code.
  • Your predictive_maintenance.py, predictive_maintenance_multiclass.py, and all generated PNG images remain saved on your machine.
  • When you return, reopen the files in VS Code and run them again.

Delete - I don't want to use this again

Remove all project files and uninstall the package you added during this project.

  • Delete the project files and generated images by running this command from the folder where your scripts are saved:

macOS

rm predictive_maintenance.py predictive_maintenance_multiclass.py
rm class_distribution.png feature_distributions.png correlation_heatmap.png confusion_matrices.png feature_importance.png multiclass_confusion_matrix.png multiclass_feature_importance.png

Windows

del predictive_maintenance.py, predictive_maintenance_multiclass.py
del class_distribution.png, feature_distributions.png, correlation_heatmap.png, confusion_matrices.png, feature_importance.png, multiclass_confusion_matrix.png, multiclass_feature_importance.png
  • Uninstall the ucimlrepo package by running:
pip uninstall ucimlrepo
  • Confirm the uninstall when prompted by typing y.

Nice Work!

Nice Work!

Nice work! You just built a predictive maintenance classifier that identifies equipment at risk of failure using real industrial sensor data.

You've learned how to:

  • Train and compare classification models (Logistic Regression, Decision Tree, Random Forest) on imbalanced industrial data, using class weights and evaluating with precision, recall, and F1 instead of misleading accuracy.
  • Engineer domain-specific features (temperature difference, power, wear-torque interaction) from raw sensor readings and connect Random Forest feature importance back to mechanical engineering failure mechanisms.
  • Explore and visualize a real UCI ML Repository dataset, identifying class imbalance and feature distributions that signal impending equipment failure.
  • Secret Mission: Extended to multi-class failure mode prediction, diagnosing not just IF but WHAT type of failure will occur (TWF, HDF, PWF, OSF, RNF).

Ready to quiz yourself?