| 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() |
|
|
