Add FINAL-2 pipeline, overnight search scripts and results (11 Nov 2026)

2c2dd6c verified 4 days ago

14.9 kB

"""
nnlab - shared harness for the overnight model search (2026-11-11).
Every model family writes, for each protocol draw (n, rep), ONE array of class
probabilities over ALL 25,992 Amsterdam pixels (row order = the feature table):
    lab.save_pred("my_method", n, rep, P)      # P: (25992, 4), columns = class 1..4
and nothing else. Scoring is done here, identically for everybody:
    python tools/nnlab/lab.py score my_method [other_method ...]
    python tools/nnlab/lab.py board
RULES every module must keep (reviewers check these):
  * Amsterdam LABELS may be read only through lab.shot_labels(d, idx) for the
    indices returned by lab.shots_for(d, n, rep). Never touch d["ya"] directly.
  * Unlabelled Amsterdam features / coordinates may be used freely (transductive).
  * Hyper-parameters are fixed in advance or chosen on Madrid (lab.madrid_split
    gives a spatial pseudo-target). Never chosen by looking at an Amsterdam score.
  * The held-out organisers' data is never opened.
The draws are the protocol's: seed 42, rng(42 + 1000 n + rep), n in
(5, 25, 50, 100, 200). Shots are removed from the scored set.
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent.parent
# ROOT itself holds winners.py when these scripts run from the published bundle
for _p in (ROOT, HERE, ROOT / "src", ROOT / "tools"):
    if str(_p) not in sys.path:
        sys.path.insert(0, str(_p))
THREADS = int(os.environ.get("NNLAB_THREADS", "3"))
os.environ.setdefault("OMP_NUM_THREADS", str(THREADS))
os.environ.setdefault("MKL_NUM_THREADS", str(THREADS))

import numpy as np                      # noqa: E402
import pandas as pd                     # noqa: E402

import winners as W                     # noqa: E402  (verified port of the protocol + rule)

PROC = ROOT / "data" / "processed"
OUT = ROOT / "reports" / "nnlab"
PRED = OUT / "preds"
SCORES = OUT / "scores"
CACHE = OUT / "cache"                   # embeddings, time-series tensors, ...
for _d in (PRED, SCORES, CACHE):
    _d.mkdir(parents=True, exist_ok=True)

SEED = 42
N_LIST = (5, 25, 50, 100, 200)
REPS = int(os.environ.get("NNLAB_REPS", "5"))   # reps 0..4 by default; NNLAB_REPS=12 adds the confirmation draws
CLASSES = np.array([1, 2, 3, 4])
PID = ["px_key", "py_key"]
BLOCK = 67                              # 67 px = 2 km
SMOOTH_R, SMOOTH_ALPHA = 6, 0.5         # teammate's choice, made on Madrid


def torch_setup(seed: int = 0):
    import torch
    torch.set_num_threads(THREADS)
    torch.manual_seed(seed)
    np.random.seed(seed)
    return torch


# --------------------------------------------------------------------------- #
# Data
# --------------------------------------------------------------------------- #
def load(official60: bool = True) -> dict:
    """Madrid + Amsterdam tables. Xm/Xa: our 393 features (float32, raw scale);
    Xm60/Xa60: the organisers' official 60 (first valid observation per year);
    pm/pa: integer pixel coordinates (30 m units); ym: Madrid labels;
    prior: Madrid class shares. Amsterdam labels are kept under a private key."""
    feats = W.feature_names()
    mad = pd.read_parquet(PROC / "madrid_features.parquet")
    ams = pd.read_parquet(PROC / "amsterdam_features.parquet")
    d = dict(
        feats=feats,
        Xm=mad[feats].to_numpy(np.float32), Xa=ams[feats].to_numpy(np.float32),
        ym=mad["age_class"].to_numpy().astype(int),
        pm=mad[PID].to_numpy().astype(np.int64), pa=ams[PID].to_numpy().astype(np.int64),
        _ya=ams["age_class"].to_numpy().astype(int),
    )
    d["prior"] = np.array([(d["ym"] == c).mean() for c in CLASSES])
    if official60:
        import config as C
        f60 = list(C.FEATURE_COLS_60)
        for city, ref, key in (("madrid", mad, "Xm60"), ("amsterdam", ams, "Xa60")):
            o = pd.read_parquet(PROC / f"{city}_official60_first.parquet")
            o = ref[PID].merge(o, on=PID, how="left")
            assert len(o) == len(ref) and o[f60].notna().all().all()
            d[key] = o[f60].to_numpy(np.float32)
        d["feats60"] = f60
    return d


def shots_for(d: dict, n: int, rep: int) -> np.ndarray:
    """Row indices of the n-per-class labelled Amsterdam pixels of draw (n, rep)."""
    return W.shots_for(d["_ya"], n, rep)


def shot_labels(d: dict, idx: np.ndarray) -> np.ndarray:
    """Labels (1..4) of shot rows. The ONLY sanctioned way to read Amsterdam labels."""
    return d["_ya"][np.asarray(idx)]


def draws(reps: int = REPS, n_list=N_LIST):
    for n in n_list:
        for r in range(reps):
            yield n, r


def balanced_class_weights(y) -> np.ndarray:
    return W.balanced_class_weights(y)


# per-city input transforms (label-free) ------------------------------------- #
def robustz(X: np.ndarray, clip: float | None = 10.0) -> np.ndarray:
    """Median/IQR scaling with this matrix's own statistics (call once per city)."""
    X = np.asarray(X, np.float64)
    med = np.median(X, 0)
    q75, q25 = np.percentile(X, [75, 25], 0)
    iqr = np.where((q75 - q25) > 1e-12, q75 - q25, 1.0)
    Z = (X - med) / iqr
    if clip:
        Z = np.clip(Z, -clip, clip)
    return Z.astype(np.float32)


