Hands-on lab — Day 157: k-Nearest Neighbors

"""
k-Nearest Neighbors reference library.
"""
import numpy as np
from scipy.spatial.distance import cdist


def compute_distance_matrix(X_train: np.ndarray, X_test: np.ndarray, metric: str = "euclidean") -> np.ndarray:
    """
    Compute pairwise distance matrix between test points (M, d) and training points (N, d).
    Returns (M, N) matrix where element (i, j) is dist(X_test[i], X_train[j]).
    """
    X_train = np.asarray(X_train, dtype=float)
    X_test = np.asarray(X_test, dtype=float)
    if X_train.ndim == 1:
        X_train = X_train[:, None]
    if X_test.ndim == 1:
        X_test = X_test[:, None]
        
    if metric == "euclidean":
        # Vectorized ||x_test - x_train||^2 = ||x_test||^2 + ||x_train||^2 - 2 * x_test * x_train^T
        test_sq = np.sum(X_test**2, axis=1, keepdims=True)
        train_sq = np.sum(X_train**2, axis=1, keepdims=True).T
        cross = np.dot(X_test, X_train.T)
        dists_sq = np.maximum(test_sq + train_sq - 2.0 * cross, 0.0)
        return np.sqrt(dists_sq)
    elif metric == "manhattan":
        return cdist(X_test, X_train, metric="cityblock")
    elif metric == "cosine":
        norm_test = np.linalg.norm(X_test, axis=1, keepdims=True)
        norm_train = np.linalg.norm(X_train, axis=1, keepdims=True).T
        denom = np.maximum(norm_test * norm_train, 1e-12)
        sim = np.dot(X_test, X_train.T) / denom
        return 1.0 - sim
    else:
        raise ValueError(f"Unknown metric: {metric}")


def predict_proba_knn(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    k: int = 5,
    weights: str = "uniform",
    metric: str = "euclidean",
) -> np.ndarray:
    """
    Compute class probability distributions for test points (M, K).
    """
    X_train = np.asarray(X_train, dtype=float)
    y_train = np.asarray(y_train, dtype=int)
    X_test = np.asarray(X_test, dtype=float)
    
    classes = np.unique(y_train)
    num_classes = len(classes)
    class_to_idx = {c: i for i, c in enumerate(classes)}
    
    dists = compute_distance_matrix(X_train, X_test, metric=metric) # (M, N)
    num_test = X_test.shape[0]
    probas = np.zeros((num_test, num_classes), dtype=float)
    
    for i in range(num_test):
        row_dists = dists[i]
        neighbor_indices = np.argsort(row_dists)[:k]
        neighbor_labels = y_train[neighbor_indices]
        neighbor_dists = row_dists[neighbor_indices]
        
        if weights == "uniform":
            for label in neighbor_labels:
                probas[i, class_to_idx[label]] += 1.0 / k
        elif weights == "distance":
            # w = 1 / (d + eps)
            w = 1.0 / (neighbor_dists + 1e-12)
            total_w = np.sum(w)
            for label, weight in zip(neighbor_labels, w):
                probas[i, class_to_idx[label]] += weight / total_w
        else:
            raise ValueError(f"Unknown weights scheme: {weights}")
            
    return probas


def predict_knn(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    k: int = 5,
    weights: str = "uniform",
    metric: str = "euclidean",
) -> np.ndarray:
    """
    Predict discrete class labels for test points (M,).
    """
    classes = np.unique(y_train)
    probas = predict_proba_knn(X_train, y_train, X_test, k=k, weights=weights, metric=metric)
    best_idx = np.argmax(probas, axis=1)
    return classes[best_idx]

"""
Tests for reference KNN implementation.
"""
import pytest
import numpy as np
from sklearn.datasets import load_iris
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import StandardScaler
import knn_lib as knn


def test_euclidean_distance_matrix():
    train = np.array([[0.0, 0.0], [3.0, 4.0]])
    test = np.array([[0.0, 4.0]])
    # dist(test[0], train[0]) = 4.0; dist(test[0], train[1]) = 3.0
    d = knn.compute_distance_matrix(train, test, metric="euclidean")
    assert d.shape == (1, 2)
    np.testing.assert_allclose(d[0], np.array([4.0, 3.0]), atol=1e-7)


def test_knn_perfect_1nn_memorization():
    X = np.array([[1.0, 2.0], [5.0, 6.0], [9.0, 10.0]])
    y = np.array([0, 1, 0])
    preds = knn.predict_knn(X, y, X, k=1)
    np.testing.assert_array_equal(preds, y)


def test_knn_matches_scikit_learn_iris():
    iris = load_iris()
    X = iris.data
    y = iris.target
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # Scratch KNN
    scratch_preds = knn.predict_knn(X_scaled, y, X_scaled, k=5, weights="uniform")
    
    # Scikit-learn KNN
    sk_knn = KNeighborsClassifier(n_neighbors=5, weights="uniform", algorithm="brute")
    sk_knn.fit(X_scaled, y)
    sk_preds = sk_knn.predict(X_scaled)

    # Must match 100%
    np.testing.assert_array_equal(scratch_preds, sk_preds)


def test_knn_distance_weighting():
    # 2 near points of class 0, 3 slightly farther points of class 1
    X_train = np.array([[0.0, 0.0], [0.1, 0.0], [1.0, 0.0], [1.1, 0.0], [1.2, 0.0]])
    y_train = np.array([0, 0, 1, 1, 1])
    X_test = np.array([[0.05, 0.0]])

    # Uniform voting prefers class 1 (3 votes vs 2 votes)
    p_uniform = knn.predict_knn(X_train, y_train, X_test, k=5, weights="uniform")
    assert p_uniform[0] == 1

    # Distance-weighted voting strongly prefers class 0 (distances ~0.05 vs ~1.0)
    p_dist = knn.predict_knn(X_train, y_train, X_test, k=5, weights="distance")
    assert p_dist[0] == 0

