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

2c2dd6c verified 4 days ago

19.3 kB

"""
metric.py - METRIC LEARNING family: an embedding trained on Madrid so that distance
means "same age class", then adapted per protocol draw with the n Amsterdam shots.
    python tools/nnlab/metric.py [--smoke] [--only <substring>]
Encoders (MLP 393-256-256-64, BN/ReLU/Dropout 0.2, AdamW lr 1e-3 wd 1e-4, 2 seeds):
    supcon   supervised contrastive (T = 0.1), class-balanced batches of 512
    proto    prototypical episodes (4-way, k in {5, 25}, 50 queries/class, scale 10),
             cross-block: support from one random half of the 2 km blocks, queries from the other
    protoda  proto + deep-CORAL (weight 1) to an unlabelled Amsterdam batch every step
    cos      Baseline++ cosine classifier (scale 16), class-balanced batches
Heads (encoder frozen, CL2N per city, seeds' probabilities averaged):
    pn, maha, tpn, knn, ftproj   ->  methods met_{encoder}_{head}
Training length per encoder is chosen ONCE on Madrid (lab.madrid_split pseudo-target,
25-shot prototype macro F1, candidates evaluated as checkpoints of one constant-LR run),
then the encoder is retrained on all Madrid. Amsterdam labels: only lab.shot_labels.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import lab                                    # noqa: E402
import numpy as np                            # noqa: E402

T0 = time.time()


def log(msg):
    print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True)


ap = argparse.ArgumentParser()
ap.add_argument("--smoke", action="store_true")
ap.add_argument("--only", default=None)
args = ap.parse_args()
SMOKE = args.smoke
SFX = "_SMOKE" if SMOKE else ""

ENCODERS = ["supcon", "proto", "protoda", "cos"]
HEADS = ["pn", "maha", "tpn", "knn", "ftproj"]

# configuration ---------------------------------------------------------------- #
# candidate training lengths (optimizer steps; one step = one batch / one episode)
# (sized to the measured ~0.09-0.2 s/step on the shared CPU so the full run fits ~60 min;
#  1500 steps x 512 = ~10 Madrid epochs)
CAND = {"supcon": [300, 700, 1500], "cos": [300, 700, 1500],
        "proto": [400, 800, 1600], "protoda": [400, 800, 1600]}
if SMOKE:
    CAND = {k: [15, 30] for k in CAND}
SEEDS = [0] if SMOKE else [0, 1]
FT_SEEDS = 1                         # ftproj uses only the first seed's encoder (time budget)
SEL_SEED = 100                       # seed of the Madrid-only selection run
SEL_DRAWS = 2 if SMOKE else 5        # pseudo-target 25-shot draws per checkpoint
DRAWS = [(5, 0), (200, 0)] if SMOKE else list(lab.draws())
BATCH = 512
TEMP = 0.1
PROTO_SCALE = 10.0
PROTO_K = (5, 25)
PROTO_Q = 50
CORAL_W = 1.0
DA_BATCH = 300                       # Amsterdam batch for CORAL ~ size of a Madrid episode
COS_SCALE = 16.0
LR, WD = 1e-3, 1e-4
# heads
K_MIX, S_PN = 5.0, 20.0
SHRINK, RIDGE = 50.0, 1e-3
TPN_IT, TPN_CAP = 10, 5.0
KNN_K = 15
FT_STEPS, FT_LR, FT_L2SP = 100, 1e-3, 1e-2

MC = lab.CACHE / "metric"
MC.mkdir(parents=True, exist_ok=True)


def mname(enc, head):
    return f"met_{enc}_{head}{SFX}"


WANT = {e: [h for h in HEADS if args.only is None or args.only in f"met_{e}_{h}"] for e in ENCODERS}
WANT = {e: hs for e, hs in WANT.items() if hs}

# data ------------------------------------------------------------------------- #
d = lab.load(official60=False)
ym, pm = d["ym"], d["pm"]
Xm = lab.robustz(d["Xm"])
Xa = lab.robustz(d["Xa"])
del d["Xa"]                         # raw Madrid kept for the selection split
nM, nA = len(Xm), len(Xa)
log(f"loaded Madrid {Xm.shape} Amsterdam {Xa.shape}; encoders {list(WANT)}")

torch = lab.torch_setup(0)
nn = torch.nn
F = torch.nn.functional


# --------------------------------------------------------------------------- #
# Encoder + losses
# --------------------------------------------------------------------------- #
class Enc(nn.Module):
    def __init__(self, din=393):
        super().__init__()
        self.body = nn.Sequential(
            nn.Linear(din, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(256, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2))
        self.proj = nn.Linear(256, 64)

    def forward(self, x):
        return self.proj(self.body(x))


def supcon_loss(z, y, T=TEMP):
    z = F.normalize(z, dim=1)
    sim = z @ z.T / T
    n = len(y)
    eye = torch.eye(n, dtype=torch.bool)
    sim = sim.masked_fill(eye, -1e9)
    logp = sim - torch.logsumexp(sim, 1, keepdim=True)
    pos = (y[:, None] == y[None, :]) & ~eye
    return -((logp * pos).sum(1) / pos.sum(1).clamp(min=1)).mean()


def proto_loss(zs, ys, zq, yq, scale=PROTO_SCALE, ncls=4):
    zs, zq = F.normalize(zs, dim=1), F.normalize(zq, dim=1)
    P = torch.stack([zs[ys == c].mean(0) for c in range(ncls)])
    d2 = ((zq[:, None, :] - P[None]) ** 2).sum(-1)
    return F.cross_entropy(-scale * d2, yq)


def coral(a, b):
    def cov(x):
        x = x - x.mean(0, keepdim=True)
        return x.T @ x / (len(x) - 1)
    dd = a.shape[1]
    return ((cov(a) - cov(b)) ** 2).sum() / (4 * dd * dd)


class Sampler:
    """Class-balanced batches and cross-block episodes over one labelled pool."""

    def __init__(self, y, p, rng):
        self.y, self.rng = y, rng
        self.cls = [np.flatnonzero(y == c) for c in lab.CLASSES]
        b = lab.block_ids(p)
        _, self.code = np.unique(b, return_inverse=True)
        self.nb = self.code.max() + 1
        self.ccode = [self.code[ix] for ix in self.cls]

    def batch(self, size=BATCH):
        m = size // 4
        return np.concatenate([ix[self.rng.integers(0, len(ix), m)] for ix in self.cls])

    def episode(self, k=None):
        """Cross-block episode. Sampling is with replacement inside pools of thousands of
        pixels (duplicates are rare; it avoids a full permutation per class per step)."""
        k = k or int(self.rng.choice(PROTO_K))
        half = self.rng.random(self.nb) < 0.5
        S, Q = [], []
        for ix, cc in zip(self.cls, self.ccode):
            m = half[cc]
            a, b = ix[m], ix[~m]
            if len(a) < k or len(b) < PROTO_Q:     # degenerate split: fall back to any pixel
                a = b = ix
            S.append(a[self.rng.integers(0, len(a), k)])
            Q.append(b[self.rng.integers(0, len(b), PROTO_Q)])
        ys = np.repeat(np.arange(4), k)
        yq = np.repeat(np.arange(4), PROTO_Q)
        return np.concatenate(S), ys, np.concatenate(Q), yq


@torch.no_grad()
def embed_H(model, X, bs=8192):
    model.eval()
    out = [model.body(torch.from_numpy(X[i:i + bs])).numpy() for i in range(0, len(X), bs)]
    return np.vstack(out).astype(np.float32)


def train_encoder(kind, X, y, p, Xt, lengths, seed, eval_fn=None):
    """Constant-LR training; returns (model, {length: eval score}) with checkpoints at
    `lengths` (the model returned is the one at max(lengths))."""
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    model = Enc(X.shape[1])
    params = list(model.parameters())
    Wc = None
    if kind == "cos":
        Wc = nn.Parameter(torch.randn(4, 64) * 0.1)
        params.append(Wc)
    opt = torch.optim.AdamW(params, lr=LR, weight_decay=WD)
    smp = Sampler(y, p, rng)
    Xt_t = torch.from_numpy(X)
    yt = torch.from_numpy(y - 1)
    scores = {}
    t0 = time.time()
    for step in range(1, max(lengths) + 1):
        model.train()
        if kind in ("supcon", "cos"):
            ix = smp.batch()
            z = model(Xt_t[ix])
            if kind == "supcon":
                loss = supcon_loss(z, yt[ix])
            else:
                logits = COS_SCALE * F.normalize(z, dim=1) @ F.normalize(Wc, dim=1).T
                loss = F.cross_entropy(logits, yt[ix])
        else:
            S, ys, Q, yq = smp.episode()
            z = model(Xt_t[np.concatenate([S, Q])])
            zs, zq = z[:len(S)], z[len(S):]
            loss = proto_loss(zs, torch.from_numpy(ys), zq, torch.from_numpy(yq))
            if kind == "protoda":
                it = rng.integers(0, len(Xt), DA_BATCH)
                zt = model(torch.from_numpy(Xt[it]))
                loss = loss + CORAL_W * coral(z, zt)
        opt.zero_grad()
        loss.backward()
        opt.step()
        if step in lengths:
            if eval_fn is not None:
                scores[step] = eval_fn(model)
                log(f"  {kind} seed{seed} step {step}: loss {loss.item():.4f} "
                    f"pseudo-target F1 {scores[step]:.4f} ({time.time() - t0:.0f}s)")
            else:
                log(f"  {kind} seed{seed} step {step}: loss {loss.item():.4f} ({time.time() - t0:.0f}s)")
    return model, scores


# --------------------------------------------------------------------------- #
# Embedding utilities
# --------------------------------------------------------------------------- #
def l2n(Z):
    return Z / np.maximum(np.linalg.norm(Z, axis=1, keepdims=True), 1e-12)


def cl2n(Z):
    """SimpleShot CL2N with this city's own mean (label-free)."""
    E = l2n(np.asarray(Z, np.float64))
    return l2n(E - E.mean(0))