def standardise(X: np.ndarray) -> np.ndarray:
    X = np.asarray(X, np.float64)
    sd = X.std(0)
    sd[sd == 0] = 1
    return ((X - X.mean(0)) / sd).astype(np.float32)


def rankgauss(X: np.ndarray) -> np.ndarray:
    """Per-column rank -> standard normal quantile, within this matrix (one city)."""
    from scipy.special import ndtri
    from scipy.stats import rankdata
    X = np.asarray(X, np.float64)
    R = rankdata(X, axis=0, method="average")
    return ndtri((R - 0.5) / len(X)).astype(np.float32)


# Madrid-only model selection ------------------------------------------------- #
def block_ids(p: np.ndarray, block: int = BLOCK) -> np.ndarray:
    b = p // block
    return (b[:, 0] * 100_003 + b[:, 1]).astype(np.int64)


def madrid_split(pm: np.ndarray, frac: float = 0.25, axis: int = 0, high: bool = True):
    """A spatial pseudo-target inside Madrid: the `frac` of 2 km blocks furthest
    along one axis. Returns (source_idx, pseudo_target_idx). Use it to choose
    epochs / regularisation without touching Amsterdam labels."""
    b = pm[:, axis] // BLOCK
    cut = np.quantile(b, 1 - frac) if high else np.quantile(b, frac)
    tgt = b > cut if high else b < cut
    return np.flatnonzero(~tgt), np.flatnonzero(tgt)


# --------------------------------------------------------------------------- #
# Heads on a frozen backbone
# --------------------------------------------------------------------------- #
def softmax(L: np.ndarray) -> np.ndarray:
    L = L - L.max(1, keepdims=True)
    E = np.exp(L)
    return E / E.sum(1, keepdims=True)