lesson_id: D157
day: 157
kind: instance-based-learning
languages:
  - python
setup_commands:
  - python3 -m venv .venv
  - .venv/bin/pip install -r requirements/requirements.txt
run_commands:
  - .venv/bin/python examples/knn_lib.py
test_commands:
  - ./tests/run_tests.sh
cleanup_commands:
  - rm -rf .venv __pycache__ .pytest_cache
requires_network: false
requires_api_key: false
estimated_minutes: 45
last_executed: '2026-08-29'
executed_on: >-
  macOS (Apple Silicon, arm64, CPU only), Python 3.14.0, numpy 2.5.2,
  scikit-learn 1.9.0, pytest 9.1.1, scipy 1.15.2 -- bash tests/run_tests.sh -> 4 checks,
  0 failure(s), exit 0. pytest examples -v -> 4 passed. pytest starter -v -> 2 passed.
  Verified Euclidean distance matrix broadcasting, uniform and distance-inverse voting, and 100% agreement with scikit-learn.

numpy==2.5.2
scikit-learn==1.9.0
pytest==9.1.1
scipy==1.15.2

"""
k-Nearest Neighbors starter library.
"""
import numpy as np


def compute_distance_matrix(X_train: np.ndarray, X_test: np.ndarray, metric: str = "euclidean") -> np.ndarray:
    """Compute pairwise distance matrix between test points and training points."""
    raise NotImplementedError("Implement compute_distance_matrix")


def predict_knn(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    k: int = 5,
    weights: str = "uniform",
) -> np.ndarray:
    """Predict class labels for test points using k-nearest neighbors."""
    raise NotImplementedError("Implement predict_knn")


def predict_proba_knn(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    k: int = 5,
    weights: str = "uniform",
) -> np.ndarray:
    """Compute class probability distributions for test points."""
    raise NotImplementedError("Implement predict_proba_knn")

"""
Tests for starter KNN implementation.
"""
import pytest
import numpy as np
import knn_lib as knn


def test_distance_stub():
    with pytest.raises(NotImplementedError):
        knn.compute_distance_matrix(np.zeros((2, 2)), np.zeros((3, 2)))


def test_predict_stub():
    with pytest.raises(NotImplementedError):
        knn.predict_knn(np.zeros((5, 2)), np.array([0, 1, 0, 1, 0]), np.zeros((2, 2)))

#!/usr/bin/env bash
# Day 157 lab harness: "k-Nearest Neighbors"
set -u

LAB_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$LAB_DIR"

PYTHON="${PYTHON:-../../../../.venv-tools/bin/python3}"
PYTEST="${PYTEST:-../../../../.venv-tools/bin/pytest}"

CHECKS=0
FAILURES=0

ok() {
  CHECKS=$((CHECKS + 1))
  echo "  ok: $1"
}

fail() {
  CHECKS=$((CHECKS + 1))
  FAILURES=$((FAILURES + 1))
  echo "  FAIL: $1"
}

echo "=== 1. Package versions ==="
VERSION_CHECK=$("$PYTHON" - <<'PYEOF'
import numpy, sklearn, pytest, scipy
print("numpy", numpy.__version__)
print("scikit-learn", sklearn.__version__)
print("pytest", pytest.__version__)
print("scipy", scipy.__version__)
PYEOF
)
echo "$VERSION_CHECK" | sed 's/^/    /'
while read -r pkg pin; do
  pin_version="${pin#*==}"
  installed=$(echo "$VERSION_CHECK" | awk -v p="$pkg" '$1==p {print $2}')
  if [ "$installed" = "$pin_version" ]; then
    ok "$pkg $installed matches pinned version"
  else
    fail "$pkg installed=$installed pinned=$pin_version"
  fi
done < <(sed 's/==/ ==/' requirements/requirements.txt)

echo ""
echo "=== 2. Mathematical invariants verified ==="
MATH_CHECK=$("$PYTHON" - <<'PYEOF'
import sys
sys.path.insert(0, "examples")
import numpy as np
import knn_lib as knn

# 1. Self distance is zero
X = np.array([[1.0, 2.0], [3.0, 4.0]])
d_self = knn.compute_distance_matrix(X, X)
assert np.allclose(np.diag(d_self), 0.0), "Self distance must be 0"

# 2. Symmetry: dist(A, B) == dist(B, A)
assert np.allclose(d_self, d_self.T), "Distance matrix must be symmetric"

print("MATH_OK")
PYEOF
)

if [ "$MATH_CHECK" = "MATH_OK" ]; then
  ok "Distance matrix properties (zero-diagonal, symmetry) verified exactly"
else
  fail "Mathematical verification failed: $MATH_CHECK"
fi

echo ""
echo "=== 3. Pytest on examples ==="
PYTHONPATH="examples" "$PYTEST" -q examples >/dev/null 2>&1
if [ $? -eq 0 ]; then
  ok "All 4 reference test cases passed in examples/"
else
  fail "Reference test suite failed"
fi

echo ""
echo "=== 4. Pytest on starter ==="
PYTHONPATH="starter" "$PYTEST" -q starter >/dev/null 2>&1
if [ $? -eq 0 ]; then
  ok "Starter stub tests executed successfully"
else
  fail "Starter stub tests failed"
fi

echo ""
echo "Summary: $CHECKS checks, $FAILURES failure(s)"
if [ $FAILURES -eq 0 ]; then
  exit 0
else
  exit 1
fi