| import os |
|
|
| for _v in ( |
| "OMP_NUM_THREADS", |
| "OPENBLAS_NUM_THREADS", |
| "MKL_NUM_THREADS", |
| "NUMEXPR_NUM_THREADS", |
| ): |
| os.environ.setdefault(_v, "1") |
| import sys, time, json, argparse |
| import numpy as np |
| import scipy.sparse as sp |
| from scipy.sparse.linalg import spsolve |
|
|
| N = 96 |
| CH8 = [ |
| "mask", |
| "q_source", |
| "k_field", |
| "h_conv", |
| "T_amb", |
| "x_coord", |
| "y_coord", |
| "log_domain_L", |
| ] |
| SEED_TRAIN = 12340000 |
| SEED_HOT = 55000000 |
| SEED_HOLDOUT = 900000000 |
| SEED_EDGE = 900000000 + 500000000 |
| RANGES = dict( |
| domain_L=(0.02, 0.12), |
| cyl_rows=(2, 5), |
| cyl_cols=(2, 5), |
| cell_radius_frac=(0.3, 0.48), |
| aspect=(0.4, 2.5), |
| cell_fill=(0.55, 0.9), |
| current=(1.0, 60.0), |
| soc=(0.05, 1.0), |
| R0=(0.005, 0.05), |
| beta=(0.5, 4.0), |
| k_cell=(1.0, 30.0), |
| k_coolant=(0.1, 1.5), |
| h_conv=(5.0, 200.0), |
| T_amb=(15.0, 40.0), |
| ) |
|
|
|
|
| def build_system(N, hg, k_field, q_field, h_conv, T_amb): |
| n = N * N |
| dx2 = hg * hg |
| k = k_field.astype(np.float64) |
| q = q_field.astype(np.float64) |
|
|
| def hmean(a, b): |
| return 2.0 * a * b / (a + b + 1e-30) |
|
|
| ids = np.arange(n).reshape(N, N) |
| rows = [] |
| cols = [] |
| vals = [] |
| diag = np.zeros((N, N)) |
| b = q * dx2 |
| kf = hmean(k[1:, :], k[:-1, :]) |
| p = ids[1:, :].ravel() |
| nb = ids[:-1, :].ravel() |
| kk = kf.ravel() |
| rows.append(p) |
| cols.append(nb) |
| vals.append(-kk) |
| diag[1:, :] += kf |
| p = ids[:-1, :].ravel() |
| nb = ids[1:, :].ravel() |
| rows.append(p) |
| cols.append(nb) |
| vals.append(-kk) |
| diag[:-1, :] += kf |
| kf = hmean(k[:, 1:], k[:, :-1]) |
| p = ids[:, 1:].ravel() |
| nb = ids[:, :-1].ravel() |
| kk = kf.ravel() |
| rows.append(p) |
| cols.append(nb) |
| vals.append(-kk) |
| diag[:, 1:] += kf |
| p = ids[:, :-1].ravel() |
| nb = ids[:, 1:].ravel() |
| rows.append(p) |
| cols.append(nb) |
| vals.append(-kk) |
| diag[:, :-1] += kf |
| hdx = h_conv * hg |
| bmask = np.zeros((N, N)) |
| bmask[0, :] += 1 |
| bmask[-1, :] += 1 |
| bmask[:, 0] += 1 |
| bmask[:, -1] += 1 |
| diag += hdx * bmask |
| b += hdx * T_amb * bmask |
| rows.append(ids.ravel()) |
| cols.append(ids.ravel()) |
| vals.append(diag.ravel()) |
| rows = np.concatenate(rows) |
| cols = np.concatenate(cols) |
| vals = np.concatenate(vals) |
| A = sp.csr_matrix((vals, (rows, cols)), shape=(n, n)) |
| return (A, b.ravel()) |
|
|
|
|
| def solve_steady(N, hg, k_field, q_field, h_conv, T_amb): |
| A, b = build_system(N, hg, k_field, q_field, h_conv, T_amb) |
| return spsolve(A.tocsc(), b).reshape(N, N) |
|
|
|
|
| def sample_params(rng): |
| p = {} |
| p["domain_L"] = rng.uniform(*RANGES["domain_L"]) |
| p["geom"] = "cyl" if rng.random() < 0.6 else "pris" |
| p["current"] = rng.uniform(*RANGES["current"]) |
| p["soc"] = rng.uniform(*RANGES["soc"]) |
| p["R0"] = rng.uniform(*RANGES["R0"]) |
| p["beta"] = rng.uniform(*RANGES["beta"]) |
| p["k_cell"] = rng.uniform(*RANGES["k_cell"]) |
| p["k_coolant"] = rng.uniform(*RANGES["k_coolant"]) |
| p["h_conv"] = rng.uniform(*RANGES["h_conv"]) |
| p["T_amb"] = rng.uniform(*RANGES["T_amb"]) |
| if p["geom"] == "cyl": |
| p["rows"] = int(rng.integers(RANGES["cyl_rows"][0], RANGES["cyl_rows"][1] + 1)) |
| p["cols"] = int(rng.integers(RANGES["cyl_cols"][0], RANGES["cyl_cols"][1] + 1)) |
| p["rfrac"] = rng.uniform(*RANGES["cell_radius_frac"]) |
| else: |
| p["aspect"] = rng.uniform(*RANGES["aspect"]) |
| p["fill"] = rng.uniform(*RANGES["cell_fill"]) |
| return p |
|
|
|
|
| def build_fields(N, p): |
| L = p["domain_L"] |
| hg = L / (N - 1) |
| yy, xx = np.meshgrid(np.linspace(0, 1, N), np.linspace(0, 1, N), indexing="ij") |
| mask = np.zeros((N, N), dtype=np.float64) |
| R_int = p["R0"] * (1.0 + p["beta"] * (1.0 - p["soc"]) ** 2) |
| P_total = p["current"] ** 2 * R_int |
| if p["geom"] == "cyl": |
| rows, cols = (p["rows"], p["cols"]) |
| pitch_x = 1.0 / cols |
| pitch_y = 1.0 / rows |
| r = p["rfrac"] * min(pitch_x, pitch_y) |
| for i in range(rows): |
| cy = (i + 0.5) * pitch_y |
| for j in range(cols): |
| cx = (j + 0.5) * pitch_x |
| d = (xx - cx) ** 2 + (yy - cy) ** 2 |
| mask[d <= r * r] = 1.0 |
| else: |
| a = p["aspect"] |
| f = p["fill"] |
| hh = min(0.98, np.sqrt(f / a)) |
| ww = min(0.98, a * hh) |
| x0, x1 = (0.5 - ww / 2, 0.5 + ww / 2) |
| y0, y1 = (0.5 - hh / 2, 0.5 + hh / 2) |
| mask[(xx >= x0) & (xx <= x1) & (yy >= y0) & (yy <= y1)] = 1.0 |
| cell_area_phys = mask.sum() * hg * hg |
| if cell_area_phys <= 0: |
| cell_area_phys = hg * hg |
| q_field = mask * (P_total / cell_area_phys) |
| k_field = np.where(mask > 0, p["k_cell"], p["k_coolant"]) |
| return (mask, k_field, q_field, hg) |
|
|
|
|
| def build_X(p, mask, k_field, q_field): |
| yy, xx = np.meshgrid(np.linspace(0, 1, N), np.linspace(0, 1, N), indexing="ij") |
| logL = np.full((N, N), np.log(p["domain_L"]), dtype=np.float64) |
| return np.stack( |
| [ |
| mask, |
| q_field, |
| k_field, |
| np.full((N, N), p["h_conv"]), |
| np.full((N, N), p["T_amb"]), |
| xx, |
| yy, |
| logL, |
| ], |
| 0, |
| ).astype(np.float32) |
|
|
|
|
| def gen_uniform(seed): |
| rng = np.random.default_rng(seed) |
| p = sample_params(rng) |
| mask, k, q, hg = build_fields(N, p) |
| T = solve_steady(N, hg, k, q, p["h_conv"], p["T_amb"]) |
| return (build_X(p, mask, k, q), T.astype(np.float32)) |
|
|
|
|
| def run(fn, jobs, workers): |
| if workers <= 1: |
| out = [fn(j) for j in jobs] |
| else: |
| from multiprocessing import Pool |
|
|
| with Pool(workers) as pool: |
| out = list(pool.imap(fn, jobs, chunksize=8)) |
| Xs, Ys = zip(*out) |
| return (np.stack(Xs), np.stack(Ys)) |
|
|
|
|
| def make_split(n_train, workers): |
| Xtr, Ytr = run(gen_uniform, [SEED_TRAIN + i for i in range(n_train)], workers) |
| n = Xtr.shape[0] |
| perm = np.random.default_rng(0).permutation(n) |
| ntr = int(0.8 * n) |
| nva = int(0.1 * n) |
| sp_idx = { |
| "train": perm[:ntr], |
| "val": perm[ntr : ntr + nva], |
| "test": perm[ntr + nva :], |
| } |
| return (Xtr, Ytr, sp_idx) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument( |
| "--data", default="/workspace/pde_surrogate/battery/battery_thermal_v3.h5" |
| ) |
| ap.add_argument( |
| "--hot", default="/workspace/pde_surrogate/battery/hot_enrich_v3.h5" |
| ) |
| ap.add_argument( |
| "--out", default="/workspace/pde_surrogate/battery/fno_ckpt_v3_final.pt" |
| ) |
| ap.add_argument("--n_train", type=int, default=4000) |
| ap.add_argument("--n_hot", type=int, default=5000) |
| ap.add_argument("--workers", type=int, default=32) |
| ap.add_argument("--epochs", type=int, default=300) |
| ap.add_argument("--bs", type=int, default=128) |
| ap.add_argument("--lr", type=float, default=0.002) |
| ap.add_argument("--modes", type=int, default=32) |
| ap.add_argument("--width", type=int, default=80) |
| ap.add_argument("--layers", type=int, default=4) |
| ap.add_argument("--scale_floor", type=float, default=2.0) |
| ap.add_argument("--grad_w", type=float, default=3.0) |
| ap.add_argument("--rise_w", type=float, default=1.5) |
| ap.add_argument("--scale_w", type=float, default=3.0) |
| ap.add_argument("--smoke", action="store_true") |
| args = ap.parse_args() |
| if args.smoke: |
| Xtr, Ytr, sp_idx = make_split(args.n_train, args.workers) |
| n = Xtr.shape[0] |
| assert Xtr.shape[1] == 8 and Xtr.shape[2] == N and (Xtr.shape[3] == N) |
| assert ( |
| len(sp_idx["train"]) == 3200 |
| and len(sp_idx["val"]) == 400 |
| and (len(sp_idx["test"]) == 400) |
| ) |
| xu, yu = gen_uniform(SEED_HOLDOUT) |
| assert xu.shape == (8, N, N) and np.isfinite(yu).all() and (yu.max() > yu.min()) |
| print( |
| f"[Smoke] n={n} train={len(sp_idx['train'])} val={len(sp_idx['val'])} test={len(sp_idx['test'])} ch={len(CH8)} holdoutT[{yu.min():.2f},{yu.max():.2f}]", |
| flush=True, |
| ) |
| return |
| import h5py, torch, torch.nn as nn, torch.nn.functional as F |
|
|
| sys.path.insert(0, "/workspace/pde_surrogate/scripts") |
| from train_fno_v32 import FNO2dV32, phys_prior, spatial_grad_mag |
|
|
| AMBCH = 4 |
|
|
| def load_h5(path, split): |
| with h5py.File(path, "r") as f: |
| return ( |
| torch.from_numpy(f[split]["X"][:].astype(np.float32)), |
| torch.from_numpy(f[split]["Y"][:].astype(np.float32)), |
| ) |
|
|
| dev = "cuda" |
| with h5py.File(args.data, "r") as f: |
| x_mean = np.array(f.attrs["x_mean"], np.float32) |
| x_std = np.array(f.attrs["x_std"], np.float32) |
| y_mean = float(f.attrs["y_mean"]) |
| y_std = float(f.attrs["y_std"]) |
| channels = json.loads(f.attrs["channels"]) |
| base_in = f["train"]["X"].shape[1] |
| Xtr, Ytr = load_h5(args.data, "train") |
| Xva, Yva = load_h5(args.data, "val") |
| with h5py.File(args.hot, "r") as f: |
| Xh = torch.from_numpy(f["hot"]["X"][:].astype(np.float32)) |
| Yh = torch.from_numpy(f["hot"]["Y"][:].astype(np.float32)) |
| Xtr = torch.cat([Xtr, Xh], 0) |
| Ytr = torch.cat([Ytr, Yh], 0) |
| xm_ = Xtr.reshape(Xtr.shape[0], base_in, -1).mean((0, 2)) |
| xs_ = Xtr.reshape(Xtr.shape[0], base_in, -1).std((0, 2)) + 1e-06 |
| x_mean = xm_.numpy() |
| x_std = xs_.numpy() |
| Xtr = Xtr.to(dev) |
| Ytr = Ytr.to(dev) |
| Xva = Xva.to(dev) |
| Yva = Yva.to(dev) |
| ntr = Xtr.shape[0] |
| xm = torch.tensor(x_mean, device=dev)[None, :, None, None] |
| xs = torch.tensor(x_std, device=dev)[None, :, None, None] |
| model = FNO2dV32( |
| in_ch=base_in, modes=args.modes, width=args.width, L=args.layers |
| ).to(dev) |
| opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-05) |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, args.epochs) |
| print( |
| f"[V32] params={sum((p.numel() for p in model.parameters())) / 1000000.0:.2f}M ntr={ntr}", |
| flush=True, |
| ) |
|
|
| def prep(Xb): |
| dTc, dTd = phys_prior(Xb) |
| return ((Xb - xm) / xs, dTc, dTd) |
|
|
| def val_binned(): |
| model.eval() |
| rls = [] |
| rises = [] |
| mae = 0.0 |
| vb = 0 |
| with torch.no_grad(): |
| for s in range(0, Xva.shape[0], args.bs): |
| Xb = Xva[s : s + args.bs] |
| Yb = Yva[s : s + args.bs][:, None] |
| Xin, dTc, dTd = prep(Xb) |
| field_ps, log_scale, _ = model(Xin, dTc, dTd) |
| amb = Xb[:, AMBCH].flatten(1).mean(1) |
| log_scale = log_scale.clamp( |
| float(np.log(args.scale_floor)), float(np.log(2000.0)) |
| ) |
| pp = ( |
| amb[:, None, None, None] |
| + torch.exp(log_scale)[:, None, None, None] * field_ps |
| ) |
| num = torch.linalg.norm((pp - Yb).flatten(1), dim=1) |
| den = torch.linalg.norm(Yb.flatten(1), dim=1).clamp_min(1e-08) |
| rls.append((num / den).cpu()) |
| rises.append((Yb.flatten(1).amax(1) - amb).cpu()) |
| mae += torch.abs(pp - Yb).mean().item() |
| vb += 1 |
| rls = torch.cat(rls) |
| rises = torch.cat(rises) |
| mae /= vb |
| hot = rls[rises > 10] |
| return ( |
| float(rls.median()), |
| float(hot.median()) if hot.numel() else float("nan"), |
| mae, |
| int((rises > 10).sum()), |
| ) |
|
|
| best = 1000000000.0 |
| for ep in range(args.epochs): |
| model.train() |
| perm = torch.randperm(ntr, device=dev) |
| tl = 0.0 |
| nb = 0 |
| for s in range(0, ntr, args.bs): |
| idx = perm[s : s + args.bs] |
| Xb = Xtr[idx] |
| Yb = Ytr[idx][:, None] |
| amb = Xb[:, AMBCH].flatten(1).mean(1) |
| flat = Yb.flatten(1) |
| scale = torch.quantile(flat - amb[:, None], 0.95, dim=1).clamp_min( |
| args.scale_floor |
| ) |
| Yps = (Yb - amb.view(-1, 1, 1, 1)) / scale.view(-1, 1, 1, 1) |
| Xin, dTc, dTd = prep(Xb) |
| opt.zero_grad() |
| field_ps, log_scale, prior = model(Xin, dTc, dTd) |
| gm = spatial_grad_mag(Yps) |
| gw = 1.0 + args.grad_w * gm / ( |
| gm.flatten(1).mean(1)[:, None, None, None] + 1e-06 |
| ) |
| rise = (flat.amax(1) - amb).clamp_min(0.0) |
| sw = ((rise + 1.0) ** args.rise_w)[:, None, None, None] |
| sw = sw / sw.mean() |
| field_loss = (gw * sw * (field_ps - Yps) ** 2).mean() |
| scale_loss = F.smooth_l1_loss(log_scale, torch.log(scale), beta=0.1) |
| loss = field_loss + args.scale_w * scale_loss |
| loss.backward() |
| opt.step() |
| tl += loss.item() |
| nb += 1 |
| sched.step() |
| overall, hotm, mae, nhot = val_binned() |
| score = hotm if hotm == hotm else overall |
| if score < best: |
| best = score |
| torch.save( |
| { |
| "model": model.state_dict(), |
| "args": vars(args), |
| "x_mean": x_mean, |
| "x_std": x_std, |
| "y_mean": y_mean, |
| "y_std": y_std, |
| "channels": channels, |
| "scale_floor": args.scale_floor, |
| "amb_channel": AMBCH, |
| "in_ch": base_in, |
| }, |
| args.out, |
| ) |
| if ep % 10 == 0 or ep == args.epochs - 1: |
| print( |
| f"Ep{ep:3d} loss={tl / nb:.4e} overall={overall:.4f} HOTSPOT(>10C,n={nhot})={hotm:.4f} MAE={mae:.3f}C (best {best:.4f})", |
| flush=True, |
| ) |
| print(f"[V32] DONE best HOTSPOT {best:.4f} -> {args.out}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
