machine learning
data visualization
scikit-learn
decision boundary
classification

How can I visualize border/decision function of two classes using scikit-learn

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Visualizing a binary classifier’s decision boundary is one of the quickest ways to understand what the model has learned. In scikit-learn, the standard pattern is to train on two features, build a mesh grid across that 2D space, score the grid, and then plot either predicted classes, probabilities, or the raw decision function.

Train a Model on Two Features

Decision-boundary plots are easiest when the training data has exactly two input features. A small synthetic dataset keeps the example focused on the plotting pattern.

python
1import numpy as np
2import matplotlib.pyplot as plt
3from sklearn.datasets import make_classification
4from sklearn.pipeline import make_pipeline
5from sklearn.preprocessing import StandardScaler
6from sklearn.svm import SVC
7
8X, y = make_classification(
9    n_samples=300,
10    n_features=2,
11    n_redundant=0,
12    n_informative=2,
13    n_clusters_per_class=1,
14    class_sep=1.2,
15    random_state=42,
16)
17
18model = make_pipeline(
19    StandardScaler(),
20    SVC(kernel="rbf", probability=True),
21)
22
23model.fit(X, y)

Using a pipeline is important because the grid points you plot will go through the same preprocessing as the training data.

Build the Mesh Grid

The next step is to create a dense grid of coordinates covering the region where your data lives.

python
1xx, yy = np.meshgrid(
2    np.linspace(X[:, 0].min() - 1, X[:, 0].max() + 1, 300),
3    np.linspace(X[:, 1].min() - 1, X[:, 1].max() + 1, 300),
4)
5
6grid = np.c_[xx.ravel(), yy.ravel()]

Each row in grid is a synthetic point in feature space. Once the classifier predicts on those points, you can reshape the result back to the mesh and plot it.

Plot Predicted Class Regions

The simplest picture is a region plot showing which class the model predicts at each location.

python
1predicted = model.predict(grid).reshape(xx.shape)
2
3plt.figure(figsize=(8, 6))
4plt.contourf(xx, yy, predicted, alpha=0.25, cmap="coolwarm")
5plt.scatter(X[:, 0], X[:, 1], c=y, cmap="coolwarm", edgecolor="k")
6plt.xlabel("feature 1")
7plt.ylabel("feature 2")
8plt.title("Predicted class regions")
9plt.show()

This is a good first view, but it hides how strongly the classifier prefers one class over the other near the boundary.

Plot the Decision Function

For models that expose decision_function, plotting the score surface is usually more informative. The contour where the score equals zero is the decision boundary.

python
1scores = model.decision_function(grid).reshape(xx.shape)
2
3plt.figure(figsize=(8, 6))
4plt.contourf(xx, yy, scores, levels=30, cmap="coolwarm", alpha=0.35)
5plt.contour(xx, yy, scores, levels=[0], colors="black", linewidths=2)
6plt.scatter(X[:, 0], X[:, 1], c=y, cmap="coolwarm", edgecolor="k")
7plt.xlabel("feature 1")
8plt.ylabel("feature 2")
9plt.title("Decision function with boundary at score 0")
10plt.colorbar(label="decision score")
11plt.show()

This view tells you more than class labels alone. It shows how sharply the model separates the two classes and where the margin is narrow.

Plot Probabilities When Available

If the estimator supports predict_proba, you can visualize the probability surface instead of raw decision scores.

python
1probabilities = model.predict_proba(grid)[:, 1].reshape(xx.shape)
2
3plt.figure(figsize=(8, 6))
4plt.contourf(xx, yy, probabilities, levels=30, cmap="viridis", alpha=0.4)
5plt.contour(xx, yy, probabilities, levels=[0.5], colors="white", linewidths=2)
6plt.scatter(X[:, 0], X[:, 1], c=y, cmap="coolwarm", edgecolor="k")
7plt.xlabel("feature 1")
8plt.ylabel("feature 2")
9plt.title("Probability surface for class 1")
10plt.colorbar(label="P(class 1)")
11plt.show()

This is useful when your discussion is about confidence or thresholding rather than just classification regions.

What Changes with Different Models

The plotting workflow stays almost the same across classifiers, but the geometry changes:

  • linear models produce straight boundaries in feature space,
  • kernel methods can produce curved shapes,
  • decision trees create axis-aligned regions,
  • ensemble models can create more jagged surfaces.

That means the plot is not just decoration. It is a direct picture of the model family you chose.

What If the Original Data Has More Than Two Features

You cannot draw the true boundary of a high-dimensional classifier on a flat 2D chart. In practice, people use one of these strategies:

  • train a model on only two selected features,
  • project the data to two dimensions for visualization,
  • hold the other features constant and plot a slice.

All of these are approximations. The plot can still be useful, but it is important to say that it is a projection or slice rather than the full decision geometry.

Common Pitfalls

  • Forgetting to apply the same preprocessing to the mesh grid that the model used during training.
  • Using a very coarse grid and getting a blocky, misleading boundary.
  • Plotting hard class labels only and missing the confidence structure around the separating line.
  • Treating a 2D projection as if it were the full boundary of a high-dimensional model.
  • Assuming predict_proba and decision_function mean the same thing for every estimator.

Summary

  • Train on two features or create a deliberate 2D slice before plotting a decision boundary.
  • Build a dense mesh grid, score it with the model, and reshape the output for plotting.
  • Use predict for class regions, decision_function for margin-like scores, and predict_proba for probability surfaces.
  • Keep preprocessing inside a pipeline so the plotted boundary matches the trained model.
  • Be explicit when the visualization is only a slice or projection of a higher-dimensional problem.

Course illustration
Course illustration

All Rights Reserved.