Winner models, code, notebooks, data and results

2b9ed19 verified 4 days ago

23.3 kB

"""
The official evaluation protocol, plus the decision-rule and prior-shift tools
that every model in this project shares.
Protocol (organisers' Notebook 1, section 1.5):
  1. train a Stage-1 model on Madrid only;
  2. adapt it with n labelled Amsterdam pixels per class, n in {5,25,50,100,200};
  3. score macro F1 on the *rest* of Amsterdam;
  4. repeat over many random draws and report mean +- std for every n.
`modelCreate` is any callable with this signature:
    predict_proba_fn = model_create(
        X_src, y_src,            # Madrid, labelled
        X_tgt_unlabelled,        # all target features, NO labels
        X_shot, y_shot,          # the n-per-class labelled target sample
    )
and must return a function mapping X -> (n, 4) probabilities. Keeping it a
closure means every adaptation trick (recalibration, prototypes, continued
boosting, retrieval prompts) fits the same harness.
"""
from __future__ import annotations

import numpy as np
from sklearn.metrics import confusion_matrix, f1_score

CLASSES = np.array([1, 2, 3, 4])


# --------------------------------------------------------------------------- #
# Metrics
# --------------------------------------------------------------------------- #
def macro_f1(y_true, y_pred) -> float:
    return float(f1_score(y_true, y_pred, labels=CLASSES, average="macro", zero_division=0))


def per_class_f1(y_true, y_pred) -> np.ndarray:
    return f1_score(y_true, y_pred, labels=CLASSES, average=None, zero_division=0)


def confusion_proportions(y_true, y_pred) -> np.ndarray:
    cm = confusion_matrix(y_true, y_pred, labels=CLASSES).astype(float)
    return cm / np.clip(cm.sum(axis=1, keepdims=True), 1, None)


# --------------------------------------------------------------------------- #
# Decision rules.  argmax is NOT the macro-F1 optimum.
# --------------------------------------------------------------------------- #
def predict_weighted(proba: np.ndarray, weights: np.ndarray) -> np.ndarray:
    """argmax of w_k * p_k  ->  class label."""
    return CLASSES[np.argmax(proba * np.asarray(weights, dtype=float), axis=1)]


def macro_f1_weighted(y_true, y_pred, sample_weight=None) -> float:
    """
    Macro F1 computed from a *weighted* confusion matrix.
    This exists because the few-shot sample is drawn n-per-class: on it, a rare
    class looks far more precise than it really is, so tuning a decision rule
    on the raw sample over-predicts rare classes. Feeding
    `gbdt.stratification_weights(...)` here makes the estimate behave like macro
    F1 on a random sample of the city, using only legal information.
    """
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    w = np.ones(len(y_true)) if sample_weight is None else np.asarray(sample_weight, float)
    f1s = []
    for c in CLASSES:
        tp = w[(y_true == c) & (y_pred == c)].sum()
        fp = w[(y_true != c) & (y_pred == c)].sum()
        fn = w[(y_true == c) & (y_pred != c)].sum()
        denom = 2 * tp + fp + fn
        f1s.append(0.0 if denom <= 0 else 2 * tp / denom)
    return float(np.mean(f1s))


def tune_class_weights(proba: np.ndarray, y_true, n_rounds: int = 6, grid=None,
                       sample_weight=None, shrink: float = 0.0) -> np.ndarray:
    """
    Coordinate ascent on per-class multipliers to maximise macro F1.
    Fit this on a validation split of the TRAINING city, or on the labelled
    few-shot sample - never on the full target labels.
    shrink: pull the result back towards 1 (no adjustment). With only 20
    labelled rows (n=5 per class) the search overfits, so shrink hard; with 800
    rows it can be left at 0.
    """
    if grid is None:
        grid = np.exp(np.linspace(np.log(0.25), np.log(4.0), 25))
    w = np.ones(4)
    score = lambda ww: macro_f1_weighted(y_true, predict_weighted(proba, ww),  # noqa: E731
                                         sample_weight)
    best = score(w)
    for _ in range(n_rounds):
        improved = False
        for k in range(4):
            cur = w[k]
            for g in grid:
                w[k] = g
                s = score(w)
                if s > best + 1e-6:
                    best, cur, improved = s, g, True
            w[k] = cur
        if not improved:
            break
    w = w / w.mean()
    if shrink > 0:
        w = np.exp((1 - shrink) * np.log(w))
    return w / w.mean()


def prior_correct(proba: np.ndarray, train_prior: np.ndarray,
                  target_prior: np.ndarray) -> np.ndarray:
    """Re-weight probabilities from one class prior to another (label shift)."""
    adj = proba * (np.asarray(target_prior) / np.asarray(train_prior))[None, :]
    return adj / np.clip(adj.sum(axis=1, keepdims=True), 1e-12, None)


def em_target_prior(proba: np.ndarray, train_prior: np.ndarray,
                    n_iter: int = 200, tol: float = 1e-7,
                    pin: dict | None = None, shrink: float = 0.0) -> np.ndarray:
    """
    Saerens-Latinne-Decaestecker EM estimate of the target class prior from
    UNLABELLED target probabilities. Uses X_amsterdam only - no labels - so it
    stays inside the rules, but it only works if the model is roughly calibrated.
    pin:    fix some class shares to known values and renormalise the rest, e.g.
            {4: 0.111} - the organisers publish Amsterdam's class-4 share in the
            competitors' guide, so pinning it is citable, not leakage. Pinning
            also bounds the failure mode where EM drives a class to zero.
    shrink: pull the estimate back towards the training prior,
            q = (1-shrink)*q_EM + shrink*train_prior. EM is confident and
            frequently wrong under covariate shift; 0.5-0.7 is a sane default
            and should be chosen on the simulated-shift harness.
    """
    train_prior = np.asarray(train_prior, dtype=float)
    train_prior = train_prior / train_prior.sum()

    def apply_pin(p):
        if not pin:
            return p / p.sum()
        p = p.copy()
        free = [k for k in range(4) if (k + 1) not in pin]
        fixed_mass = sum(pin.values())
        for c, v in pin.items():
            p[c - 1] = v
        rest = p[free].sum()
        if rest > 0:
            p[free] *= (1.0 - fixed_mass) / rest
        else:
            p[free] = (1.0 - fixed_mass) / max(len(free), 1)
        return p

    prior = apply_pin(train_prior.copy())
    for _ in range(n_iter):
        post = proba * (prior / train_prior)[None, :]
        post /= np.clip(post.sum(axis=1, keepdims=True), 1e-12, None)
        new = apply_pin(post.mean(axis=0))
        if np.max(np.abs(new - prior)) < tol:
            prior = new
            break
        prior = new
    if shrink > 0:
        prior = apply_pin((1 - shrink) * prior + shrink * train_prior)
    return prior


