Concepts 🚀 Project 🧮 Labs

Decision Tree visualized — the splits + overfitting (iris)

SEE the actual tree structure with its split rules. SEE how depth 3 vs depth 10 changes the decision regions. The most interpretable model.

script.py v3.1

import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.model_selection import train_test_split

# ── Load iris ──
iris = load_iris()
X = iris.data[:, [2, 3]]  # petal length, petal width
y = iris.target
feature_names = ['petal length', 'petal width']

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

print(f'Dataset: iris (train={len(X_train)}, test={len(X_test)})')

# ════════════════════════════════════════════════════════════════════
#  Compare 3 tree depths: underfit → good → overfit
# ════════════════════════════════════════════════════════════════════
depths = [1, 3, 10]
fig, axes = plt.subplots(2, 3, figsize=(18, 11))

for i, depth in enumerate(depths):
    model = DecisionTreeClassifier(max_depth=depth, random_state=42)
    model.fit(X_train, y_train)

    train_acc = model.score(X_train, y_train)
    test_acc = model.score(X_test, y_test)

    # TOP ROW: decision regions
    ax = axes[0, i]
    x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5
    y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 200), np.linspace(y_min, y_max, 200))
    Z = model.predict(np.c_[xx.ravel(), yy.ravel()]).reshape(xx.shape)

    ax.contourf(xx, yy, Z, alpha=0.2, cmap=plt.cm.Set1)
    colors = ['#e74c3c', '#2ecc71', '#3498db']
    for j, species in enumerate(iris.target_names):
        mask = y == j
        ax.scatter(X[mask, 0], X[mask, 1], c=colors[j], label=species, s=30, edgecolors='black')
    ax.set_xlabel(feature_names[0])
    ax.set_ylabel(feature_names[1])
    fit_label = 'UNDERFIT' if depth == 1 else ('GOOD' if depth == 3 else 'OVERFIT')
    ax.set_title(f'Depth {depth} ({fit_label})\ntrain={train_acc:.3f}, test={test_acc:.3f}')
    ax.legend(fontsize=8)

    # BOTTOM ROW: the actual tree structure
    ax2 = axes[1, i]
    plot_tree(model, ax=ax2, feature_names=feature_names,
             class_names=list(iris.target_names), filled=True, fontsize=7,
             rounded=True, precision=2)
    ax2.set_title(f'Tree structure (depth {depth})', fontsize=10)

plt.suptitle('Decision Tree: depth controls overfitting', fontsize=14, y=1.01)
plt.tight_layout()

print(f'\nResults:')
for depth in depths:
    model = DecisionTreeClassifier(max_depth=depth, random_state=42)
    model.fit(X_train, y_train)
    print(f'  Depth {depth:2d}: train={model.score(X_train, y_train):.3f}, test={model.score(X_test, y_test):.3f}, '
          f'n_leaves={model.get_n_leaves()}')

print(f'\nHOW TO READ THE TREE:')
print(f'  Each box is a NODE. Top box = root.')
print(f'  "petal length <= 2.45" = the SPLIT RULE')
print(f'  gini = impurity (0 = pure, 0.5 = mixed)')
print(f'  samples = how many rows reached this node')
print(f'  value = [n_class0, n_class1, n_class2]')
print(f'  class = the majority class (the prediction)')
print(f'\n  LEFT = condition is TRUE. RIGHT = condition is FALSE.')
print(f'  Leaf nodes (no children) = final predictions.')
print(f'\nTHE MATHS:')
print(f'  At each node, find the split that MINIMIZES GINI impurity:')
print(f'  Gini = 1 - Σ(p_i)²  where p_i = probability of class i')
print(f'  Pure node (one class): gini = 0')
print(f'  50/50 split: gini = 0.5')
print(f'  The tree GREEDILY picks the split that reduces gini the most.')
print(f'\n  Depth 1 (UNDERFIT): single split — only separates setosa from rest')
print(f'  Depth 3 (GOOD): enough splits to capture the pattern')
print(f'  Depth 10 (OVERFIT): too many splits — memorizes training data')

Python output click Run to execute

Click "Run" to execute the code on the left.
The output will appear here.

🧠 How it works

Decision Tree — VISUALIZED: Each tree node asks a YES/NO question: 'petal length <= 2.45?' LEFT = yes, RIGHT = no. Follow the path to a LEAF = your prediction. GINI IMPURITY: measures how mixed the classes are at a node. Gini = 1 - Σ(p_i)². Pure node (one class) = 0. 50/50 = 0.5. The tree picks the split that reduces gini the MOST (greedy). DEPTH = overfitting dial: Depth 1 (UNDERFIT): 1 split — only separates setosa. Train + test both low. Depth 3 (GOOD): 3-4 splits — captures the real pattern. Train ≈ test. Depth 10 (OVERFIT): 20+ splits — memorizes training. Train=1.0, test drops. DECISION REGIONS: trees create RECTANGULAR regions (axis-aligned splits). This is why trees can't learn diagonal boundaries well — need many splits. WHY TREES ARE GREAT: (1) Interpretable — you can EXPLAIN any prediction by following the path. (2) No scaling needed — splits are threshold-based. (3) Handles categorical + numerical naturally. (4) Feature importance is built-in (which features are used near the root?). WHY TREES ARE BAD (alone): (1) High variance — small data change → totally different tree. (2) Overfit easily — no max_depth = memorizes everything. (3) Can't extrapolate — predictions bounded by training data range. SOLUTION: RandomForest (many trees + voting) fixes the variance problem.

✏️ Try these changes

Change max_depth to None (unlimited)
→ tree grows until all leaves are pure — extreme overfitting
Change the split criterion from 'gini' to 'entropy'
→ uses information gain instead of gini — usually similar results
Change min_samples_leaf from default (1) to 10
→ each leaf must have ≥10 samples — prevents tiny overfit leaves
Use all 4 features: X = iris.data
→ tree uses more features — deeper tree needed for same accuracy

← Previous lab

SVM visualized — hyperplane + margins + support vectors (iris)

Next lab →

Random Forest visualized — many trees + feature importance (wine)