Chapter 36: Categorical Embeddings

import numpy as np
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import KFold

from kaggle_companion.activities._common import clean

SEED = 36
REPLICATES = 6                       # independent worlds per control value; results are means over them
SEEN, NEW = 1500, 1500               # categories that can appear in training, and a pool of new ones for test
N_TRAIN, N_TEST = 10_000, 6_000
M = 10.0                             # smoothing strength of the chapter's smoothed mean


def make_world(rng):
    """Zipf popularity for the training categories; a category effect plus a rarity effect (rare categories run hotter)."""
    p = 1 / np.arange(1, SEEN + 1) ** 0.9
    p /= p.sum()
    log_p = np.concatenate([np.log(p), np.full(NEW, np.log(p[-1]))])        # new categories are tail-rare
    rarity = -0.45 * (log_p - np.log(p).mean()) / np.log(p).std()
    return p, rng.normal(0, 0.7, SEEN + NEW) + rarity


def draw_rows(categories, effect, rng):
    x = rng.normal(size=(len(categories), 2))
    logit = -1.2 + 0.7 * x[:, 0] - 0.3 * x[:, 1] + effect[categories]
    return x, (rng.random(len(categories)) < 1 / (1 + np.exp(-logit))).astype(int)


def smoothed_means(cat, y, prior):
    """e_c = (s_c + m*mu) / (n_c + m). A category with n_c = 0 gets exactly the prior mu."""
    s = np.bincount(cat, weights=y, minlength=SEEN + NEW)
    n = np.bincount(cat, minlength=SEEN + NEW)
    return (s + M * prior) / (n + M)


def one_world(unseen_share, seed):
    rng = np.random.default_rng(seed)
    p, effect = make_world(rng)
    cat_tr = rng.choice(SEEN, N_TRAIN, p=p)
    x_tr, y_tr = draw_rows(cat_tr, effect, rng)
    n_new = int(round(unseen_share * N_TEST))
    cat_te = np.concatenate([rng.choice(SEEN, N_TEST - n_new, p=p), SEEN + rng.integers(0, NEW, n_new)])
    x_te, y_te = draw_rows(cat_te, effect, rng)

    known = np.zeros(SEEN + NEW, bool)
    known[cat_tr] = True
    unseen = ~known[cat_te]                                  # unseen means: absent from the training dictionary
    counts = np.bincount(cat_tr, minlength=SEEN + NEW)       # frequency encoding from training rows only; unseen -> 0
    codes = -np.ones(SEEN + NEW, int)                        # label codes in order of first appearance; unseen -> -1
    first = np.unique(cat_tr, return_index=True)
    order = first[0][np.argsort(first[1])]                   # in a shuffled table, first appearance follows popularity
    codes[order] = np.arange(known.sum())
    random_codes = -np.ones(SEEN + NEW, int)                 # the same categories numbered in random order; unseen -> -1
    random_codes[np.random.default_rng([seed, 1]).permutation(order)] = np.arange(known.sum())

    prior = y_tr.mean()
    te_train = np.zeros(N_TRAIN)                             # nested: each training row is encoded from the other folds
    for fit, out in KFold(5, shuffle=True, random_state=0).split(cat_tr):
        te_train[out] = smoothed_means(cat_tr[fit], y_tr[fit], y_tr[fit].mean())[cat_tr[out]]
    full_map = smoothed_means(cat_tr, y_tr, prior)
    te_known = np.where(known[cat_te], full_map[cat_te], prior)    # the documented rule: unknown -> training prior
    te_zero = np.where(known[cat_te], full_map[cat_te], 0.0)       # the shortcut: unknown -> 0

    arms = {                                                 # (training features, test features)
        "none": (x_tr, x_te),
        "label": (np.c_[x_tr, codes[cat_tr]], np.c_[x_te, codes[cat_te]]),
        "label_random": (np.c_[x_tr, random_codes[cat_tr]], np.c_[x_te, random_codes[cat_te]]),
        "frequency": (np.c_[x_tr, counts[cat_tr]], np.c_[x_te, counts[cat_te]]),
        "target_prior": (np.c_[x_tr, te_train, counts[cat_tr]], np.c_[x_te, te_known, counts[cat_te]]),
        "target_zero": (np.c_[x_tr, te_train, counts[cat_tr]], np.c_[x_te, te_zero, counts[cat_te]]),
    }
    scores = {}
    for name, (a, b) in arms.items():
        model = HistGradientBoostingClassifier(max_iter=80, max_leaf_nodes=15, random_state=0).fit(a, y_tr)
        pred = model.predict_proba(b)[:, 1]
        scores[name] = [roc_auc_score(y_te, pred), roc_auc_score(y_te[~unseen], pred[~unseen]),
                        roc_auc_score(y_te[unseen], pred[unseen])]
    return scores, float(unseen.mean())


def run(unseen_share):
    worlds = [one_world(unseen_share, SEED * 100 + r) for r in range(REPLICATES)]
    names = list(worlds[0][0])
    mean = {k: np.mean([w[0][k] for w in worlds], axis=0) for k in names}
    below = {k: float(np.mean([w[0][k][0] < w[0]["none"][0] for w in worlds])) for k in names}
    return clean({
        "unseen_share": unseen_share,
        "measured_unseen_share": float(np.mean([w[1] for w in worlds])),
        "auc": {k: mean[k][0] for k in names}, "auc_seen": {k: mean[k][1] for k in names},
        "auc_unseen": {k: mean[k][2] for k in names}, "below_no_feature": below,
        "replicates": REPLICATES, "test_rows": N_TEST,
    })

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

VALUES = [0, 0.1, 0.25, 0.5]
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.ch36 import draw, NCOLS, HEIGHT

DEFAULT = 0.25
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()