def macro_f1(y, yhat):
    from sklearn.metrics import f1_score
    return float(f1_score(y, yhat, average="macro"))


def proto_f1_eval(Xpt, ypt, ndraw=SEL_DRAWS, shots=25):
    """Pseudo-target score: 25-shot pure prototype classifier on CL2N embeddings."""
    rs = np.random.default_rng(7)
    draws_ = []
    for _ in range(ndraw):
        sh = np.concatenate([rs.choice(np.flatnonzero(ypt == c), shots, replace=False) for c in lab.CLASSES])
        draws_.append(sh)

    def fn(model):
        H = embed_H(model, Xpt)
        with torch.no_grad():
            Z = model.proj(torch.from_numpy(H)).numpy()
        E = cl2n(Z)
        out = []
        for sh in draws_:
            P = np.stack([E[sh][ypt[sh] == c].mean(0) for c in lab.CLASSES])
            keep = np.setdiff1d(np.arange(len(E)), sh)
            pred = lab.CLASSES[(-((E[keep, None, :] - P[None]) ** 2).sum(-1)).argmax(1)]
            out.append(macro_f1(ypt[keep], pred))
        return float(np.mean(out))
    return fn


def select_length(kind):
    p = MC / f"{kind}_select{SFX}.json"
    if p.exists():
        return json.loads(p.read_text())
    src, pt = lab.madrid_split(pm)
    Xs, Xpt = lab.robustz(d["Xm"][src]), lab.robustz(d["Xm"][pt])
    log(f"{kind}: selecting length on Madrid split (src {len(src)}, pseudo-target {len(pt)}) "
        f"from {CAND[kind]}")
    # protoda: the unlabelled 'target' of the selection run is the pseudo-target itself
    _, sc = train_encoder(kind, Xs, ym[src], pm[src], Xpt, CAND[kind], SEL_SEED,
                          eval_fn=proto_f1_eval(Xpt, ym[pt]))
    best = max(sc, key=lambda s: (sc[s], -s))
    res = {"candidates": CAND[kind], "pseudo_target_f1": {str(k): v for k, v in sc.items()},
           "chosen": int(best), "seed": SEL_SEED, "split": "lab.madrid_split(pm) default",
           "da_target_in_selection": "pseudo-target (unlabelled)" if kind == "protoda" else None}
    p.write_text(json.dumps(res, indent=1))
    log(f"{kind}: chose {best} steps")
    return res