def logit_adjust(proba: np.ndarray, train_prior: np.ndarray,
                 tau: float = 0.4) -> np.ndarray:
    """
    tau-scaled logit adjustment (Menon et al. 2021): divide by the training
    prior raised to tau.
    tau = 1 is Fisher-consistent for *balanced error rate*, i.e. macro RECALL.
    Macro F1 is a different metric and punishes over-predicting a rare class,
    so its optimum sits lower - tune tau on training-city CV and expect
    0.25-0.5. Leaving tau at 1.0 reliably loses macro F1 here.
    """
    p = np.asarray(proba, dtype=float)
    tp = np.asarray(train_prior, dtype=float)
    adj = p / np.clip(tp, 1e-12, None) ** float(tau)
    return adj / np.clip(adj.sum(axis=1, keepdims=True), 1e-12, None)


def macro_f1_plugin(proba: np.ndarray, prior: np.ndarray, n_iter: int = 50):
    """
    The plug-in decision rule for macro F1 (Ramaswamy et al., "Consistent
    algorithms for multiclass F-measures").
    argmax of a calibrated posterior maximises accuracy, not macro F1. The
    macro-F1-optimal rule is
            predict argmax_k  a_k * (p_k(x) - F1_k / 2),    a_k = 2 / (pi_k + q_k)
    where q_k is the predicted share of class k and F1_k the resulting per-class
    F1 - so it is a fixed point, solved by iteration. It needs the probabilities
    and a prior estimate, and NO target labels, which is why it is the right
    tool here.
    Returns (labels, diagnostics).
    """
    P = np.asarray(proba, dtype=float)
    pi = np.asarray(prior, dtype=float)
    pi = pi / pi.sum()
    f1 = np.zeros(4)
    best = None
    for it in range(n_iter):
        q = np.clip(np.bincount(np.argmax(P, axis=1), minlength=4) / len(P), 1e-6, None) \
            if it == 0 else q_new
        a = 2.0 / np.clip(pi + q, 1e-9, None)
        score = a[None, :] * (P - f1[None, :] / 2.0)
        lab_idx = np.argmax(score, axis=1)
        q_new = np.clip(np.bincount(lab_idx, minlength=4) / len(P), 1e-6, None)
        # soft per-class F1 implied by the current assignment and the posterior
        tp = np.array([P[lab_idx == k, k].sum() for k in range(4)]) / len(P)
        f1_new = 2 * tp / np.clip(pi + q_new, 1e-9, None)
        obj = float(np.mean(f1_new))
        if best is None or obj > best[0]:
            best = (obj, CLASSES[lab_idx].copy(), q_new.copy(), f1_new.copy())
        if np.max(np.abs(f1_new - f1)) < 1e-6 and it > 0:
            f1 = f1_new
            break
        f1 = f1_new
    obj, labels, q, f1 = best
    return labels, dict(expected_macro_f1=obj, predicted_share=q, per_class_f1=f1)


def adapt_pipeline(proba, src_prior, pin=None, tau: float = 0.5,
                   shrink: float = 0.7, use_em: bool = False,
                   use_plugin: bool = True):
    """
    The label-free adaptation chain: tau-logit-adjust -> (optional prior
    re-estimation) -> macro-F1 plug-in.
    `use_em` defaults to FALSE on purpose. On the simulated-shift harness, free
    SLD-EM was the worst thing tested: under covariate shift it drove the prior
    to [0.85, 0.10, 0.04, 0.02] against a truth of [0.34, 0.32, 0.22, 0.11] and
    cost 0.19 macro F1 on average, 0.35 at worst. Pinning class 4 at the
    published 0.111 and shrinking hard towards the source prior bounds the
    damage but does not repair it: EM assumes pure label shift, and this problem
    has a large covariate shift on top. Re-run `shift_harness` on the real
    Madrid features before turning it on.
    What survived the harness: tau-logit-adjustment at tau = 0.25-0.5 (small but
    positive in every regime) and the macro-F1 plug-in rule driven by the
    SOURCE prior (+0.033 mean, +0.038 worst case).
    """
    src_prior = np.asarray(src_prior, dtype=float)
    src_prior = src_prior / src_prior.sum()
    p = logit_adjust(proba, src_prior, tau=tau) if tau else np.asarray(proba, float)
    prior = src_prior
    if use_em:
        prior = em_target_prior(p, src_prior, pin=pin, shrink=shrink)
        p = prior_correct(p, src_prior, prior)
    if use_plugin:
        labels, diag = macro_f1_plugin(p, prior)
    else:
        labels, diag = CLASSES[np.argmax(p, axis=1)], {}
    diag["prior_used"] = prior
    return labels, p, diag


# --------------------------------------------------------------------------- #
# Spatial smoothing (stage 2): neighbouring pixels share a construction era
# --------------------------------------------------------------------------- #
def spatial_smooth(proba: np.ndarray, coords: np.ndarray, radius: float = 4.0,
                   alpha: float = 0.5, n_iter: int = 1) -> np.ndarray:
    """
    Blend each pixel's probabilities with a Gaussian-weighted average of its
    neighbours: p <- (1-alpha) * p + alpha * mean_neighbours(p).
    `radius` is in pixel units (1 unit = 30 m). Choose it on Madrid only.
    """
    from sklearn.neighbors import NearestNeighbors

    coords = np.asarray(coords, dtype=float)
    nn = NearestNeighbors(radius=radius * 2.0).fit(coords)
    dist, idx = nn.radius_neighbors(coords, radius=radius * 2.0)
    out = proba.copy()
    for _ in range(n_iter):
        nxt = out.copy()
        for i, (d, j) in enumerate(zip(dist, idx)):
            if len(j) <= 1:
                continue
            w = np.exp(-0.5 * (d / radius) ** 2)
            keep = j != i
            if not keep.any():
                continue
            avg = np.average(out[j[keep]], axis=0, weights=w[keep])
            nxt[i] = (1 - alpha) * out[i] + alpha * avg
        out = nxt / np.clip(nxt.sum(axis=1, keepdims=True), 1e-12, None)
    return out