def finetune_last_layer(W0, b0, Zs, ys, lam: float = 1.0, class_balanced: bool = True,
                        max_iter: int = 200):
    """True last-layer fine-tuning: start from the Madrid-trained last layer
    (W0: (4, k), b0: (4,)) and minimise, on the labelled shots only,
        sum_i w_i CE(softmax(W z_i + b), y_i)  +  lam * (||W - W0||^2 + ||b - b0||^2).
    The CE is a SUM, so the pull toward the source weights fades as n grows.
    Returns (W, b). Probabilities: softmax(Z @ W.T + b)."""
    torch = torch_setup(0)
    Z = torch.tensor(np.asarray(Zs), dtype=torch.float64)
    y = torch.tensor(np.asarray(ys) - 1, dtype=torch.long)
    w = torch.ones(len(y), dtype=torch.float64)
    if class_balanced:
        w = torch.tensor(balanced_class_weights(np.asarray(ys))[np.asarray(ys) - 1], dtype=torch.float64)
    W0t = torch.tensor(np.asarray(W0), dtype=torch.float64)
    b0t = torch.tensor(np.asarray(b0), dtype=torch.float64)
    Wt, bt = W0t.clone().requires_grad_(True), b0t.clone().requires_grad_(True)
    opt = torch.optim.LBFGS([Wt, bt], max_iter=max_iter, line_search_fn="strong_wolfe")

    def closure():
        opt.zero_grad()
        ce = torch.nn.functional.cross_entropy(Z @ Wt.T + bt, y, reduction="none")
        loss = (w * ce).sum() + lam * (((Wt - W0t) ** 2).sum() + ((bt - b0t) ** 2).sum())
        loss.backward()
        return loss

    opt.step(closure)
    return Wt.detach().numpy(), bt.detach().numpy()


def joint_head(Zm, ym, Zs, ys, shot_weight: float = 30.0, C: float = 1.0, max_iter: int = 400):
    """Last layer re-fitted on frozen embeddings of Madrid + shots (shots
    weighted `shot_weight`x, classes balanced) - the frozen-backbone analogue
    of the retrained LightGBM. Returns a fitted sklearn LogisticRegression
    (predict_proba columns = classes 1..4)."""
    from sklearn.linear_model import LogisticRegression
    X = np.vstack([Zm, Zs])
    y = np.r_[np.asarray(ym), np.asarray(ys)]
    w = np.r_[np.ones(len(ym)), np.full(len(ys), shot_weight)]
    w = w * balanced_class_weights(y)[y - 1]
    return LogisticRegression(C=C, max_iter=max_iter).fit(X, y, sample_weight=w)


# --------------------------------------------------------------------------- #
# Prediction store
# --------------------------------------------------------------------------- #
def pred_path(method: str, n: int, rep: int) -> Path:
    return PRED / method / f"n{n}_r{rep}.npy"


def have_pred(method: str, n: int, rep: int) -> bool:
    return pred_path(method, n, rep).exists()


def save_pred(method: str, n: int, rep: int, P: np.ndarray, meta: dict | None = None):
    """P: (25992, 4) probabilities for ALL Amsterdam pixels, feature-table row
    order. A method that ignores the shots saves once with n=0, rep=0."""
    P = np.asarray(P, np.float32)
    assert P.shape == (25992, 4) and np.isfinite(P).all(), P.shape
    p = pred_path(method, n, rep)
    p.parent.mkdir(parents=True, exist_ok=True)
    np.save(p, P)
    if meta is not None:
        (p.parent / "meta.json").write_text(json.dumps(meta, indent=1, default=str))


def load_pred(method: str, n: int, rep: int) -> np.ndarray:
    p = pred_path(method, n, rep)
    if not p.exists():
        p = pred_path(method, 0, 0)       # shot-independent method
    return np.load(p).astype(np.float64)


# --------------------------------------------------------------------------- #
# Scoring (the only place Amsterdam labels are compared with predictions)
# --------------------------------------------------------------------------- #
_ADJ = {}


def adjacency(coords: np.ndarray, r: float = SMOOTH_R):
    """Sparse 0/1 matrix: pixels within r (Euclidean, pixel units), self included."""
    from scipy.sparse import coo_matrix, identity
    from scipy.spatial import cKDTree
    key = (coords.shape[0], float(r), int(coords[:, 0].sum()), int(coords[:, 1].sum()))
    if key not in _ADJ:
        pr = cKDTree(coords.astype(float)).query_pairs(r, output_type="ndarray")
        n = len(coords)
        A = coo_matrix((np.ones(len(pr)), (pr[:, 0], pr[:, 1])), shape=(n, n)).tocsr()
        _ADJ[key] = (A + A.T + identity(n, format="csr")).tocsr()
    return _ADJ[key]


def smooth(P: np.ndarray, coords: np.ndarray, r: float = SMOOTH_R, alpha: float = SMOOTH_ALPHA,
           keep: np.ndarray | None = None) -> np.ndarray:
    """Teammate's neighbour smoothing: alpha * own + (1 - alpha) * mean over
    pixels within r (self included). With `keep`, only those rows take part
    (P must then already be P_all[keep]) - shots are never neighbours."""
    A = adjacency(coords, r)
    if keep is not None:
        A = A[keep][:, keep]
    deg = np.asarray(A.sum(1)).ravel()
    return alpha * P + (1 - alpha) * (A @ P) / deg[:, None]