def encoder_cache(kind, seed, length):
    p = MC / f"{kind}_s{seed}_L{length}{SFX}.npz"
    if p.exists():
        z = np.load(p)
        return {k: z[k] for k in z.files}
    log(f"{kind}: training seed {seed} on all Madrid for {length} steps")
    model, _ = train_encoder(kind, Xm, ym, pm, Xa, [length], seed)
    out = dict(Hm=embed_H(model, Xm), Ha=embed_H(model, Xa),
               W=model.proj.weight.detach().numpy().copy(), b=model.proj.bias.detach().numpy().copy())
    np.savez(p, **out)
    return out


# --------------------------------------------------------------------------- #
# Heads
# --------------------------------------------------------------------------- #
def class_means(E, y):
    return np.stack([E[y == c].mean(0) for c in lab.CLASSES])


def within_cov(E, y):
    R = np.vstack([E[y == c] - E[y == c].mean(0) for c in lab.CLASSES])
    return R.T @ R / max(len(R) - 4, 1)


def pn_protos(Ea, sh, ysh, mu_m):
    P = []
    for j, c in enumerate(lab.CLASSES):
        S = Ea[sh[ysh == c]]
        P.append((len(S) * S.mean(0) + K_MIX * mu_m[j]) / (len(S) + K_MIX))
    return np.stack(P)


def sqd(E, P):
    return (E ** 2).sum(1)[:, None] - 2 * E @ P.T + (P ** 2).sum(1)[None]


def head_pn(Ea, sh, ysh, mu_m):
    return lab.softmax(-S_PN * sqd(Ea, pn_protos(Ea, sh, ysh, mu_m)))


def head_maha(Ea, sh, ysh, mu_m, S_m):
    from scipy.linalg import cho_factor, cho_solve
    P = pn_protos(Ea, sh, ysh, mu_m)
    nt = len(sh)
    S = (nt * within_cov(Ea[sh], ysh) + SHRINK * S_m) / (nt + SHRINK) + RIDGE * np.eye(Ea.shape[1])
    cf = cho_factor(S)
    D = np.empty((len(Ea), 4))
    for j in range(4):
        R = Ea - P[j]
        D[:, j] = (R * cho_solve(cf, R.T).T).sum(1)
    return lab.softmax(-0.5 * D)


def head_tpn(Ea, sh, ysh, mu_m):
    P = pn_protos(Ea, sh, ysh, mu_m)
    unl = np.setdiff1d(np.arange(len(Ea)), sh)
    Eu = Ea[unl]
    sums = np.stack([Ea[sh[ysh == c]].sum(0) for c in lab.CLASSES])
    nc = np.array([(ysh == c).sum() for c in lab.CLASSES], float)
    for _ in range(TPN_IT):
        R = lab.softmax(-S_PN * sqd(Eu, P))           # soft assignments of the unlabelled pool
        mass = R.sum(0)
        w = np.minimum(1.0, TPN_CAP * nc / np.maximum(mass, 1e-12))
        P = (sums + w[:, None] * (R.T @ Eu)) / (nc + w * mass)[:, None]
    return lab.softmax(-S_PN * sqd(Ea, P))


def head_knn(Ea, sh, ysh, mu_m, n):
    if n < 25:
        return head_pn(Ea, sh, ysh, mu_m)
    Es = Ea[sh]
    C = Ea @ Es.T                                       # cosine (unit vectors)
    k = min(KNN_K, len(sh))
    nb = np.argpartition(-C, k - 1, axis=1)[:, :k]
    cs = np.take_along_axis(C, nb, 1)
    w = 1.0 / (1.0 - cs + 1e-3)                         # inverse cosine distance
    lab_nb = ysh[nb] - 1
    P = np.zeros((len(Ea), 4))
    for j in range(4):
        P[:, j] = (w * (lab_nb == j)).sum(1)
    P += 1e-3
    return P / P.sum(1, keepdims=True)


