import sys, time, os, numpy as np, pandas as pd, torch, torch.nn as nn, torch.nn.functional as F dev = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(2024) np.random.seed(2024) if dev == "cpu" and os.environ.get("OMP_NUM_THREADS"): torch.set_num_threads(int(os.environ["OMP_NUM_THREADS"])) valflags = {"--save", "--arch", "--lr"} pos = [] av = sys.argv[1:] i = 0 while i < len(av): if av[i] in valflags: i += 2 continue if av[i].startswith("--"): i += 1 continue pos.append(av[i]) i += 1 track = pos[0] dL, dm, dD, dn = {"solar": (6, 3, 64, 4), "wind": (12, 6, 96, 3)}[track] L = int(pos[1]) if len(pos) > 1 else dL modes = int(pos[2]) if len(pos) > 2 else dm D = int(pos[3]) if len(pos) > 3 else dD nblk = int(pos[4]) if len(pos) > 4 else dn ep = int(pos[5]) if len(pos) > 5 else 80 save = sys.argv[sys.argv.index("--save") + 1] if "--save" in sys.argv else None arch = sys.argv[sys.argv.index("--arch") + 1] if "--arch" in sys.argv else "dual" lr = float(sys.argv[sys.argv.index("--lr") + 1]) if "--lr" in sys.argv else 0.002 smoke = "--smoke" in sys.argv prep = "/workspace/gefcom/prep" qs = np.arange(1, 100) / 100.0 nwp = { "solar": [ "VAR78", "VAR79", "VAR134", "VAR157", "VAR164", "VAR165", "VAR166", "VAR167", "VAR169", "VAR175", "VAR178", "VAR228", ], "wind": ["U10", "V10", "U100", "V100"], }[track] def build(): df = ( pd.read_parquet(f"{prep}/{track}.parquet") .sort_values(["ZONEID", "TIMESTAMP"]) .reset_index(drop=True) ) df["hour"] = df.TIMESTAMP.dt.hour df["doy"] = df.TIMESTAMP.dt.dayofyear df["hsin"] = np.sin(2 * np.pi * df.hour / 24) df["hcos"] = np.cos(2 * np.pi * df.hour / 24) df["dsin"] = np.sin(2 * np.pi * df.doy / 365.25) df["dcos"] = np.cos(2 * np.pi * df.doy / 365.25) feats = list(nwp) + ["hsin", "hcos", "dsin", "dcos"] if track == "wind": df["ws10"] = np.hypot(df.U10, df.V10) df["ws100"] = np.hypot(df.U100, df.V100) df["wd100"] = np.arctan2(df.V100, df.U100) df["wds"] = np.sin(df.wd100) df["wdc"] = np.cos(df.wd100) df["ws100_2"] = df.ws100**2 df["ws100_3"] = df.ws100**3 df["shear"] = df.ws100 - df.ws10 feats += ["ws10", "ws100", "wds", "wdc", "ws100_2", "ws100_3", "shear"] else: df["csi"] = df.VAR169 / (df.VAR178 + 1000.0) df["cloud2"] = df.VAR164**2 df["temp_c"] = df.VAR167 - 273.15 df["daylight"] = (df.VAR178 > 10000.0).astype(np.float32) feats += ["csi", "cloud2", "temp_c", "daylight"] tr = df.is_test == 0 mu = df.loc[tr, feats].mean() sd = df.loc[tr, feats].std() + 1e-06 df[feats] = (df[feats] - mu) / sd df["y"] = df["y"].fillna(0.0) X, Y, M = ([], [], []) half = L // 2 for z, g in df.groupby("ZONEID"): g = g.reset_index(drop=True) fv = g[feats].values.astype(np.float32) powr = g["y"].values.astype(np.float32) iste = g["is_test"].values T = len(g) for i in range(T): a = i - half b = i - half + L if a < 0 or b > T: continue X.append(fv[a:b]) Y.append(powr[i]) M.append(iste[i]) X = torch.tensor(np.array(X, dtype=np.float32)) Y = torch.tensor(np.array(Y, dtype=np.float32)) return (X, Y, np.array(M), df, feats) class RevIN(nn.Module): def __init__(self, C): super().__init__() self.g = nn.Parameter(torch.ones(C)) self.b = nn.Parameter(torch.zeros(C)) def norm(self, x): self.m = x.mean(1, keepdim=True) self.s = x.std(1, keepdim=True) + 1e-05 return (x - self.m) / self.s * self.g + self.b class FNO1D(nn.Module): def __init__(self, D, modes): super().__init__() self.modes = modes s = 1 / (D * D) self.w = nn.Parameter(s * torch.rand(modes, D, D, dtype=torch.cfloat)) def forward(self, x): P = x.shape[1] xf = torch.fft.rfft(x, dim=1) m = min(self.modes, xf.shape[1]) o = torch.zeros_like(xf) o[:, :m] = torch.einsum("bpd,pde->bpe", xf[:, :m], self.w[:m]) return torch.fft.irfft(o, n=P, dim=1) class Block(nn.Module): def __init__(self, D, modes, ff=2, drop=0.1): super().__init__() self.n1 = nn.LayerNorm(D) self.fno = FNO1D(D, modes) self.d1 = nn.Dropout(drop) self.n2 = nn.LayerNorm(D) self.ff = nn.Sequential( nn.Linear(D, D * ff), nn.GELU(), nn.Dropout(drop), nn.Linear(D * ff, D) ) def forward(self, x): x = x + self.d1(self.fno(self.n1(x))) return x + self.ff(self.n2(x)) class FELA_Grid(nn.Module): def __init__(self, Fin, L, D=96, modes=6, nblk=3, nq=99, arch="dual"): super().__init__() self.L = L self.arch = arch self.center = L // 2 self.nq = nq self.revin = RevIN(Fin) self.embed = nn.Linear(Fin, D) self.pos = nn.Parameter(0.02 * torch.randn(1, L, D)) self.blocks = nn.ModuleList([Block(D, modes) for _ in range(nblk)]) self.norm = nn.LayerNorm(D) if arch == "dual": self.direct = nn.Sequential( nn.Linear(Fin, D), nn.GELU(), nn.Linear(D, D), nn.GELU() ) fuse_in = 2 * D else: fuse_in = D self.med = nn.Linear(fuse_in, 1) self.spread = nn.Linear(fuse_in, nq) self.register_buffer("qidx", torch.arange(nq)) def forward(self, x): xc = x[:, self.center] xn = self.revin.norm(x) h = self.embed(xn) + self.pos for b in self.blocks: h = b(h) ctx = self.norm(h)[:, self.center] z = torch.cat([ctx, self.direct(xc)], dim=1) if self.arch == "dual" else ctx med = torch.sigmoid(self.med(z)) w = F.softplus(self.spread(z)) half = self.nq // 2 below = -torch.flip(torch.cumsum(torch.flip(w[:, :half], [1]), 1), [1]) above = torch.cumsum(w[:, half + 1 :], 1) offs = ( torch.cat([below, torch.zeros_like(w[:, half : half + 1]), above], dim=1) * 0.05 ) return torch.clamp(med + offs, 0, 1) def pinball(pred, y, qs): y = y[:, None] e = y - pred return torch.maximum(qs[None, :] * e, (qs[None, :] - 1) * e).mean() def main(): X, Y, M, df, feats = build() Fin = X.shape[2] tr_idx = np.where(M == 0)[0] te_idx = np.where(M == 1)[0] yte = Y[te_idx].numpy() keep = ~np.isnan(yte) te_idx = te_idx[keep] nval = int(len(tr_idx) * 0.06) va_idx = tr_idx[-nval:] tr_idx = tr_idx[:-nval] assert len(te_idx) == {"solar": 2154, "wind": 7390}[track] if smoke: print( f"{track} Test_windows {len(te_idx)} train {len(tr_idx)} val {len(va_idx)} zones {df.ZONEID.nunique()}" ) return qst = torch.tensor(qs, dtype=torch.float32, device=dev) Xtr, Ytr = (X[tr_idx].to(dev), Y[tr_idx].to(dev)) Xva, Yva = (X[va_idx].to(dev), Y[va_idx].to(dev)) Xte, Yte = (X[te_idx].to(dev), Y[te_idx].to(dev)) m = FELA_Grid(Fin, L, D=D, modes=modes, nblk=nblk, arch=arch).to(dev) npar = sum((p.numel() for p in m.parameters())) print( f"[{track}] arch={arch} L={L} modes={modes} D={D} nblk={nblk} Fin={Fin} | train {len(Xtr)} val {len(Xva)} test {len(Xte)} | {npar / 1000.0:.0f}K" ) opt = torch.optim.Adam(m.parameters(), lr=lr, weight_decay=1e-05) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, ep) bs = 512 def ev(Xs, Ys): m.eval() with torch.no_grad(): ps = [m(Xs[i : i + 8192]) for i in range(0, len(Xs), 8192)] p = torch.cat(ps) return (pinball(p, Ys, qst).item(), p) best = 1000000000.0 bstate = None bad = 0 for e in range(ep): m.train() perm = torch.randperm(len(Xtr), device=dev) for i in range(0, len(Xtr) - bs, bs): idx = perm[i : i + bs] loss = pinball(m(Xtr[idx]), Ytr[idx], qst) opt.zero_grad() loss.backward() opt.step() sched.step() vpb, _ = ev(Xva, Yva) if vpb < best: best = vpb bstate = {k: v.detach().clone() for k, v in m.state_dict().items()} bad = 0 else: bad += 1 if bad >= 15: break m.load_state_dict(bstate) tpb, _ = ev(Xte, Yte) bench = {"solar": 0.0285, "wind": 0.0792}[track] print( f"RESULT {track} pinball {tpb:.5f} (off.bench {bench}) | {npar / 1000.0:.0f}K" ) if save: torch.save( { "state": m.state_dict(), "cfg": dict(Fin=Fin, L=L, D=D, modes=modes, nblk=nblk), "feats": feats, "track": track, "test_pinball": tpb, "off_bench": bench, "npar": npar, }, save, ) print(f"SAVED {save}") if __name__ == "__main__": main()