def half_A(pa: np.ndarray) -> np.ndarray:
    """Checkerboard of 2 km blocks. Half A is for comparing variants; half B is
    looked at only to confirm the final picks."""
    b = pa // BLOCK
    return ((b[:, 0] + b[:, 1]) % 2) == 0


def score_method(d: dict, method: str, reps: int = REPS, n_list=N_LIST, verbose: bool = True) -> dict:
    ya, pa, prior = d["_ya"], d["pa"], d["prior"]
    inA = half_A(pa)
    rows = []
    for n, rep in draws(reps, n_list):
        if not (have_pred(method, n, rep) or have_pred(method, 0, 0)):
            continue
        P_all = load_pred(method, n, rep)
        shot = shots_for(d, n, rep)
        keep = np.setdiff1d(np.arange(len(ya)), shot)
        P, y = P_all[keep], ya[keep]
        lab = {
            "argmax": CLASSES[P.argmax(1)],
            "rule": W.decide(P, prior),
            "rule_smooth": W.decide(smooth(P, pa, keep=keep), prior),
        }
        row = {"n": n, "rep": rep}
        for k, l in lab.items():
            row[k] = W.macro_f1(y, l)
            row[k + "_A"] = W.macro_f1(y[inA[keep]], l[inA[keep]])
            row[k + "_B"] = W.macro_f1(y[~inA[keep]], l[~inA[keep]])
        rows.append(row)
    if not rows:
        return {}
    df = pd.DataFrame(rows)
    summ = df.groupby("n").mean(numeric_only=True).drop(columns="rep")
    std = df.groupby("n").std(numeric_only=True).drop(columns="rep")
    out = {"method": method, "draws": rows, "n_draws": int(len(rows)),
           "mean": {c: {int(n): float(v) for n, v in summ[c].items()} for c in summ.columns},
           "std": {c: {int(n): float(v) for n, v in std[c].items()} for c in std.columns}}
    (SCORES / f"{method}.json").write_text(json.dumps(out, indent=1))
    if verbose:
        for k in ("argmax", "rule", "rule_smooth"):
            m = summ[k]
            print(f"{method:<34s} {k:<12s} " + "  ".join(f"n{n}={m[n]:.4f}" for n in m.index)
                  + f"  avg={m.mean():.4f}  peak={m.max():.4f}  ({len(rows)} draws)", flush=True)
    return out


def board(metric_suffix: str = "") -> pd.DataFrame:
    """All scored methods, one row per (method, decision). suffix '' / '_A' / '_B'."""
    rows = []
    for f in sorted(SCORES.glob("*.json")):
        s = json.loads(f.read_text())
        for k in ("argmax", "rule", "rule_smooth"):
            m = s["mean"].get(k + metric_suffix)
            if not m:
                continue
            r = {"method": s["method"], "decision": k, **{f"n{n}": m.get(str(n), np.nan) for n in N_LIST}}
            v = [r[f"n{n}"] for n in N_LIST]
            r["avg"], r["peak"], r["draws"] = float(np.nanmean(v)), float(np.nanmax(v)), s["n_draws"]
            rows.append(r)
    df = pd.DataFrame(rows).sort_values("avg", ascending=False)
    df.to_csv(OUT / f"leaderboard{metric_suffix}.csv", index=False)
    return df