def head_ftproj(enc, Hm, Ha, W0, b0, sh, ysh, n, seed):
    """Fine-tune ONLY the projection Linear(256,64) on cached frozen activations."""
    torch.manual_seed(seed + 1000 * n)
    rng = np.random.default_rng(seed + 1000 * n)
    Wt = torch.tensor(W0, requires_grad=True)
    bt = torch.tensor(b0, requires_grad=True)
    W0t, b0t = torch.tensor(W0), torch.tensor(b0)
    opt = torch.optim.Adam([Wt, bt], lr=FT_LR)
    Hs = torch.from_numpy(Ha[sh])
    Hm_t = torch.from_numpy(Hm)
    smp = MAD_SMP
    smp.rng = rng
    cls_ix = [np.flatnonzero(ysh == c) for c in lab.CLASSES]
    for _ in range(FT_STEPS):
        S, Q = [], []
        for ix in cls_ix:
            pr = rng.permutation(ix)
            ns = 3 if len(ix) <= 5 else len(ix) // 2
            S.append(pr[:ns]); Q.append(pr[ns:])
        ys = torch.from_numpy(np.concatenate([np.full(len(s), j) for j, s in enumerate(S)]))
        yq = torch.from_numpy(np.concatenate([np.full(len(q), j) for j, q in enumerate(Q)]))
        S, Q = np.concatenate(S), np.concatenate(Q)
        loss = proto_loss(Hs[S] @ Wt.T + bt, ys, Hs[Q] @ Wt.T + bt, yq)
        mS, mys, mQ, myq = smp.episode()
        loss = loss + proto_loss(Hm_t[mS] @ Wt.T + bt, torch.from_numpy(mys),
                                 Hm_t[mQ] @ Wt.T + bt, torch.from_numpy(myq))
        loss = loss + FT_L2SP * (((Wt - W0t) ** 2).sum() + ((bt - b0t) ** 2).sum())
        opt.zero_grad()
        loss.backward()
        opt.step()
    W, b = Wt.detach().numpy(), bt.detach().numpy()
    Em, Ea = cl2n(Hm @ W.T + b), cl2n(Ha @ W.T + b)
    return head_pn(Ea, sh, ysh, class_means(Em, ym))


# --------------------------------------------------------------------------- #
# Main loop
# --------------------------------------------------------------------------- #
def meta_for(enc, head, sel):
    return dict(family="metric", encoder=enc, head=head, smoke=SMOKE,
                seeds=SEEDS[:FT_SEEDS] if head == "ftproj" else SEEDS,
                input="lab.robustz per city (clip 10)",
                encoder_arch="393-256-BN-ReLU-Do0.2-256-BN-ReLU-Do0.2-64, L2-normalised",
                optim=f"AdamW lr {LR} wd {WD}, constant LR",
                encoder_cfg=dict(batch=BATCH, supcon_T=TEMP, proto_scale=PROTO_SCALE, proto_k=PROTO_K,
                                 proto_q=PROTO_Q, cross_block="2 km blocks random halves per episode",
                                 coral_w=CORAL_W, da_batch=DA_BATCH, cos_scale=COS_SCALE,
                                 cos_balance="class-balanced batches (128 per class), plain CE"),
                length_selection=sel,
                madrid_choices=f"training steps chosen on lab.madrid_split pseudo-target: {sel['chosen']}",
                head_cfg=dict(CL2N="per-city mean, re-normalised", k_mix=K_MIX, s=S_PN, shrink=SHRINK,
                              ridge=RIDGE, tpn_iters=TPN_IT, tpn_cap=TPN_CAP, knn_k=KNN_K,
                              knn_weight="1/(1-cos+1e-3)", knn_n5="pn fallback",
                              ft_steps=FT_STEPS, ft_lr=FT_LR, ft_l2sp=FT_L2SP,
                              ft_episode="shots half/half (n=5: 3/2) + one Madrid cross-block episode"),
                hyperparams_from_amsterdam="none")


MAD_SMP = None


def main():
    global MAD_SMP
    MAD_SMP = Sampler(ym, pm, np.random.default_rng(0))
    for enc, heads in WANT.items():
        todo = [(n, r) for n, r in DRAWS if not all(lab.have_pred(mname(enc, h), n, r) for h in heads)]
        if not todo:
            log(f"{enc}: all predictions exist, skipping")
            continue
        sel = select_length(enc)
        L = sel["chosen"]
        caches = [encoder_cache(enc, s, L) for s in SEEDS]
        prep = []
        for c in caches:
            Zm, Za = c["Hm"] @ c["W"].T + c["b"], c["Ha"] @ c["W"].T + c["b"]
            Em, Ea = cl2n(Zm), cl2n(Za)
            prep.append(dict(Ea=Ea, mu=class_means(Em, ym), Sm=within_cov(Em, ym)))
            del Em
        metas = {h: meta_for(enc, h, sel) for h in heads}
        for n, r in todo:
            t0 = time.time()
            sh = lab.shots_for(d, n, r)
            ysh = lab.shot_labels(d, sh)
            for h in heads:
                if lab.have_pred(mname(enc, h), n, r):
                    continue
                Ps = []
                for si, (c, pp) in enumerate(zip(caches, prep)):
                    if h == "ftproj" and si >= FT_SEEDS:
                        break
                    Ea, mu = pp["Ea"], pp["mu"]
                    if h == "pn":
                        P = head_pn(Ea, sh, ysh, mu)
                    elif h == "maha":
                        P = head_maha(Ea, sh, ysh, mu, pp["Sm"])
                    elif h == "tpn":
                        P = head_tpn(Ea, sh, ysh, mu)
                    elif h == "knn":
                        P = head_knn(Ea, sh, ysh, mu, n)
                    else:
                        P = head_ftproj(enc, c["Hm"], c["Ha"], c["W"], c["b"], sh, ysh, n, SEEDS[si] + 10 * r)
                    Ps.append(P)
                lab.save_pred(mname(enc, h), n, r, np.mean(Ps, 0), metas[h])
            log(f"{enc} n={n} rep={r}: saved {len(heads)} heads ({time.time() - t0:.1f}s)")
        del caches, prep
    log("done")


