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