# --------------------------------------------------------------------------- #
# The protocol harness
# --------------------------------------------------------------------------- #
def draw_shots(y_tgt, n_per_class: int, rng) -> np.ndarray:
    """Indices of a stratified n-per-class sample of the target city."""
    picks = []
    for c in CLASSES:
        pool = np.flatnonzero(np.asarray(y_tgt) == c)
        if len(pool) == 0:
            continue
        take = min(n_per_class, len(pool))
        picks.append(rng.choice(pool, size=take, replace=False))
    return np.concatenate(picks)


def evaluate_protocol(model_create, X_src, y_src, X_tgt, y_tgt,
                      n_list=(5, 25, 50, 100, 200), n_repeats: int = 20,
                      seed: int = 42, decision="argmax", verbose: bool = True,
                      eval_subsample: int | None = None):
    """
    Run the official few-shot protocol and return
    {n: {"scores": [...], "mean": float, "std": float, "per_class": array}}.
    `y_tgt` is touched ONLY here: to draw the n-per-class sample that the rules
    allow, and to score the held-out remainder. No tuning may read it.
    """
    X_src, y_src = np.asarray(X_src), np.asarray(y_src)
    X_tgt, y_tgt = np.asarray(X_tgt), np.asarray(y_tgt)

    # One FIXED evaluation subsample, shared by every n and every repeat.
    #
    # Scoring all 25,992 target rows on every one of 100 draws costs about 19
    # seconds a draw for a 1200-tree LightGBM over ~390 features - 32 minutes
    # per variant, hours for a table. A fixed subsample cuts that and is also a
    # better experiment: the +- then reflects only the draw of the labelled
    # sample, and every variant is compared on identical pixels. It is a plain
    # random sample, so it keeps the target city's real class proportions,
    # which is what macro F1 must be measured against.
    eval_pool = np.arange(len(y_tgt))
    if eval_subsample and eval_subsample < len(y_tgt):
        eval_pool = np.sort(np.random.default_rng(seed).choice(
            len(y_tgt), eval_subsample, replace=False))

    results = {}
    pool_sizes = {int(c): int((y_tgt == c).sum()) for c in CLASSES}
    for n in n_list:
        short = {c: k for c, k in pool_sizes.items() if k <= n}
        if short:
            print(f"  n={n}: only {short} pixels available in those classes - "
                  "the draw is capped and the held-out remainder shrinks")
        scores, per_class = [], []
        for r in range(n_repeats):
            rng = np.random.default_rng(seed + 1000 * n + r)
            shot_idx = draw_shots(y_tgt, n, rng)
            keep = np.setdiff1d(eval_pool, shot_idx, assume_unique=False)
            if len(keep) < 4 or len(np.unique(y_tgt[keep])) < 2:
                print(f"  n={n} rep={r}: nothing left to score after drawing "
                      "the sample - skipped")
                continue

            fn = model_create(X_src, y_src, X_tgt,
                              X_tgt[shot_idx], y_tgt[shot_idx])
            proba = np.asarray(fn(X_tgt[keep]), dtype=float)
            mask = keep
            if callable(decision):
                pred = decision(proba)
            else:
                pred = CLASSES[np.argmax(proba, axis=1)]
            scores.append(macro_f1(y_tgt[mask], pred))
            per_class.append(per_class_f1(y_tgt[mask], pred))
        if not scores:
            results[n] = {"scores": np.array([]), "mean": float("nan"),
                          "std": float("nan"), "per_class": np.full(4, np.nan)}
            continue
        results[n] = {
            "scores": np.array(scores),
            "mean": float(np.mean(scores)),
            "std": float(np.std(scores, ddof=1)) if len(scores) > 1 else 0.0,
            "per_class": np.mean(per_class, axis=0),
        }
        if verbose:
            print(f"  n={n:<4d} macro F1 = {results[n]['mean']:.4f} "
                  f"+- {results[n]['std']:.4f}   per-class "
                  f"{np.round(results[n]['per_class'], 3)}")
    return results


# --------------------------------------------------------------------------- #
# The simulated-shift harness: choose every adaptation knob on MADRID only.
#
# Nothing above can be tuned honestly without this. A step that helps under
# pure label shift can lose 0.06-0.10 macro F1 under covariate shift, and
# without a simulator there is no way to know which regime you are in.
# --------------------------------------------------------------------------- #
def resample_to_prior(y, prior, rng, n: int | None = None) -> np.ndarray:
    """Indices of a subsample of `y` whose class shares match `prior`."""
    y = np.asarray(y)
    prior = np.asarray(prior, dtype=float)
    prior = prior / prior.sum()
    pools = {c: np.flatnonzero(y == c) for c in CLASSES}
    if n is None:
        n = int(min(len(pools[c]) / prior[i] for i, c in enumerate(CLASSES)
                    if prior[i] > 0 and len(pools[c])))
    picks = []
    for i, c in enumerate(CLASSES):
        k = int(round(prior[i] * n))
        if k and len(pools[c]):
            picks.append(rng.choice(pools[c], size=min(k, len(pools[c])), replace=False))
    return np.concatenate(picks) if picks else np.array([], dtype=int)