if __name__ == "__main__":
    main()

	"""
	metric.py - METRIC LEARNING family: an embedding trained on Madrid so that distance
	means "same age class", then adapted per protocol draw with the n Amsterdam shots.

	python tools/nnlab/metric.py [--smoke] [--only <substring>]

	Encoders (MLP 393-256-256-64, BN/ReLU/Dropout 0.2, AdamW lr 1e-3 wd 1e-4, 2 seeds):
	supcon supervised contrastive (T = 0.1), class-balanced batches of 512
	proto prototypical episodes (4-way, k in {5, 25}, 50 queries/class, scale 10),
	cross-block: support from one random half of the 2 km blocks, queries from the other
	protoda proto + deep-CORAL (weight 1) to an unlabelled Amsterdam batch every step
	cos Baseline++ cosine classifier (scale 16), class-balanced batches
	Heads (encoder frozen, CL2N per city, seeds' probabilities averaged):
	pn, maha, tpn, knn, ftproj -> methods met_{encoder}_{head}

	Training length per encoder is chosen ONCE on Madrid (lab.madrid_split pseudo-target,
	25-shot prototype macro F1, candidates evaluated as checkpoints of one constant-LR run),
	then the encoder is retrained on all Madrid. Amsterdam labels: only lab.shot_labels.
	"""
	from __future__ import annotations

	import argparse
	import json
	import sys
	import time
	from pathlib import Path

	sys.path.insert(0, str(Path(__file__).resolve().parent))
	import lab # noqa: E402
	import numpy as np # noqa: E402

	T0 = time.time()


	def log(msg):
	print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True)


	ap = argparse.ArgumentParser()
	ap.add_argument("--smoke", action="store_true")
	ap.add_argument("--only", default=None)
	args = ap.parse_args()
	SMOKE = args.smoke
	SFX = "_SMOKE" if SMOKE else ""

	ENCODERS = ["supcon", "proto", "protoda", "cos"]
	HEADS = ["pn", "maha", "tpn", "knn", "ftproj"]

	# configuration ---------------------------------------------------------------- #
	# candidate training lengths (optimizer steps; one step = one batch / one episode)
	# (sized to the measured ~0.09-0.2 s/step on the shared CPU so the full run fits ~60 min;
	# 1500 steps x 512 = ~10 Madrid epochs)
	CAND = {"supcon": [300, 700, 1500], "cos": [300, 700, 1500],
	"proto": [400, 800, 1600], "protoda": [400, 800, 1600]}
	if SMOKE:
	CAND = {k: [15, 30] for k in CAND}
	SEEDS = [0] if SMOKE else [0, 1]
	FT_SEEDS = 1 # ftproj uses only the first seed's encoder (time budget)
	SEL_SEED = 100 # seed of the Madrid-only selection run
	SEL_DRAWS = 2 if SMOKE else 5 # pseudo-target 25-shot draws per checkpoint
	DRAWS = [(5, 0), (200, 0)] if SMOKE else list(lab.draws())
	BATCH = 512
	TEMP = 0.1
	PROTO_SCALE = 10.0
	PROTO_K = (5, 25)
	PROTO_Q = 50
	CORAL_W = 1.0
	DA_BATCH = 300 # Amsterdam batch for CORAL ~ size of a Madrid episode
	COS_SCALE = 16.0
	LR, WD = 1e-3, 1e-4
	# heads
	K_MIX, S_PN = 5.0, 20.0
	SHRINK, RIDGE = 50.0, 1e-3
	TPN_IT, TPN_CAP = 10, 5.0
	KNN_K = 15
	FT_STEPS, FT_LR, FT_L2SP = 100, 1e-3, 1e-2

	MC = lab.CACHE / "metric"
	MC.mkdir(parents=True, exist_ok=True)


	def mname(enc, head):
	return f"met_{enc}_{head}{SFX}"


	WANT = {e: [h for h in HEADS if args.only is None or args.only in f"met_{e}_{h}"] for e in ENCODERS}
	WANT = {e: hs for e, hs in WANT.items() if hs}

	# data ------------------------------------------------------------------------- #
	d = lab.load(official60=False)
	ym, pm = d["ym"], d["pm"]
	Xm = lab.robustz(d["Xm"])
	Xa = lab.robustz(d["Xa"])
	del d["Xa"] # raw Madrid kept for the selection split
	nM, nA = len(Xm), len(Xa)
	log(f"loaded Madrid {Xm.shape} Amsterdam {Xa.shape}; encoders {list(WANT)}")

	torch = lab.torch_setup(0)
	nn = torch.nn
	F = torch.nn.functional


	# --------------------------------------------------------------------------- #
	# Encoder + losses
	# --------------------------------------------------------------------------- #
	class Enc(nn.Module):
	def __init__(self, din=393):
	super().__init__()
	self.body = nn.Sequential(
	nn.Linear(din, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
	nn.Linear(256, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2))
	self.proj = nn.Linear(256, 64)

	def forward(self, x):
	return self.proj(self.body(x))


	def supcon_loss(z, y, T=TEMP):
	z = F.normalize(z, dim=1)
	sim = z @ z.T / T
	n = len(y)
	eye = torch.eye(n, dtype=torch.bool)
	sim = sim.masked_fill(eye, -1e9)
	logp = sim - torch.logsumexp(sim, 1, keepdim=True)
	pos = (y[:, None] == y[None, :]) & ~eye
	return -((logp * pos).sum(1) / pos.sum(1).clamp(min=1)).mean()


	def proto_loss(zs, ys, zq, yq, scale=PROTO_SCALE, ncls=4):
	zs, zq = F.normalize(zs, dim=1), F.normalize(zq, dim=1)
	P = torch.stack([zs[ys == c].mean(0) for c in range(ncls)])
	d2 = ((zq[:, None, :] - P[None]) ** 2).sum(-1)
	return F.cross_entropy(-scale * d2, yq)


	def coral(a, b):
	def cov(x):
	x = x - x.mean(0, keepdim=True)
	return x.T @ x / (len(x) - 1)
	dd = a.shape[1]
	return ((cov(a) - cov(b)) ** 2).sum() / (4 * dd * dd)


	class Sampler:
	"""Class-balanced batches and cross-block episodes over one labelled pool."""

	def __init__(self, y, p, rng):
	self.y, self.rng = y, rng
	self.cls = [np.flatnonzero(y == c) for c in lab.CLASSES]
	b = lab.block_ids(p)
	_, self.code = np.unique(b, return_inverse=True)
	self.nb = self.code.max() + 1
	self.ccode = [self.code[ix] for ix in self.cls]

	def batch(self, size=BATCH):
	m = size // 4
	return np.concatenate([ix[self.rng.integers(0, len(ix), m)] for ix in self.cls])

	def episode(self, k=None):
	"""Cross-block episode. Sampling is with replacement inside pools of thousands of
	pixels (duplicates are rare; it avoids a full permutation per class per step)."""
	k = k or int(self.rng.choice(PROTO_K))
	half = self.rng.random(self.nb) < 0.5
	S, Q = [], []
	for ix, cc in zip(self.cls, self.ccode):
	m = half[cc]
	a, b = ix[m], ix[~m]
	if len(a) < k or len(b) < PROTO_Q: # degenerate split: fall back to any pixel
	a = b = ix
	S.append(a[self.rng.integers(0, len(a), k)])
	Q.append(b[self.rng.integers(0, len(b), PROTO_Q)])
	ys = np.repeat(np.arange(4), k)
	yq = np.repeat(np.arange(4), PROTO_Q)
	return np.concatenate(S), ys, np.concatenate(Q), yq


	@torch.no_grad()
	def embed_H(model, X, bs=8192):
	model.eval()
	out = [model.body(torch.from_numpy(X[i:i + bs])).numpy() for i in range(0, len(X), bs)]
	return np.vstack(out).astype(np.float32)


	def train_encoder(kind, X, y, p, Xt, lengths, seed, eval_fn=None):
	"""Constant-LR training; returns (model, {length: eval score}) with checkpoints at
	`lengths` (the model returned is the one at max(lengths))."""
	torch.manual_seed(seed)
	rng = np.random.default_rng(seed)
	model = Enc(X.shape[1])
	params = list(model.parameters())
	Wc = None
	if kind == "cos":
	Wc = nn.Parameter(torch.randn(4, 64) * 0.1)
	params.append(Wc)
	opt = torch.optim.AdamW(params, lr=LR, weight_decay=WD)
	smp = Sampler(y, p, rng)
	Xt_t = torch.from_numpy(X)
	yt = torch.from_numpy(y - 1)
	scores = {}
	t0 = time.time()
	for step in range(1, max(lengths) + 1):
	model.train()
	if kind in ("supcon", "cos"):
	ix = smp.batch()
	z = model(Xt_t[ix])
	if kind == "supcon":
	loss = supcon_loss(z, yt[ix])
	else:
	logits = COS_SCALE * F.normalize(z, dim=1) @ F.normalize(Wc, dim=1).T
	loss = F.cross_entropy(logits, yt[ix])
	else:
	S, ys, Q, yq = smp.episode()
	z = model(Xt_t[np.concatenate([S, Q])])
	zs, zq = z[:len(S)], z[len(S):]
	loss = proto_loss(zs, torch.from_numpy(ys), zq, torch.from_numpy(yq))
	if kind == "protoda":
	it = rng.integers(0, len(Xt), DA_BATCH)
	zt = model(torch.from_numpy(Xt[it]))
	loss = loss + CORAL_W * coral(z, zt)
	opt.zero_grad()
	loss.backward()
	opt.step()
	if step in lengths:
	if eval_fn is not None:
	scores[step] = eval_fn(model)
	log(f" {kind} seed{seed} step {step}: loss {loss.item():.4f} "
	f"pseudo-target F1 {scores[step]:.4f} ({time.time() - t0:.0f}s)")
	else:
	log(f" {kind} seed{seed} step {step}: loss {loss.item():.4f} ({time.time() - t0:.0f}s)")
	return model, scores


	# --------------------------------------------------------------------------- #
	# Embedding utilities
	# --------------------------------------------------------------------------- #
	def l2n(Z):
	return Z / np.maximum(np.linalg.norm(Z, axis=1, keepdims=True), 1e-12)


	def cl2n(Z):
	"""SimpleShot CL2N with this city's own mean (label-free)."""
	E = l2n(np.asarray(Z, np.float64))
	return l2n(E - E.mean(0))


	def macro_f1(y, yhat):
	from sklearn.metrics import f1_score
	return float(f1_score(y, yhat, average="macro"))


	def proto_f1_eval(Xpt, ypt, ndraw=SEL_DRAWS, shots=25):
	"""Pseudo-target score: 25-shot pure prototype classifier on CL2N embeddings."""
	rs = np.random.default_rng(7)
	draws_ = []
	for _ in range(ndraw):
	sh = np.concatenate([rs.choice(np.flatnonzero(ypt == c), shots, replace=False) for c in lab.CLASSES])
	draws_.append(sh)

	def fn(model):
	H = embed_H(model, Xpt)
	with torch.no_grad():
	Z = model.proj(torch.from_numpy(H)).numpy()
	E = cl2n(Z)
	out = []
	for sh in draws_:
	P = np.stack([E[sh][ypt[sh] == c].mean(0) for c in lab.CLASSES])
	keep = np.setdiff1d(np.arange(len(E)), sh)
	pred = lab.CLASSES[(-((E[keep, None, :] - P[None]) ** 2).sum(-1)).argmax(1)]
	out.append(macro_f1(ypt[keep], pred))
	return float(np.mean(out))
	return fn


	def select_length(kind):
	p = MC / f"{kind}_select{SFX}.json"
	if p.exists():
	return json.loads(p.read_text())
	src, pt = lab.madrid_split(pm)
	Xs, Xpt = lab.robustz(d["Xm"][src]), lab.robustz(d["Xm"][pt])
	log(f"{kind}: selecting length on Madrid split (src {len(src)}, pseudo-target {len(pt)}) "
	f"from {CAND[kind]}")
	# protoda: the unlabelled 'target' of the selection run is the pseudo-target itself
	_, sc = train_encoder(kind, Xs, ym[src], pm[src], Xpt, CAND[kind], SEL_SEED,
	eval_fn=proto_f1_eval(Xpt, ym[pt]))
	best = max(sc, key=lambda s: (sc[s], -s))
	res = {"candidates": CAND[kind], "pseudo_target_f1": {str(k): v for k, v in sc.items()},
	"chosen": int(best), "seed": SEL_SEED, "split": "lab.madrid_split(pm) default",
	"da_target_in_selection": "pseudo-target (unlabelled)" if kind == "protoda" else None}
	p.write_text(json.dumps(res, indent=1))
	log(f"{kind}: chose {best} steps")
	return res


	def encoder_cache(kind, seed, length):
	p = MC / f"{kind}_s{seed}_L{length}{SFX}.npz"
	if p.exists():
	z = np.load(p)
	return {k: z[k] for k in z.files}
	log(f"{kind}: training seed {seed} on all Madrid for {length} steps")
	model, _ = train_encoder(kind, Xm, ym, pm, Xa, [length], seed)
	out = dict(Hm=embed_H(model, Xm), Ha=embed_H(model, Xa),
	W=model.proj.weight.detach().numpy().copy(), b=model.proj.bias.detach().numpy().copy())
	np.savez(p, **out)
	return out


	# --------------------------------------------------------------------------- #
	# Heads
	# --------------------------------------------------------------------------- #
	def class_means(E, y):
	return np.stack([E[y == c].mean(0) for c in lab.CLASSES])


	def within_cov(E, y):
	R = np.vstack([E[y == c] - E[y == c].mean(0) for c in lab.CLASSES])
	return R.T @ R / max(len(R) - 4, 1)


	def pn_protos(Ea, sh, ysh, mu_m):
	P = []
	for j, c in enumerate(lab.CLASSES):
	S = Ea[sh[ysh == c]]
	P.append((len(S) * S.mean(0) + K_MIX * mu_m[j]) / (len(S) + K_MIX))
	return np.stack(P)


	def sqd(E, P):
	return (E ** 2).sum(1)[:, None] - 2 * E @ P.T + (P ** 2).sum(1)[None]


	def head_pn(Ea, sh, ysh, mu_m):
	return lab.softmax(-S_PN * sqd(Ea, pn_protos(Ea, sh, ysh, mu_m)))


	def head_maha(Ea, sh, ysh, mu_m, S_m):
	from scipy.linalg import cho_factor, cho_solve
	P = pn_protos(Ea, sh, ysh, mu_m)
	nt = len(sh)
	S = (nt * within_cov(Ea[sh], ysh) + SHRINK * S_m) / (nt + SHRINK) + RIDGE * np.eye(Ea.shape[1])
	cf = cho_factor(S)
	D = np.empty((len(Ea), 4))
	for j in range(4):
	R = Ea - P[j]
	D[:, j] = (R * cho_solve(cf, R.T).T).sum(1)
	return lab.softmax(-0.5 * D)


	def head_tpn(Ea, sh, ysh, mu_m):
	P = pn_protos(Ea, sh, ysh, mu_m)
	unl = np.setdiff1d(np.arange(len(Ea)), sh)
	Eu = Ea[unl]
	sums = np.stack([Ea[sh[ysh == c]].sum(0) for c in lab.CLASSES])
	nc = np.array([(ysh == c).sum() for c in lab.CLASSES], float)
	for _ in range(TPN_IT):
	R = lab.softmax(-S_PN * sqd(Eu, P)) # soft assignments of the unlabelled pool
	mass = R.sum(0)
	w = np.minimum(1.0, TPN_CAP * nc / np.maximum(mass, 1e-12))
	P = (sums + w[:, None] * (R.T @ Eu)) / (nc + w * mass)[:, None]
	return lab.softmax(-S_PN * sqd(Ea, P))


	def head_knn(Ea, sh, ysh, mu_m, n):
	if n < 25:
	return head_pn(Ea, sh, ysh, mu_m)
	Es = Ea[sh]
	C = Ea @ Es.T # cosine (unit vectors)
	k = min(KNN_K, len(sh))
	nb = np.argpartition(-C, k - 1, axis=1)[:, :k]
	cs = np.take_along_axis(C, nb, 1)
	w = 1.0 / (1.0 - cs + 1e-3) # inverse cosine distance
	lab_nb = ysh[nb] - 1
	P = np.zeros((len(Ea), 4))
	for j in range(4):
	P[:, j] = (w * (lab_nb == j)).sum(1)
	P += 1e-3
	return P / P.sum(1, keepdims=True)


	def head_ftproj(enc, Hm, Ha, W0, b0, sh, ysh, n, seed):
	"""Fine-tune ONLY the projection Linear(256,64) on cached frozen activations."""
	torch.manual_seed(seed + 1000 * n)
	rng = np.random.default_rng(seed + 1000 * n)
	Wt = torch.tensor(W0, requires_grad=True)
	bt = torch.tensor(b0, requires_grad=True)
	W0t, b0t = torch.tensor(W0), torch.tensor(b0)
	opt = torch.optim.Adam([Wt, bt], lr=FT_LR)
	Hs = torch.from_numpy(Ha[sh])
	Hm_t = torch.from_numpy(Hm)
	smp = MAD_SMP
	smp.rng = rng
	cls_ix = [np.flatnonzero(ysh == c) for c in lab.CLASSES]
	for _ in range(FT_STEPS):
	S, Q = [], []
	for ix in cls_ix:
	pr = rng.permutation(ix)
	ns = 3 if len(ix) <= 5 else len(ix) // 2
	S.append(pr[:ns]); Q.append(pr[ns:])
	ys = torch.from_numpy(np.concatenate([np.full(len(s), j) for j, s in enumerate(S)]))
	yq = torch.from_numpy(np.concatenate([np.full(len(q), j) for j, q in enumerate(Q)]))
	S, Q = np.concatenate(S), np.concatenate(Q)
	loss = proto_loss(Hs[S] @ Wt.T + bt, ys, Hs[Q] @ Wt.T + bt, yq)
	mS, mys, mQ, myq = smp.episode()
	loss = loss + proto_loss(Hm_t[mS] @ Wt.T + bt, torch.from_numpy(mys),
	Hm_t[mQ] @ Wt.T + bt, torch.from_numpy(myq))
	loss = loss + FT_L2SP * (((Wt - W0t) 2).sum() + ((bt - b0t) 2).sum())
	opt.zero_grad()
	loss.backward()
	opt.step()
	W, b = Wt.detach().numpy(), bt.detach().numpy()
	Em, Ea = cl2n(Hm @ W.T + b), cl2n(Ha @ W.T + b)
	return head_pn(Ea, sh, ysh, class_means(Em, ym))


	# --------------------------------------------------------------------------- #
	# Main loop
	# --------------------------------------------------------------------------- #
	def meta_for(enc, head, sel):
	return dict(family="metric", encoder=enc, head=head, smoke=SMOKE,
	seeds=SEEDS[:FT_SEEDS] if head == "ftproj" else SEEDS,
	input="lab.robustz per city (clip 10)",
	encoder_arch="393-256-BN-ReLU-Do0.2-256-BN-ReLU-Do0.2-64, L2-normalised",
	optim=f"AdamW lr {LR} wd {WD}, constant LR",
	encoder_cfg=dict(batch=BATCH, supcon_T=TEMP, proto_scale=PROTO_SCALE, proto_k=PROTO_K,
	proto_q=PROTO_Q, cross_block="2 km blocks random halves per episode",
	coral_w=CORAL_W, da_batch=DA_BATCH, cos_scale=COS_SCALE,
	cos_balance="class-balanced batches (128 per class), plain CE"),
	length_selection=sel,
	madrid_choices=f"training steps chosen on lab.madrid_split pseudo-target: {sel['chosen']}",
	head_cfg=dict(CL2N="per-city mean, re-normalised", k_mix=K_MIX, s=S_PN, shrink=SHRINK,
	ridge=RIDGE, tpn_iters=TPN_IT, tpn_cap=TPN_CAP, knn_k=KNN_K,
	knn_weight="1/(1-cos+1e-3)", knn_n5="pn fallback",
	ft_steps=FT_STEPS, ft_lr=FT_LR, ft_l2sp=FT_L2SP,
	ft_episode="shots half/half (n=5: 3/2) + one Madrid cross-block episode"),
	hyperparams_from_amsterdam="none")


	MAD_SMP = None


	def main():
	global MAD_SMP
	MAD_SMP = Sampler(ym, pm, np.random.default_rng(0))
	for enc, heads in WANT.items():
	todo = [(n, r) for n, r in DRAWS if not all(lab.have_pred(mname(enc, h), n, r) for h in heads)]
	if not todo:
	log(f"{enc}: all predictions exist, skipping")
	continue
	sel = select_length(enc)
	L = sel["chosen"]
	caches = [encoder_cache(enc, s, L) for s in SEEDS]
	prep = []
	for c in caches:
	Zm, Za = c["Hm"] @ c["W"].T + c["b"], c["Ha"] @ c["W"].T + c["b"]
	Em, Ea = cl2n(Zm), cl2n(Za)
	prep.append(dict(Ea=Ea, mu=class_means(Em, ym), Sm=within_cov(Em, ym)))
	del Em
	metas = {h: meta_for(enc, h, sel) for h in heads}
	for n, r in todo:
	t0 = time.time()
	sh = lab.shots_for(d, n, r)
	ysh = lab.shot_labels(d, sh)
	for h in heads:
	if lab.have_pred(mname(enc, h), n, r):
	continue
	Ps = []
	for si, (c, pp) in enumerate(zip(caches, prep)):
	if h == "ftproj" and si >= FT_SEEDS:
	break
	Ea, mu = pp["Ea"], pp["mu"]
	if h == "pn":
	P = head_pn(Ea, sh, ysh, mu)
	elif h == "maha":
	P = head_maha(Ea, sh, ysh, mu, pp["Sm"])
	elif h == "tpn":
	P = head_tpn(Ea, sh, ysh, mu)
	elif h == "knn":
	P = head_knn(Ea, sh, ysh, mu, n)
	else:
	P = head_ftproj(enc, c["Hm"], c["Ha"], c["W"], c["b"], sh, ysh, n, SEEDS[si] + 10 * r)
	Ps.append(P)
	lab.save_pred(mname(enc, h), n, r, np.mean(Ps, 0), metas[h])
	log(f"{enc} n={n} rep={r}: saved {len(heads)} heads ({time.time() - t0:.1f}s)")
	del caches, prep
	log("done")


	if __name__ == "__main__":
	main()