Chapter 28: Nearest-Neighbor Features

import numpy as np
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import KFold
from sklearn.neighbors import NearestNeighbors
from sklearn.preprocessing import StandardScaler

from kaggle_companion.activities._common import clean

SEED = 28
REPLICATES = 6          # independent constructed datasets; every estimate is their mean
DIMS, BLOBS = 5, 40
N_TRAIN, N_NEW = 1500, 4000


def generate(rng):
    """Forty tight clusters in five dimensions, each with its own event rate (0.12 or 0.88).
    The label is a property of the neighbourhood, so neighbours' labels carry real signal."""
    centers = rng.uniform(-3, 3, (BLOBS, DIMS))
    rate = np.where(rng.random(BLOBS) < 0.5, 0.12, 0.88)

    def rows(n):
        b = rng.integers(0, BLOBS, n)
        return centers[b] + rng.normal(0, 0.55, (n, DIMS)), (rng.random(n) < rate[b]).astype(int)

    return rows(N_TRAIN), rows(N_NEW)


def neighbour_mean(X_db, y_db, X_query, k, exclude_self):
    """Mean label of the k nearest database rows. The scaler and the database come from X_db only.
    With exclude_self the query rows ARE the database rows, so each one's first neighbour is itself and is dropped."""
    scaler = StandardScaler().fit(X_db)
    nn = NearestNeighbors(n_neighbors=k + int(exclude_self)).fit(scaler.transform(X_db))
    ids = nn.kneighbors(scaler.transform(X_query))[1]
    return y_db[ids[:, int(exclude_self):]].mean(axis=1)


def auc(X_fit, y_fit, X_eval, y_eval):
    model = HistGradientBoostingClassifier(max_iter=60, max_leaf_nodes=8, random_state=0).fit(X_fit, y_fit)
    return roc_auc_score(y_eval, model.predict_proba(X_eval)[:, 1])


def with_feature(X, f):
    return np.column_stack([X, f])


def one_replicate(k, seed):
    (X, y), (X_new, y_new) = generate(np.random.default_rng(seed))
    outer = list(KFold(5, shuffle=True, random_state=1).split(X))
    cv = lambda col: np.mean([auc(with_feature(X[a], col[a]), y[a], with_feature(X[b], col[b]), y[b]) for a, b in outer])

    baseline = (np.mean([auc(X[a], y[a], X[b], y[b]) for a, b in outer]), auc(X, y, X_new, y_new))
    query_new = neighbour_mean(X, y, X_new, k, exclude_self=False)    # new rows are queried against the training rows

    # 1. Self included: the training column is built from the same rows it describes, each row's own label among them.
    col_self = neighbour_mean(X, y, X, k, exclude_self=False)
    # 2. Self excluded once on the full training set, then ordinary cross-validation on that column.
    col_loo = neighbour_mean(X, y, X, k, exclude_self=True)
    # 3. Nested: inside each outer fold, exclude self among the outer-training rows and query the validation rows from them.
    nested = np.mean([auc(with_feature(X[a], neighbour_mean(X[a], y[a], X[a], k, True)), y[a],
                          with_feature(X[b], neighbour_mean(X[a], y[a], X[b], k, False)), y[b]) for a, b in outer])
    deployed_loo = auc(with_feature(X, col_loo), y, with_feature(X_new, query_new), y_new)
    return {"baseline": baseline,
            "self_included": (cv(col_self), auc(with_feature(X, col_self), y, with_feature(X_new, query_new), y_new)),
            "self_excluded": (cv(col_loo), deployed_loo),
            "nested": (nested, deployed_loo)}      # pipelines 2 and 3 deploy the same model; they differ in the estimate


def run(k):
    reps = [one_replicate(k, SEED * 100 + i) for i in range(REPLICATES)]
    out = {"neighbors": k, "replicates": REPLICATES, "training_rows": N_TRAIN, "new_rows": N_NEW}
    for name in ("baseline", "self_included", "self_excluded", "nested"):
        out[name] = {"cv": float(np.mean([r[name][0] for r in reps])), "new_rows": float(np.mean([r[name][1] for r in reps]))}
    out["per_replicate_cv"] = {n: [r[n][0] for r in reps] for n in ("self_included", "self_excluded", "nested")}
    return clean(out)

import pandas as pd
from kaggle_companion.activities.ch28 import explain

VALUES = [1, 5, 20, 50]
results = {value: run(value) for value in VALUES}
table = pd.DataFrame({str(v): explain(results[v], v)['metrics'] for v in VALUES})
print(table.to_string())

import matplotlib.pyplot as plt
from kaggle_companion.activities.ch28 import draw, NCOLS, HEIGHT

DEFAULT = 5
fig, axes = plt.subplots(1, NCOLS, figsize=(6.6 if NCOLS == 1 else 10.4, HEIGHT), layout='constrained')
draw(axes, results[DEFAULT], DEFAULT)
plt.show()