def match_marginals(X_fold, X_src, X_tgt, strength: float = 1.0) -> np.ndarray:
    """
    Quantile-map each column of `X_fold` from the source city's marginal onto
    the target city's marginal. Uses X_amsterdam features only, so it is legal,
    and it turns a held-out Madrid fold into a stand-in for Amsterdam that keeps
    its true labels.
    """
    X_fold = np.asarray(X_fold, dtype=float)
    qs = np.linspace(0, 100, 201)
    gs = np.percentile(np.asarray(X_src, float), qs, axis=0)
    gt = np.percentile(np.asarray(X_tgt, float), qs, axis=0)
    out = np.empty_like(X_fold)
    for j in range(X_fold.shape[1]):
        r = np.interp(X_fold[:, j], gs[:, j], qs, left=0.0, right=100.0)
        out[:, j] = np.interp(r, qs, gt[:, j])
    return X_fold + strength * (out - X_fold)


def shift_harness(fit_fn, X, y, blocks, X_tgt, strategies: dict,
                  n_folds: int = 5, n_rep: int = 4, seed: int = 0,
                  verbose: bool = True):
    """
    For each held-out spatial block, build three versions of it:
      "clean"      - as-is (no shift),
      "label"      - resampled to a random Dirichlet class prior,
      "covariate"  - feature marginals mapped onto the target city,
      "both"       - label shift and covariate shift together,
    then score every decision strategy on each. A strategy is
    `fn(proba, src_prior) -> labels`.
    Returns a DataFrame of mean and worst-case macro F1 per (strategy, regime).
    Ship only the strategies that are non-negative in the WORST case.
    """
    import pandas as pd

    X, y = np.asarray(X, float), np.asarray(y)
    uniq = np.unique(blocks)
    rng0 = np.random.default_rng(seed)
    folds = np.array_split(rng0.permutation(uniq), n_folds)
    rows = []
    for fi, f in enumerate(folds):
        te = np.isin(blocks, f)
        tr = ~te
        model = fit_fn(X[tr], y[tr])
        src_prior = np.array([(y[tr] == c).mean() for c in CLASSES])
        Xte_clean = X[te]
        Xte_cov = match_marginals(Xte_clean, X[tr], X_tgt)
        for rep in range(n_rep):
            rng = np.random.default_rng(seed + 100 * fi + rep)
            alpha = rng.dirichlet(np.ones(4) * 1.5)
            sel_lab = resample_to_prior(y[te], alpha, rng)
            regimes = {
                "clean": (Xte_clean, y[te]),
                "label": (Xte_clean[sel_lab], y[te][sel_lab]),
                "covariate": (Xte_cov, y[te]),
                "both": (Xte_cov[sel_lab], y[te][sel_lab]),
            }
            for rname, (Xr, yr) in regimes.items():
                if len(yr) == 0:
                    continue
                proba = model.predict_proba(Xr)
                for sname, fn in strategies.items():
                    rows.append(dict(fold=fi, rep=rep, regime=rname,
                                     strategy=sname,
                                     macro_f1=macro_f1(yr, fn(proba, src_prior))))
    df = pd.DataFrame(rows)
    piv = df.pivot_table(index="strategy", columns="regime", values="macro_f1",
                         aggfunc="mean")
    worst = df.pivot_table(index="strategy", columns="regime", values="macro_f1",
                           aggfunc="min")
    if "argmax" in piv.index:
        regimes_cols = list(piv.columns)
        piv["delta_vs_argmax_mean"] = (piv[regimes_cols].mean(axis=1)
                                       - piv.loc["argmax", regimes_cols].mean())
        # PER-REGIME deltas, then the worst of them. The previous version took
        # each strategy's own worst cell anywhere and subtracted argmax's own
        # worst cell anywhere, which compares two different regimes - and two
        # different folds, since `worst` is a min over folds. A rule can then
        # look safe because its bad regime happens to be argmax's bad regime
        # too. What the decision actually needs is "in the regime where this
        # rule does worst RELATIVE to argmax, how much does it lose".
        dmean = piv[regimes_cols].sub(piv.loc["argmax", regimes_cols], axis=1)
        piv["delta_vs_argmax_worst_regime"] = dmean.min(axis=1)
        dworst = worst[regimes_cols].sub(worst.loc["argmax", regimes_cols], axis=1)
        piv["delta_vs_argmax_worst_fold"] = dworst.min(axis=1)
    if verbose:
        print(piv.round(4).to_string())
    return df, piv


DEFAULT_STRATEGIES = {
    "argmax": lambda p, sp: CLASSES[np.argmax(p, axis=1)],
    "logit_adj_tau0.25": lambda p, sp: CLASSES[np.argmax(logit_adjust(p, sp, 0.25), 1)],
    "logit_adj_tau0.5": lambda p, sp: CLASSES[np.argmax(logit_adjust(p, sp, 0.5), 1)],
    "logit_adj_tau1.0": lambda p, sp: CLASSES[np.argmax(logit_adjust(p, sp, 1.0), 1)],
    "em_free": lambda p, sp: CLASSES[np.argmax(
        prior_correct(p, sp, em_target_prior(p, sp)), 1)],
    "em_shrunk": lambda p, sp: CLASSES[np.argmax(
        prior_correct(p, sp, em_target_prior(p, sp, shrink=0.7)), 1)],
    "em_pinned_shrunk": lambda p, sp: CLASSES[np.argmax(
        prior_correct(p, sp, em_target_prior(p, sp, pin={4: 0.111}, shrink=0.7)), 1)],
    "plugin_srcprior": lambda p, sp: macro_f1_plugin(p, sp)[0],
    "tau0.5_then_plugin": lambda p, sp: adapt_pipeline(p, sp, tau=0.5)[0],
    "tau0.5_em_plugin": lambda p, sp: adapt_pipeline(p, sp, tau=0.5, use_em=True,
                                                     pin={4: 0.111}, shrink=0.7)[0],
}


