Scenarios

prog.start("scenario table")
TIERS = np.arange(1, K + 1)
share_tr = np.bincount(MT_tr, minlength=MT_tr.max() + 1) / n_tr
share_te = np.bincount(MT_te, minlength=MT_tr.max() + 1) / len(MT_te)
W_tr = (share_te / np.maximum(share_tr, 1e-12))[MT_tr]
print("metro weights (test share / train share):", np.round(share_te / np.maximum(share_tr, 1e-12), 2))


def fit_cw(P, yy, w, mt):
    W = np.ones(K)
    for _ in range(3):
        for k in range(1, K):
            grid = W[k] * np.exp(np.linspace(-0.4, 0.4, 41))
            sc = [fast_f1(yy, (P * np.r_[W[:k], g_, W[k + 1:]]).argmax(1), w) for g_ in grid]
            W[k] = grid[int(np.argmax(sc))]
    return W


def apply_cw(P, W, mt):
    return (P * W).argmax(1)


def fit_cuts(s, yy, w, mt=None):
    cuts = np.quantile(s, np.cumsum(np.bincount(yy, minlength=K) / len(yy))[:-1])
    for _ in range(3):
        for i in range(K - 1):
            lo = cuts[i - 1] + 1e-6 if i else s.min()
            hi = cuts[i + 1] - 1e-6 if i < K - 2 else s.max()
            grid = np.linspace(max(lo, cuts[i] - 0.4), min(hi, cuts[i] + 0.4), 61)
            sc = [fast_f1(yy, np.digitize(s, np.r_[cuts[:i], g_, cuts[i + 1:]]), w) for g_ in grid]
            cuts[i] = grid[int(np.argmax(sc))]
    return cuts


def apply_cuts(s, cuts, mt):
    return np.digitize(s, cuts)


def fit_metro_cuts(s, yy, w, mt):
    g0 = fit_cuts(s, yy, w)
    C = np.tile(g0, (MT_tr.max() + 1, 1))
    pred = np.digitize(s, g0)
    cm = conf(yy, pred, w)
    for _ in range(2):
        for m_ in np.unique(mt):
            ix = mt == m_
            for i in range(K - 1):
                rest = cm - conf(yy[ix], pred[ix], w[ix])
                lo = C[m_, i - 1] + 1e-6 if i else s.min()
                hi = C[m_, i + 1] - 1e-6 if i < K - 2 else s.max()
                grid = np.linspace(max(lo, C[m_, i] - 0.5), min(hi, C[m_, i] + 0.5), 41)
                best_g, best_s = C[m_, i], -1
                for g_ in grid:
                    p_ = np.digitize(s[ix], np.r_[C[m_, :i], g_, C[m_, i + 1:]])
                    sc = f1_from_conf(rest + conf(yy[ix], p_, w[ix]))
                    if sc > best_s:
                        best_g, best_s = g_, sc
                C[m_, i] = best_g
                pred[ix] = np.digitize(s[ix], C[m_])
                cm = rest + conf(yy[ix], pred[ix], w[ix])
    return C


def apply_metro_cuts(s, C, mt):
    out = np.empty(len(s), int)
    for m_ in np.unique(mt):
        ix = mt == m_
        out[ix] = np.digitize(s[ix], C[m_])
    return out


def crossfit(o, fit, apply):
    pred = np.empty(n_tr, int)
    for f in range(FOLDS):
        m_ = fold_id == f
        prm = fit(o[~m_], y[~m_], W_tr[~m_], MT_tr[~m_])
        pred[m_] = apply(o[m_], prm, MT_tr[m_])
    return fast_f1(y, pred, W_tr), fast_f1(y, pred)


def opt_weights(models):
    Ps = [OOF[m] for m in models]
    def loss(z):
        w = np.exp(z) / np.exp(z).sum()
        return log_loss(y, np.clip(sum(wi * p for wi, p in zip(w, Ps)), 1e-7, 1))
    z = minimize(loss, np.zeros(len(Ps)), method="Nelder-Mead", options={"maxiter": 300}).x
    return np.exp(z) / np.exp(z).sum()


candidates = {}  # name -> (oof output, test output, kind)
for m_ in PROB_MODELS + ["stack"]:
    candidates[m_] = (OOF[m_], TEST[m_], "prob")
for m_ in REG_MODELS:
    candidates[m_] = (OOF[m_][:, 0], TEST[m_][:, 0], "score")
if len(PROB_MODELS) > 1:
    candidates["avg(" + "+".join(PROB_MODELS) + ")"] = (np.mean([OOF[m_] for m_ in PROB_MODELS], 0),
                                                       np.mean([TEST[m_] for m_ in PROB_MODELS], 0), "prob")
    wb = opt_weights(PROB_MODELS)
    print("optimised blend weights:", dict(zip(PROB_MODELS, wb.round(3))))
    candidates["opt_blend"] = (sum(wi * OOF[m_] for wi, m_ in zip(wb, PROB_MODELS)),
                               sum(wi * TEST[m_] for wi, m_ in zip(wb, PROB_MODELS)), "prob")
if REG_MODELS:
    reg_o = np.mean([OOF[m_][:, 0] for m_ in REG_MODELS], 0)
    reg_t = np.mean([TEST[m_][:, 0] for m_ in REG_MODELS], 0)
    if len(REG_MODELS) > 1:
        candidates["avg(" + "+".join(REG_MODELS) + ")"] = (reg_o, reg_t, "score")
    reg_opts = {m_: (OOF[m_][:, 0], TEST[m_][:, 0]) for m_ in REG_MODELS}
    if len(REG_MODELS) > 1:
        reg_opts["avg_reg"] = (reg_o, reg_t)
    for src in [c for c in ("opt_blend", "stack") if c in candidates]:
        for rname, (ro, rt) in reg_opts.items():
            for a in (0.3, 0.5, 0.7):
                candidates[f"{src} + {rname} (a={a})"] = (a * (candidates[src][0] @ TIERS) + (1 - a) * ro,
                                                          a * (candidates[src][1] @ TIERS) + (1 - a) * rt, "score")

rows = []
for name, (o, t, kind) in candidates.items():
    s_o, s_t = (o @ TIERS, t @ TIERS) if kind == "prob" else (o, t)
    if kind == "prob":
        rows.append(dict(scenario=name, decoder="A class weights", out=(o, t), fit=fit_cw, apply=apply_cw,
                         score=crossfit(o, fit_cw, apply_cw)))
    rows.append(dict(scenario=name, decoder="B cut-points", out=(s_o, s_t), fit=fit_cuts, apply=apply_cuts,
                     score=crossfit(s_o, fit_cuts, apply_cuts)))
    print(f"  scored: {name}", flush=True)
for r in sorted([r for r in rows if r["decoder"].startswith("B")], key=lambda r: -r["score"][0])[:4]:
    rows.append(dict(scenario=r["scenario"], decoder="C per-metro cut-points", out=r["out"], fit=fit_metro_cuts,
                     apply=apply_metro_cuts, score=crossfit(r["out"][0], fit_metro_cuts, apply_metro_cuts)))
    print(f"  scored per-metro: {r['scenario']}", flush=True)

table = pd.DataFrame([{"scenario": r["scenario"], "decoder": r["decoder"],
                       "CV macro F1 (test-weighted)": round(r["score"][0], 4), "CV macro F1 (plain)": round(r["score"][1], 4)}
                      for r in rows]).sort_values("CV macro F1 (test-weighted)", ascending=False).reset_index(drop=True)
prog.finish()
display(table)