scripts/cmp_nn.py

"""Etapa 2 de la comparación: lee /tmp/cmp.npz (arrays + preds LightGBM + baseline),
entrena la red neuronal multi-tarea (embeddings + softmax/KL para bloques y pasillos,
MSE para xT absoluto) e imprime la tabla final baseline / LightGBM / NN.

Corre en proceso APARTE de lightgbm (solo importa torch) para evitar el choque de
OpenMP. Uso: python scripts/cmp_nn.py
"""

from __future__ import annotations

import numpy as np
import torch
import torch.nn as nn

torch.manual_seed(7)
NPZ = "/tmp/cmp.npz"


def _kl(a, b):
    a = np.clip(a, 1e-9, None); b = np.clip(b, 1e-9, None)
    return float(np.mean(np.sum(a * (np.log(a) - np.log(b)), 1)))


def _metrics(y, p, dist=True):
    mae = float(np.mean(np.abs(y - p)))
    if dist:
        return {"MAE": round(mae, 4), "KL": round(_kl(y, p), 4)}
    return {"MAE": round(mae, 4), "RMSE": round(float(np.sqrt(np.mean((y - p) ** 2))), 4)}


class MTNet(nn.Module):
    def __init__(self, n_num, n_team, n_lg, n_form):
        super().__init__()
        self.te = nn.Embedding(n_team, 16); self.le = nn.Embedding(n_lg, 4); self.fe = nn.Embedding(n_form, 4)
        d = n_num + 16 + 16 + 4 + 4
        self.trunk = nn.Sequential(nn.Linear(d, 128), nn.ReLU(), nn.Dropout(0.3),
                                   nn.Linear(128, 64), nn.ReLU(), nn.Dropout(0.2))
        self.h_blo = nn.Linear(64, 3); self.h_pas = nn.Linear(64, 3); self.h_xt = nn.Linear(64, 3)

    def forward(self, x, ts, to, lg, fm):
        z = torch.cat([x, self.te(ts), self.te(to), self.le(lg), self.fe(fm)], 1)
        z = self.trunk(z)
        return self.h_blo(z), self.h_pas(z), self.h_xt(z)


def main():
    d = np.load(NPZ)
    Xn = d["Xn"]; Yb, Yp, Yx = d["Yb"], d["Yp"], d["Yx"]
    Bb, Bp, Bx = d["Bb"], d["Bp"], d["Bx"]
    tr, va, te = d["tr"], d["va"], d["te"]
    pb, pp, px = d["pb"], d["pp"], d["px"]
    ts, to, lg, fm = d["ts"], d["to"], d["lg"], d["fm"]
    nfeat = int(d["nfeat"])

    print("--- Red neuronal (embeddings + multi-tarea) ---", flush=True)
    T = lambda a: torch.tensor(a)
    net = MTNet(nfeat, int(d["n_team"]), int(d["n_lg"]), int(d["n_form"]))
    opt = torch.optim.Adam(net.parameters(), lr=2e-3, weight_decay=1e-4)
    klf = nn.KLDivLoss(reduction="batchmean"); msef = nn.MSELoss()
    Xt = T(Xn); tsT, toT, lgT, fmT = T(ts), T(to), T(lg), T(fm)
    Yb_t, Yp_t, Yx_t = T(Yb), T(Yp), T(Yx)
    tr_t, va_t, te_t = T(np.where(tr)[0]), T(np.where(va)[0]), T(np.where(te)[0])
    best = (1e9, None)
    for ep in range(300):
        net.train(); opt.zero_grad()
        b, p, x = net(Xt[tr_t], tsT[tr_t], toT[tr_t], lgT[tr_t], fmT[tr_t])
        loss = (klf(torch.log_softmax(b, 1), Yb_t[tr_t]) + klf(torch.log_softmax(p, 1), Yp_t[tr_t])
                + 0.5 * msef(x, Yx_t[tr_t]))
        loss.backward(); opt.step()
        if ep % 10 == 0:
            net.eval()
            with torch.no_grad():
                b, p, x = net(Xt[va_t], tsT[va_t], toT[va_t], lgT[va_t], fmT[va_t])
                v = (klf(torch.log_softmax(b, 1), Yb_t[va_t]) + klf(torch.log_softmax(p, 1), Yp_t[va_t])).item()
            if v < best[0]:
                best = (v, {k: val.clone() for k, val in net.state_dict().items()})
    net.load_state_dict(best[1]); net.eval()
    with torch.no_grad():
        b, p, x = net(Xt[te_t], tsT[te_t], toT[te_t], lgT[te_t], fmT[te_t])
        nb = torch.softmax(b, 1).numpy(); npp = torch.softmax(p, 1).numpy(); nx = np.clip(x.numpy(), 0, None)

    base = {"bloques": _metrics(Yb[te], Bb[te]), "pasillos": _metrics(Yp[te], Bp[te]), "xt": _metrics(Yx[te], Bx[te], dist=False)}
    lgb_res = {"bloques": _metrics(Yb[te], pb), "pasillos": _metrics(Yp[te], pp), "xt": _metrics(Yx[te], px, dist=False)}
    nn_res = {"bloques": _metrics(Yb[te], nb), "pasillos": _metrics(Yp[te], npp), "xt": _metrics(Yx[te], nx, dist=False)}

    print("\n================ COMPARACIÓN (test temporal) ================", flush=True)
    for k in ("bloques", "pasillos", "xt"):
        print(f"\n{k}:")
        print(f"  baseline : {base[k]}")
        print(f"  LightGBM : {lgb_res[k]}")
        print(f"  NN       : {nn_res[k]}")


if __name__ == "__main__":
    main()