def summarise(results: dict) -> str:
    lines = ["| n per class | macro F1 | std |", "|---|---|---|"]
    for n, r in results.items():
        lines.append(f"| {n} | {r['mean']:.4f} | {r['std']:.4f} |")
    means = [r["mean"] for r in results.values()]
    lines.append(f"\nAverage across n: {np.mean(means):.4f}   peak: {np.max(means):.4f}")
    return "\n".join(lines)

	"""
	The official evaluation protocol, plus the decision-rule and prior-shift tools
	that every model in this project shares.

	Protocol (organisers' Notebook 1, section 1.5):
	1. train a Stage-1 model on Madrid only;
	2. adapt it with n labelled Amsterdam pixels per class, n in {5,25,50,100,200};
	3. score macro F1 on the rest of Amsterdam;
	4. repeat over many random draws and report mean +- std for every n.

	`modelCreate` is any callable with this signature:

	predict_proba_fn = model_create(
	X_src, y_src, # Madrid, labelled
	X_tgt_unlabelled, # all target features, NO labels
	X_shot, y_shot, # the n-per-class labelled target sample
	)

	and must return a function mapping X -> (n, 4) probabilities. Keeping it a
	closure means every adaptation trick (recalibration, prototypes, continued
	boosting, retrieval prompts) fits the same harness.
	"""
	from __future__ import annotations

	import numpy as np
	from sklearn.metrics import confusion_matrix, f1_score

	CLASSES = np.array([1, 2, 3, 4])


	# --------------------------------------------------------------------------- #
	# Metrics
	# --------------------------------------------------------------------------- #
	def macro_f1(y_true, y_pred) -> float:
	return float(f1_score(y_true, y_pred, labels=CLASSES, average="macro", zero_division=0))


	def per_class_f1(y_true, y_pred) -> np.ndarray:
	return f1_score(y_true, y_pred, labels=CLASSES, average=None, zero_division=0)


	def confusion_proportions(y_true, y_pred) -> np.ndarray:
	cm = confusion_matrix(y_true, y_pred, labels=CLASSES).astype(float)
	return cm / np.clip(cm.sum(axis=1, keepdims=True), 1, None)


	# --------------------------------------------------------------------------- #
	# Decision rules. argmax is NOT the macro-F1 optimum.
	# --------------------------------------------------------------------------- #
	def predict_weighted(proba: np.ndarray, weights: np.ndarray) -> np.ndarray:
	"""argmax of w_k * p_k -> class label."""
	return CLASSES[np.argmax(proba * np.asarray(weights, dtype=float), axis=1)]


	def macro_f1_weighted(y_true, y_pred, sample_weight=None) -> float:
	"""
	Macro F1 computed from a weighted confusion matrix.

	This exists because the few-shot sample is drawn n-per-class: on it, a rare
	class looks far more precise than it really is, so tuning a decision rule
	on the raw sample over-predicts rare classes. Feeding
	`gbdt.stratification_weights(...)` here makes the estimate behave like macro
	F1 on a random sample of the city, using only legal information.
	"""
	y_true = np.asarray(y_true)
	y_pred = np.asarray(y_pred)
	w = np.ones(len(y_true)) if sample_weight is None else np.asarray(sample_weight, float)
	f1s = []
	for c in CLASSES:
	tp = w[(y_true == c) & (y_pred == c)].sum()
	fp = w[(y_true != c) & (y_pred == c)].sum()
	fn = w[(y_true == c) & (y_pred != c)].sum()
	denom = 2 * tp + fp + fn
	f1s.append(0.0 if denom <= 0 else 2 * tp / denom)
	return float(np.mean(f1s))


	def tune_class_weights(proba: np.ndarray, y_true, n_rounds: int = 6, grid=None,
	sample_weight=None, shrink: float = 0.0) -> np.ndarray:
	"""
	Coordinate ascent on per-class multipliers to maximise macro F1.

	Fit this on a validation split of the TRAINING city, or on the labelled
	few-shot sample - never on the full target labels.

	shrink: pull the result back towards 1 (no adjustment). With only 20
	labelled rows (n=5 per class) the search overfits, so shrink hard; with 800
	rows it can be left at 0.
	"""
	if grid is None:
	grid = np.exp(np.linspace(np.log(0.25), np.log(4.0), 25))
	w = np.ones(4)
	score = lambda ww: macro_f1_weighted(y_true, predict_weighted(proba, ww), # noqa: E731
	sample_weight)
	best = score(w)
	for _ in range(n_rounds):
	improved = False
	for k in range(4):
	cur = w[k]
	for g in grid:
	w[k] = g
	s = score(w)
	if s > best + 1e-6:
	best, cur, improved = s, g, True
	w[k] = cur
	if not improved:
	break
	w = w / w.mean()
	if shrink > 0:
	w = np.exp((1 - shrink) * np.log(w))
	return w / w.mean()


	def prior_correct(proba: np.ndarray, train_prior: np.ndarray,
	target_prior: np.ndarray) -> np.ndarray:
	"""Re-weight probabilities from one class prior to another (label shift)."""
	adj = proba * (np.asarray(target_prior) / np.asarray(train_prior))[None, :]
	return adj / np.clip(adj.sum(axis=1, keepdims=True), 1e-12, None)


	def em_target_prior(proba: np.ndarray, train_prior: np.ndarray,
	n_iter: int = 200, tol: float = 1e-7,
	pin: dict \| None = None, shrink: float = 0.0) -> np.ndarray:
	"""
	Saerens-Latinne-Decaestecker EM estimate of the target class prior from
	UNLABELLED target probabilities. Uses X_amsterdam only - no labels - so it
	stays inside the rules, but it only works if the model is roughly calibrated.

	pin: fix some class shares to known values and renormalise the rest, e.g.
	{4: 0.111} - the organisers publish Amsterdam's class-4 share in the
	competitors' guide, so pinning it is citable, not leakage. Pinning
	also bounds the failure mode where EM drives a class to zero.
	shrink: pull the estimate back towards the training prior,
	q = (1-shrink)q_EM + shrinktrain_prior. EM is confident and
	frequently wrong under covariate shift; 0.5-0.7 is a sane default
	and should be chosen on the simulated-shift harness.
	"""
	train_prior = np.asarray(train_prior, dtype=float)
	train_prior = train_prior / train_prior.sum()

	def apply_pin(p):
	if not pin:
	return p / p.sum()
	p = p.copy()
	free = [k for k in range(4) if (k + 1) not in pin]
	fixed_mass = sum(pin.values())
	for c, v in pin.items():
	p[c - 1] = v
	rest = p[free].sum()
	if rest > 0:
	p[free] *= (1.0 - fixed_mass) / rest
	else:
	p[free] = (1.0 - fixed_mass) / max(len(free), 1)
	return p

	prior = apply_pin(train_prior.copy())
	for _ in range(n_iter):
	post = proba * (prior / train_prior)[None, :]
	post /= np.clip(post.sum(axis=1, keepdims=True), 1e-12, None)
	new = apply_pin(post.mean(axis=0))
	if np.max(np.abs(new - prior)) < tol:
	prior = new
	break
	prior = new
	if shrink > 0:
	prior = apply_pin((1 - shrink) * prior + shrink * train_prior)
	return prior


	def logit_adjust(proba: np.ndarray, train_prior: np.ndarray,
	tau: float = 0.4) -> np.ndarray:
	"""
	tau-scaled logit adjustment (Menon et al. 2021): divide by the training
	prior raised to tau.

	tau = 1 is Fisher-consistent for balanced error rate, i.e. macro RECALL.
	Macro F1 is a different metric and punishes over-predicting a rare class,
	so its optimum sits lower - tune tau on training-city CV and expect
	0.25-0.5. Leaving tau at 1.0 reliably loses macro F1 here.
	"""
	p = np.asarray(proba, dtype=float)
	tp = np.asarray(train_prior, dtype=float)
	adj = p / np.clip(tp, 1e-12, None) ** float(tau)
	return adj / np.clip(adj.sum(axis=1, keepdims=True), 1e-12, None)


	def macro_f1_plugin(proba: np.ndarray, prior: np.ndarray, n_iter: int = 50):
	"""
	The plug-in decision rule for macro F1 (Ramaswamy et al., "Consistent
	algorithms for multiclass F-measures").

	argmax of a calibrated posterior maximises accuracy, not macro F1. The
	macro-F1-optimal rule is
	predict argmax_k a_k * (p_k(x) - F1_k / 2), a_k = 2 / (pi_k + q_k)
	where q_k is the predicted share of class k and F1_k the resulting per-class
	F1 - so it is a fixed point, solved by iteration. It needs the probabilities
	and a prior estimate, and NO target labels, which is why it is the right
	tool here.

	Returns (labels, diagnostics).
	"""
	P = np.asarray(proba, dtype=float)
	pi = np.asarray(prior, dtype=float)
	pi = pi / pi.sum()
	f1 = np.zeros(4)
	best = None
	for it in range(n_iter):
	q = np.clip(np.bincount(np.argmax(P, axis=1), minlength=4) / len(P), 1e-6, None) \
	if it == 0 else q_new
	a = 2.0 / np.clip(pi + q, 1e-9, None)
	score = a[None, :] * (P - f1[None, :] / 2.0)
	lab_idx = np.argmax(score, axis=1)
	q_new = np.clip(np.bincount(lab_idx, minlength=4) / len(P), 1e-6, None)
	# soft per-class F1 implied by the current assignment and the posterior
	tp = np.array([P[lab_idx == k, k].sum() for k in range(4)]) / len(P)
	f1_new = 2 * tp / np.clip(pi + q_new, 1e-9, None)
	obj = float(np.mean(f1_new))
	if best is None or obj > best[0]:
	best = (obj, CLASSES[lab_idx].copy(), q_new.copy(), f1_new.copy())
	if np.max(np.abs(f1_new - f1)) < 1e-6 and it > 0:
	f1 = f1_new
	break
	f1 = f1_new
	obj, labels, q, f1 = best
	return labels, dict(expected_macro_f1=obj, predicted_share=q, per_class_f1=f1)


	def adapt_pipeline(proba, src_prior, pin=None, tau: float = 0.5,
	shrink: float = 0.7, use_em: bool = False,
	use_plugin: bool = True):
	"""
	The label-free adaptation chain: tau-logit-adjust -> (optional prior
	re-estimation) -> macro-F1 plug-in.

	`use_em` defaults to FALSE on purpose. On the simulated-shift harness, free
	SLD-EM was the worst thing tested: under covariate shift it drove the prior
	to [0.85, 0.10, 0.04, 0.02] against a truth of [0.34, 0.32, 0.22, 0.11] and
	cost 0.19 macro F1 on average, 0.35 at worst. Pinning class 4 at the
	published 0.111 and shrinking hard towards the source prior bounds the
	damage but does not repair it: EM assumes pure label shift, and this problem
	has a large covariate shift on top. Re-run `shift_harness` on the real
	Madrid features before turning it on.

	What survived the harness: tau-logit-adjustment at tau = 0.25-0.5 (small but
	positive in every regime) and the macro-F1 plug-in rule driven by the
	SOURCE prior (+0.033 mean, +0.038 worst case).
	"""
	src_prior = np.asarray(src_prior, dtype=float)
	src_prior = src_prior / src_prior.sum()
	p = logit_adjust(proba, src_prior, tau=tau) if tau else np.asarray(proba, float)
	prior = src_prior
	if use_em:
	prior = em_target_prior(p, src_prior, pin=pin, shrink=shrink)
	p = prior_correct(p, src_prior, prior)
	if use_plugin:
	labels, diag = macro_f1_plugin(p, prior)
	else:
	labels, diag = CLASSES[np.argmax(p, axis=1)], {}
	diag["prior_used"] = prior
	return labels, p, diag


	# --------------------------------------------------------------------------- #
	# Spatial smoothing (stage 2): neighbouring pixels share a construction era
	# --------------------------------------------------------------------------- #
	def spatial_smooth(proba: np.ndarray, coords: np.ndarray, radius: float = 4.0,
	alpha: float = 0.5, n_iter: int = 1) -> np.ndarray:
	"""
	Blend each pixel's probabilities with a Gaussian-weighted average of its
	neighbours: p <- (1-alpha) * p + alpha * mean_neighbours(p).
	`radius` is in pixel units (1 unit = 30 m). Choose it on Madrid only.
	"""
	from sklearn.neighbors import NearestNeighbors

	coords = np.asarray(coords, dtype=float)
	nn = NearestNeighbors(radius=radius * 2.0).fit(coords)
	dist, idx = nn.radius_neighbors(coords, radius=radius * 2.0)
	out = proba.copy()
	for _ in range(n_iter):
	nxt = out.copy()
	for i, (d, j) in enumerate(zip(dist, idx)):
	if len(j) <= 1:
	continue
	w = np.exp(-0.5 * (d / radius) ** 2)
	keep = j != i
	if not keep.any():
	continue
	avg = np.average(out[j[keep]], axis=0, weights=w[keep])
	nxt[i] = (1 - alpha) * out[i] + alpha * avg
	out = nxt / np.clip(nxt.sum(axis=1, keepdims=True), 1e-12, None)
	return out


	# --------------------------------------------------------------------------- #
	# The protocol harness
	# --------------------------------------------------------------------------- #
	def draw_shots(y_tgt, n_per_class: int, rng) -> np.ndarray:
	"""Indices of a stratified n-per-class sample of the target city."""
	picks = []
	for c in CLASSES:
	pool = np.flatnonzero(np.asarray(y_tgt) == c)
	if len(pool) == 0:
	continue
	take = min(n_per_class, len(pool))
	picks.append(rng.choice(pool, size=take, replace=False))
	return np.concatenate(picks)


	def evaluate_protocol(model_create, X_src, y_src, X_tgt, y_tgt,
	n_list=(5, 25, 50, 100, 200), n_repeats: int = 20,
	seed: int = 42, decision="argmax", verbose: bool = True,
	eval_subsample: int \| None = None):
	"""
	Run the official few-shot protocol and return
	{n: {"scores": [...], "mean": float, "std": float, "per_class": array}}.

	`y_tgt` is touched ONLY here: to draw the n-per-class sample that the rules
	allow, and to score the held-out remainder. No tuning may read it.
	"""
	X_src, y_src = np.asarray(X_src), np.asarray(y_src)
	X_tgt, y_tgt = np.asarray(X_tgt), np.asarray(y_tgt)

	# One FIXED evaluation subsample, shared by every n and every repeat.
	#
	# Scoring all 25,992 target rows on every one of 100 draws costs about 19
	# seconds a draw for a 1200-tree LightGBM over ~390 features - 32 minutes
	# per variant, hours for a table. A fixed subsample cuts that and is also a
	# better experiment: the +- then reflects only the draw of the labelled
	# sample, and every variant is compared on identical pixels. It is a plain
	# random sample, so it keeps the target city's real class proportions,
	# which is what macro F1 must be measured against.
	eval_pool = np.arange(len(y_tgt))
	if eval_subsample and eval_subsample < len(y_tgt):
	eval_pool = np.sort(np.random.default_rng(seed).choice(
	len(y_tgt), eval_subsample, replace=False))

	results = {}
	pool_sizes = {int(c): int((y_tgt == c).sum()) for c in CLASSES}
	for n in n_list:
	short = {c: k for c, k in pool_sizes.items() if k <= n}
	if short:
	print(f" n={n}: only {short} pixels available in those classes - "
	"the draw is capped and the held-out remainder shrinks")
	scores, per_class = [], []
	for r in range(n_repeats):
	rng = np.random.default_rng(seed + 1000 * n + r)
	shot_idx = draw_shots(y_tgt, n, rng)
	keep = np.setdiff1d(eval_pool, shot_idx, assume_unique=False)
	if len(keep) < 4 or len(np.unique(y_tgt[keep])) < 2:
	print(f" n={n} rep={r}: nothing left to score after drawing "
	"the sample - skipped")
	continue

	fn = model_create(X_src, y_src, X_tgt,
	X_tgt[shot_idx], y_tgt[shot_idx])
	proba = np.asarray(fn(X_tgt[keep]), dtype=float)
	mask = keep
	if callable(decision):
	pred = decision(proba)
	else:
	pred = CLASSES[np.argmax(proba, axis=1)]
	scores.append(macro_f1(y_tgt[mask], pred))
	per_class.append(per_class_f1(y_tgt[mask], pred))
	if not scores:
	results[n] = {"scores": np.array([]), "mean": float("nan"),
	"std": float("nan"), "per_class": np.full(4, np.nan)}
	continue
	results[n] = {
	"scores": np.array(scores),
	"mean": float(np.mean(scores)),
	"std": float(np.std(scores, ddof=1)) if len(scores) > 1 else 0.0,
	"per_class": np.mean(per_class, axis=0),
	}
	if verbose:
	print(f" n={n:<4d} macro F1 = {results[n]['mean']:.4f} "
	f"+- {results[n]['std']:.4f} per-class "
	f"{np.round(results[n]['per_class'], 3)}")
	return results


	# --------------------------------------------------------------------------- #
	# The simulated-shift harness: choose every adaptation knob on MADRID only.
	#
	# Nothing above can be tuned honestly without this. A step that helps under
	# pure label shift can lose 0.06-0.10 macro F1 under covariate shift, and
	# without a simulator there is no way to know which regime you are in.
	# --------------------------------------------------------------------------- #
	def resample_to_prior(y, prior, rng, n: int \| None = None) -> np.ndarray:
	"""Indices of a subsample of `y` whose class shares match `prior`."""
	y = np.asarray(y)
	prior = np.asarray(prior, dtype=float)
	prior = prior / prior.sum()
	pools = {c: np.flatnonzero(y == c) for c in CLASSES}
	if n is None:
	n = int(min(len(pools[c]) / prior[i] for i, c in enumerate(CLASSES)
	if prior[i] > 0 and len(pools[c])))
	picks = []
	for i, c in enumerate(CLASSES):
	k = int(round(prior[i] * n))
	if k and len(pools[c]):
	picks.append(rng.choice(pools[c], size=min(k, len(pools[c])), replace=False))
	return np.concatenate(picks) if picks else np.array([], dtype=int)


	def match_marginals(X_fold, X_src, X_tgt, strength: float = 1.0) -> np.ndarray:
	"""
	Quantile-map each column of `X_fold` from the source city's marginal onto
	the target city's marginal. Uses X_amsterdam features only, so it is legal,
	and it turns a held-out Madrid fold into a stand-in for Amsterdam that keeps
	its true labels.
	"""
	X_fold = np.asarray(X_fold, dtype=float)
	qs = np.linspace(0, 100, 201)
	gs = np.percentile(np.asarray(X_src, float), qs, axis=0)
	gt = np.percentile(np.asarray(X_tgt, float), qs, axis=0)
	out = np.empty_like(X_fold)
	for j in range(X_fold.shape[1]):
	r = np.interp(X_fold[:, j], gs[:, j], qs, left=0.0, right=100.0)
	out[:, j] = np.interp(r, qs, gt[:, j])
	return X_fold + strength * (out - X_fold)


	def shift_harness(fit_fn, X, y, blocks, X_tgt, strategies: dict,
	n_folds: int = 5, n_rep: int = 4, seed: int = 0,
	verbose: bool = True):
	"""
	For each held-out spatial block, build three versions of it:
	"clean" - as-is (no shift),
	"label" - resampled to a random Dirichlet class prior,
	"covariate" - feature marginals mapped onto the target city,
	"both" - label shift and covariate shift together,
	then score every decision strategy on each. A strategy is
	`fn(proba, src_prior) -> labels`.

	Returns a DataFrame of mean and worst-case macro F1 per (strategy, regime).
	Ship only the strategies that are non-negative in the WORST case.
	"""
	import pandas as pd

	X, y = np.asarray(X, float), np.asarray(y)
	uniq = np.unique(blocks)
	rng0 = np.random.default_rng(seed)
	folds = np.array_split(rng0.permutation(uniq), n_folds)
	rows = []
	for fi, f in enumerate(folds):
	te = np.isin(blocks, f)
	tr = ~te
	model = fit_fn(X[tr], y[tr])
	src_prior = np.array([(y[tr] == c).mean() for c in CLASSES])
	Xte_clean = X[te]
	Xte_cov = match_marginals(Xte_clean, X[tr], X_tgt)
	for rep in range(n_rep):
	rng = np.random.default_rng(seed + 100 * fi + rep)
	alpha = rng.dirichlet(np.ones(4) * 1.5)
	sel_lab = resample_to_prior(y[te], alpha, rng)
	regimes = {
	"clean": (Xte_clean, y[te]),
	"label": (Xte_clean[sel_lab], y[te][sel_lab]),
	"covariate": (Xte_cov, y[te]),
	"both": (Xte_cov[sel_lab], y[te][sel_lab]),
	}
	for rname, (Xr, yr) in regimes.items():
	if len(yr) == 0:
	continue
	proba = model.predict_proba(Xr)
	for sname, fn in strategies.items():
	rows.append(dict(fold=fi, rep=rep, regime=rname,
	strategy=sname,
	macro_f1=macro_f1(yr, fn(proba, src_prior))))
	df = pd.DataFrame(rows)
	piv = df.pivot_table(index="strategy", columns="regime", values="macro_f1",
	aggfunc="mean")
	worst = df.pivot_table(index="strategy", columns="regime", values="macro_f1",
	aggfunc="min")
	if "argmax" in piv.index:
	regimes_cols = list(piv.columns)
	piv["delta_vs_argmax_mean"] = (piv[regimes_cols].mean(axis=1)
	- piv.loc["argmax", regimes_cols].mean())
	# PER-REGIME deltas, then the worst of them. The previous version took
	# each strategy's own worst cell anywhere and subtracted argmax's own
	# worst cell anywhere, which compares two different regimes - and two
	# different folds, since `worst` is a min over folds. A rule can then
	# look safe because its bad regime happens to be argmax's bad regime
	# too. What the decision actually needs is "in the regime where this
	# rule does worst RELATIVE to argmax, how much does it lose".
	dmean = piv[regimes_cols].sub(piv.loc["argmax", regimes_cols], axis=1)
	piv["delta_vs_argmax_worst_regime"] = dmean.min(axis=1)
	dworst = worst[regimes_cols].sub(worst.loc["argmax", regimes_cols], axis=1)
	piv["delta_vs_argmax_worst_fold"] = dworst.min(axis=1)
	if verbose:
	print(piv.round(4).to_string())
	return df, piv


	DEFAULT_STRATEGIES = {
	"argmax": lambda p, sp: CLASSES[np.argmax(p, axis=1)],
	"logit_adj_tau0.25": lambda p, sp: CLASSES[np.argmax(logit_adjust(p, sp, 0.25), 1)],
	"logit_adj_tau0.5": lambda p, sp: CLASSES[np.argmax(logit_adjust(p, sp, 0.5), 1)],
	"logit_adj_tau1.0": lambda p, sp: CLASSES[np.argmax(logit_adjust(p, sp, 1.0), 1)],
	"em_free": lambda p, sp: CLASSES[np.argmax(
	prior_correct(p, sp, em_target_prior(p, sp)), 1)],
	"em_shrunk": lambda p, sp: CLASSES[np.argmax(
	prior_correct(p, sp, em_target_prior(p, sp, shrink=0.7)), 1)],
	"em_pinned_shrunk": lambda p, sp: CLASSES[np.argmax(
	prior_correct(p, sp, em_target_prior(p, sp, pin={4: 0.111}, shrink=0.7)), 1)],
	"plugin_srcprior": lambda p, sp: macro_f1_plugin(p, sp)[0],
	"tau0.5_then_plugin": lambda p, sp: adapt_pipeline(p, sp, tau=0.5)[0],
	"tau0.5_em_plugin": lambda p, sp: adapt_pipeline(p, sp, tau=0.5, use_em=True,
	pin={4: 0.111}, shrink=0.7)[0],
	}


	def summarise(results: dict) -> str:
	lines = ["\| n per class \| macro F1 \| std \|", "\|---\|---\|---\|"]
	for n, r in results.items():
	lines.append(f"\| {n} \| {r['mean']:.4f} \| {r['std']:.4f} \|")
	means = [r["mean"] for r in results.values()]
	lines.append(f"\nAverage across n: {np.mean(means):.4f} peak: {np.max(means):.4f}")
	return "\n".join(lines)