if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "board"
    if cmd == "score":
        d = load(official60=False)
        names = sys.argv[2:] or sorted(p.name for p in PRED.iterdir() if p.is_dir())
        t0 = time.time()
        for m in names:
            score_method(d, m)
        print(f"scored {len(names)} methods in {time.time() - t0:.0f}s")
    else:
        suffix = sys.argv[2] if len(sys.argv) > 2 else ""
        pd.set_option("display.width", 250)
        print(board(suffix).round(4).to_string(index=False))

	"""
	nnlab - shared harness for the overnight model search (2026-11-11).

	Every model family writes, for each protocol draw (n, rep), ONE array of class
	probabilities over ALL 25,992 Amsterdam pixels (row order = the feature table):

	lab.save_pred("my_method", n, rep, P) # P: (25992, 4), columns = class 1..4

	and nothing else. Scoring is done here, identically for everybody:

	python tools/nnlab/lab.py score my_method [other_method ...]
	python tools/nnlab/lab.py board

	RULES every module must keep (reviewers check these):
	* Amsterdam LABELS may be read only through lab.shot_labels(d, idx) for the
	indices returned by lab.shots_for(d, n, rep). Never touch d["ya"] directly.
	* Unlabelled Amsterdam features / coordinates may be used freely (transductive).
	* Hyper-parameters are fixed in advance or chosen on Madrid (lab.madrid_split
	gives a spatial pseudo-target). Never chosen by looking at an Amsterdam score.
	* The held-out organisers' data is never opened.

	The draws are the protocol's: seed 42, rng(42 + 1000 n + rep), n in
	(5, 25, 50, 100, 200). Shots are removed from the scored set.
	"""
	from __future__ import annotations

	import json
	import os
	import sys
	import time
	from pathlib import Path

	HERE = Path(__file__).resolve().parent
	ROOT = HERE.parent.parent
	# ROOT itself holds winners.py when these scripts run from the published bundle
	for _p in (ROOT, HERE, ROOT / "src", ROOT / "tools"):
	if str(_p) not in sys.path:
	sys.path.insert(0, str(_p))
	THREADS = int(os.environ.get("NNLAB_THREADS", "3"))
	os.environ.setdefault("OMP_NUM_THREADS", str(THREADS))
	os.environ.setdefault("MKL_NUM_THREADS", str(THREADS))

	import numpy as np # noqa: E402
	import pandas as pd # noqa: E402

	import winners as W # noqa: E402 (verified port of the protocol + rule)

	PROC = ROOT / "data" / "processed"
	OUT = ROOT / "reports" / "nnlab"
	PRED = OUT / "preds"
	SCORES = OUT / "scores"
	CACHE = OUT / "cache" # embeddings, time-series tensors, ...
	for _d in (PRED, SCORES, CACHE):
	_d.mkdir(parents=True, exist_ok=True)

	SEED = 42
	N_LIST = (5, 25, 50, 100, 200)
	REPS = int(os.environ.get("NNLAB_REPS", "5")) # reps 0..4 by default; NNLAB_REPS=12 adds the confirmation draws
	CLASSES = np.array([1, 2, 3, 4])
	PID = ["px_key", "py_key"]
	BLOCK = 67 # 67 px = 2 km
	SMOOTH_R, SMOOTH_ALPHA = 6, 0.5 # teammate's choice, made on Madrid


	def torch_setup(seed: int = 0):
	import torch
	torch.set_num_threads(THREADS)
	torch.manual_seed(seed)
	np.random.seed(seed)
	return torch


	# --------------------------------------------------------------------------- #
	# Data
	# --------------------------------------------------------------------------- #
	def load(official60: bool = True) -> dict:
	"""Madrid + Amsterdam tables. Xm/Xa: our 393 features (float32, raw scale);
	Xm60/Xa60: the organisers' official 60 (first valid observation per year);
	pm/pa: integer pixel coordinates (30 m units); ym: Madrid labels;
	prior: Madrid class shares. Amsterdam labels are kept under a private key."""
	feats = W.feature_names()
	mad = pd.read_parquet(PROC / "madrid_features.parquet")
	ams = pd.read_parquet(PROC / "amsterdam_features.parquet")
	d = dict(
	feats=feats,
	Xm=mad[feats].to_numpy(np.float32), Xa=ams[feats].to_numpy(np.float32),
	ym=mad["age_class"].to_numpy().astype(int),
	pm=mad[PID].to_numpy().astype(np.int64), pa=ams[PID].to_numpy().astype(np.int64),
	_ya=ams["age_class"].to_numpy().astype(int),
	)
	d["prior"] = np.array([(d["ym"] == c).mean() for c in CLASSES])
	if official60:
	import config as C
	f60 = list(C.FEATURE_COLS_60)
	for city, ref, key in (("madrid", mad, "Xm60"), ("amsterdam", ams, "Xa60")):
	o = pd.read_parquet(PROC / f"{city}_official60_first.parquet")
	o = ref[PID].merge(o, on=PID, how="left")
	assert len(o) == len(ref) and o[f60].notna().all().all()
	d[key] = o[f60].to_numpy(np.float32)
	d["feats60"] = f60
	return d


	def shots_for(d: dict, n: int, rep: int) -> np.ndarray:
	"""Row indices of the n-per-class labelled Amsterdam pixels of draw (n, rep)."""
	return W.shots_for(d["_ya"], n, rep)


	def shot_labels(d: dict, idx: np.ndarray) -> np.ndarray:
	"""Labels (1..4) of shot rows. The ONLY sanctioned way to read Amsterdam labels."""
	return d["_ya"][np.asarray(idx)]


	def draws(reps: int = REPS, n_list=N_LIST):
	for n in n_list:
	for r in range(reps):
	yield n, r


	def balanced_class_weights(y) -> np.ndarray:
	return W.balanced_class_weights(y)


	# per-city input transforms (label-free) ------------------------------------- #
	def robustz(X: np.ndarray, clip: float \| None = 10.0) -> np.ndarray:
	"""Median/IQR scaling with this matrix's own statistics (call once per city)."""
	X = np.asarray(X, np.float64)
	med = np.median(X, 0)
	q75, q25 = np.percentile(X, [75, 25], 0)
	iqr = np.where((q75 - q25) > 1e-12, q75 - q25, 1.0)
	Z = (X - med) / iqr
	if clip:
	Z = np.clip(Z, -clip, clip)
	return Z.astype(np.float32)


	def standardise(X: np.ndarray) -> np.ndarray:
	X = np.asarray(X, np.float64)
	sd = X.std(0)
	sd[sd == 0] = 1
	return ((X - X.mean(0)) / sd).astype(np.float32)


	def rankgauss(X: np.ndarray) -> np.ndarray:
	"""Per-column rank -> standard normal quantile, within this matrix (one city)."""
	from scipy.special import ndtri
	from scipy.stats import rankdata
	X = np.asarray(X, np.float64)
	R = rankdata(X, axis=0, method="average")
	return ndtri((R - 0.5) / len(X)).astype(np.float32)


	# Madrid-only model selection ------------------------------------------------- #
	def block_ids(p: np.ndarray, block: int = BLOCK) -> np.ndarray:
	b = p // block
	return (b[:, 0] * 100_003 + b[:, 1]).astype(np.int64)


	def madrid_split(pm: np.ndarray, frac: float = 0.25, axis: int = 0, high: bool = True):
	"""A spatial pseudo-target inside Madrid: the `frac` of 2 km blocks furthest
	along one axis. Returns (source_idx, pseudo_target_idx). Use it to choose
	epochs / regularisation without touching Amsterdam labels."""
	b = pm[:, axis] // BLOCK
	cut = np.quantile(b, 1 - frac) if high else np.quantile(b, frac)
	tgt = b > cut if high else b < cut
	return np.flatnonzero(~tgt), np.flatnonzero(tgt)


	# --------------------------------------------------------------------------- #
	# Heads on a frozen backbone
	# --------------------------------------------------------------------------- #
	def softmax(L: np.ndarray) -> np.ndarray:
	L = L - L.max(1, keepdims=True)
	E = np.exp(L)
	return E / E.sum(1, keepdims=True)


	def finetune_last_layer(W0, b0, Zs, ys, lam: float = 1.0, class_balanced: bool = True,
	max_iter: int = 200):
	"""True last-layer fine-tuning: start from the Madrid-trained last layer
	(W0: (4, k), b0: (4,)) and minimise, on the labelled shots only,
	sum_i w_i CE(softmax(W z_i + b), y_i) + lam * (\|\|W - W0\|\|^2 + \|\|b - b0\|\|^2).
	The CE is a SUM, so the pull toward the source weights fades as n grows.
	Returns (W, b). Probabilities: softmax(Z @ W.T + b)."""
	torch = torch_setup(0)
	Z = torch.tensor(np.asarray(Zs), dtype=torch.float64)
	y = torch.tensor(np.asarray(ys) - 1, dtype=torch.long)
	w = torch.ones(len(y), dtype=torch.float64)
	if class_balanced:
	w = torch.tensor(balanced_class_weights(np.asarray(ys))[np.asarray(ys) - 1], dtype=torch.float64)
	W0t = torch.tensor(np.asarray(W0), dtype=torch.float64)
	b0t = torch.tensor(np.asarray(b0), dtype=torch.float64)
	Wt, bt = W0t.clone().requires_grad_(True), b0t.clone().requires_grad_(True)
	opt = torch.optim.LBFGS([Wt, bt], max_iter=max_iter, line_search_fn="strong_wolfe")

	def closure():
	opt.zero_grad()
	ce = torch.nn.functional.cross_entropy(Z @ Wt.T + bt, y, reduction="none")
	loss = (w * ce).sum() + lam * (((Wt - W0t) 2).sum() + ((bt - b0t) 2).sum())
	loss.backward()
	return loss

	opt.step(closure)
	return Wt.detach().numpy(), bt.detach().numpy()


	def joint_head(Zm, ym, Zs, ys, shot_weight: float = 30.0, C: float = 1.0, max_iter: int = 400):
	"""Last layer re-fitted on frozen embeddings of Madrid + shots (shots
	weighted `shot_weight`x, classes balanced) - the frozen-backbone analogue
	of the retrained LightGBM. Returns a fitted sklearn LogisticRegression
	(predict_proba columns = classes 1..4)."""
	from sklearn.linear_model import LogisticRegression
	X = np.vstack([Zm, Zs])
	y = np.r_[np.asarray(ym), np.asarray(ys)]
	w = np.r_[np.ones(len(ym)), np.full(len(ys), shot_weight)]
	w = w * balanced_class_weights(y)[y - 1]
	return LogisticRegression(C=C, max_iter=max_iter).fit(X, y, sample_weight=w)


	# --------------------------------------------------------------------------- #
	# Prediction store
	# --------------------------------------------------------------------------- #
	def pred_path(method: str, n: int, rep: int) -> Path:
	return PRED / method / f"n{n}_r{rep}.npy"


	def have_pred(method: str, n: int, rep: int) -> bool:
	return pred_path(method, n, rep).exists()


	def save_pred(method: str, n: int, rep: int, P: np.ndarray, meta: dict \| None = None):
	"""P: (25992, 4) probabilities for ALL Amsterdam pixels, feature-table row
	order. A method that ignores the shots saves once with n=0, rep=0."""
	P = np.asarray(P, np.float32)
	assert P.shape == (25992, 4) and np.isfinite(P).all(), P.shape
	p = pred_path(method, n, rep)
	p.parent.mkdir(parents=True, exist_ok=True)
	np.save(p, P)
	if meta is not None:
	(p.parent / "meta.json").write_text(json.dumps(meta, indent=1, default=str))


	def load_pred(method: str, n: int, rep: int) -> np.ndarray:
	p = pred_path(method, n, rep)
	if not p.exists():
	p = pred_path(method, 0, 0) # shot-independent method
	return np.load(p).astype(np.float64)


	# --------------------------------------------------------------------------- #
	# Scoring (the only place Amsterdam labels are compared with predictions)
	# --------------------------------------------------------------------------- #
	_ADJ = {}


	def adjacency(coords: np.ndarray, r: float = SMOOTH_R):
	"""Sparse 0/1 matrix: pixels within r (Euclidean, pixel units), self included."""
	from scipy.sparse import coo_matrix, identity
	from scipy.spatial import cKDTree
	key = (coords.shape[0], float(r), int(coords[:, 0].sum()), int(coords[:, 1].sum()))
	if key not in _ADJ:
	pr = cKDTree(coords.astype(float)).query_pairs(r, output_type="ndarray")
	n = len(coords)
	A = coo_matrix((np.ones(len(pr)), (pr[:, 0], pr[:, 1])), shape=(n, n)).tocsr()
	_ADJ[key] = (A + A.T + identity(n, format="csr")).tocsr()
	return _ADJ[key]


	def smooth(P: np.ndarray, coords: np.ndarray, r: float = SMOOTH_R, alpha: float = SMOOTH_ALPHA,
	keep: np.ndarray \| None = None) -> np.ndarray:
	"""Teammate's neighbour smoothing: alpha * own + (1 - alpha) * mean over
	pixels within r (self included). With `keep`, only those rows take part
	(P must then already be P_all[keep]) - shots are never neighbours."""
	A = adjacency(coords, r)
	if keep is not None:
	A = A[keep][:, keep]
	deg = np.asarray(A.sum(1)).ravel()
	return alpha * P + (1 - alpha) * (A @ P) / deg[:, None]


	def half_A(pa: np.ndarray) -> np.ndarray:
	"""Checkerboard of 2 km blocks. Half A is for comparing variants; half B is
	looked at only to confirm the final picks."""
	b = pa // BLOCK
	return ((b[:, 0] + b[:, 1]) % 2) == 0


	def score_method(d: dict, method: str, reps: int = REPS, n_list=N_LIST, verbose: bool = True) -> dict:
	ya, pa, prior = d["_ya"], d["pa"], d["prior"]
	inA = half_A(pa)
	rows = []
	for n, rep in draws(reps, n_list):
	if not (have_pred(method, n, rep) or have_pred(method, 0, 0)):
	continue
	P_all = load_pred(method, n, rep)
	shot = shots_for(d, n, rep)
	keep = np.setdiff1d(np.arange(len(ya)), shot)
	P, y = P_all[keep], ya[keep]
	lab = {
	"argmax": CLASSES[P.argmax(1)],
	"rule": W.decide(P, prior),
	"rule_smooth": W.decide(smooth(P, pa, keep=keep), prior),
	}
	row = {"n": n, "rep": rep}
	for k, l in lab.items():
	row[k] = W.macro_f1(y, l)
	row[k + "_A"] = W.macro_f1(y[inA[keep]], l[inA[keep]])
	row[k + "_B"] = W.macro_f1(y[~inA[keep]], l[~inA[keep]])
	rows.append(row)
	if not rows:
	return {}
	df = pd.DataFrame(rows)
	summ = df.groupby("n").mean(numeric_only=True).drop(columns="rep")
	std = df.groupby("n").std(numeric_only=True).drop(columns="rep")
	out = {"method": method, "draws": rows, "n_draws": int(len(rows)),
	"mean": {c: {int(n): float(v) for n, v in summ[c].items()} for c in summ.columns},
	"std": {c: {int(n): float(v) for n, v in std[c].items()} for c in std.columns}}
	(SCORES / f"{method}.json").write_text(json.dumps(out, indent=1))
	if verbose:
	for k in ("argmax", "rule", "rule_smooth"):
	m = summ[k]
	print(f"{method:<34s} {k:<12s} " + " ".join(f"n{n}={m[n]:.4f}" for n in m.index)
	+ f" avg={m.mean():.4f} peak={m.max():.4f} ({len(rows)} draws)", flush=True)
	return out


	def board(metric_suffix: str = "") -> pd.DataFrame:
	"""All scored methods, one row per (method, decision). suffix '' / '_A' / '_B'."""
	rows = []
	for f in sorted(SCORES.glob("*.json")):
	s = json.loads(f.read_text())
	for k in ("argmax", "rule", "rule_smooth"):
	m = s["mean"].get(k + metric_suffix)
	if not m:
	continue
	r = {"method": s["method"], "decision": k, **{f"n{n}": m.get(str(n), np.nan) for n in N_LIST}}
	v = [r[f"n{n}"] for n in N_LIST]
	r["avg"], r["peak"], r["draws"] = float(np.nanmean(v)), float(np.nanmax(v)), s["n_draws"]
	rows.append(r)
	df = pd.DataFrame(rows).sort_values("avg", ascending=False)
	df.to_csv(OUT / f"leaderboard{metric_suffix}.csv", index=False)
	return df


	if __name__ == "__main__":
	cmd = sys.argv[1] if len(sys.argv) > 1 else "board"
	if cmd == "score":
	d = load(official60=False)
	names = sys.argv[2:] or sorted(p.name for p in PRED.iterdir() if p.is_dir())
	t0 = time.time()
	for m in names:
	score_method(d, m)
	print(f"scored {len(names)} methods in {time.time() - t0:.0f}s")
	else:
	suffix = sys.argv[2] if len(sys.argv) > 2 else ""
	pd.set_option("display.width", 250)
	print(board(suffix).round(4).to_string(index=False))