| """Reference outputs of the three release models, computed with the ORIGINAL code. |
| |
| Runs the frozen code in ../source_snapshot on the frozen weights in ../originals, on the |
| first devices of the internal test split 0, and writes ../reference/<model>.npz. The new |
| boss_bench code must reproduce these outputs on the same devices (matched by uuid in the |
| published files). |
| |
| Usage: python make_reference.py {loss_fraction,heat_flux_iid,inverse_design} |
| |
| Everything runs on CPU in float32 with fixed noise, so the outputs are deterministic. |
| """ |
|
|
| import os |
| import sys |
| import types |
|
|
| import h5py |
| import numpy as np |
| import torch |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| SNAP = os.path.join(ROOT, "source_snapshot") |
| ORIG = os.path.join(ROOT, "originals") |
| OUT = os.path.join(ROOT, "reference") |
| DATA = "/path/to/internal/processed/confinement/confinement_data_test_split_00000.hdf5" |
| SEED = 0 |
|
|
|
|
| def uuids(rows): |
| with h5py.File(DATA, "r") as f: |
| return f["uuid"][:][rows].astype(str) |
|
|
|
|
| def loss_fraction(n=64): |
| sys.path.insert(0, os.path.join(SNAP, "artemiss_benchmarks")) |
| from benchmarks.surrogates.models_vmec.train import Predictor, load_files |
|
|
| d = os.path.join(ORIG, "loss_fraction") |
| pred = Predictor(os.path.join(d, "model_loss_frac_conv_tmax_100ms_biggg_sigmoid.pt"), |
| os.path.join(d, "model_config_loss_frac_conv_tmax_100ms_biggg_sigmoid.json")) |
| |
| data = load_files(DATA, "loss_frac", s0=0.1, t_max_override=0.1) |
|
|
| |
| with h5py.File(DATA, "r") as f: |
| valid = ~np.isnan(f["surf_xyz"][:]).any(axis=1) |
| for k in ("pressure", "current"): |
| valid &= ~np.isnan(f[k][:]).any(axis=1) |
| valid &= ~np.isnan(f["volavgB"][:]) |
| rows = np.flatnonzero(valid) |
| assert len(rows) == len(data["target"]) |
| rows = rows[:n] |
|
|
| inputs = {k: data[k][:n] for k in ("surf_r", "surf_z", "pressure", "current", "volavgB", "nfp")} |
| y = pred.predict_batch(**inputs) |
| np.savez(os.path.join(OUT, "loss_fraction.npz"), uuid=uuids(rows), rows=rows, |
| prediction=np.asarray(y), target=data["target"][:n], **inputs) |
| mae = np.abs(np.asarray(y) - data["target"][:n]).mean() |
| print(f"loss_fraction: {n} devices, MAE vs traced = {mae:.4f}") |
|
|
|
|
| def heat_flux_iid(n_configs=16, n_points=512, n_steps=50): |
| sys.path.insert(0, os.path.join(SNAP, "fusion", "heat_flux")) |
| import train_final as tf |
| from flow_model_final import FlowMatchingModel |
|
|
| d = os.path.join(ORIG, "heat_flux_iid") |
| cfg = torch.load(os.path.join(d, "model_config.pt"), weights_only=True) |
| model = FlowMatchingModel( |
| d_cond=cfg["d_cond"], x_dim=cfg.get("x_dim", tf.X_DIM), |
| embed_dim=cfg["embed_dim"], hidden=cfg["hidden_dim"], |
| full_cloud=cfg.get("full_cloud", False), pool=cfg.get("pool", "max"), |
| gate_global=cfg.get("gate_global", False), |
| n_vel_blocks=cfg.get("n_vel_blocks", 4)) |
| model.load_state_dict(torch.load(os.path.join(d, "model_best.pt"), weights_only=True)) |
| model.eval() |
| stats = dict(np.load(os.path.join(d, "cond_stats.npz"))) |
|
|
| |
| scal, surf, v2u, sinit, points, lf, ni, _ = tf.load_single_file(DATA, min_lost=0) |
| keep = v2u < n_configs |
| scal, surf = scal[:n_configs], surf[:n_configs] |
| v2u, sinit = v2u[keep], sinit[keep] |
| points = [p for p, k in zip(points, keep) if k] |
| lf, ni = lf[keep], ni[keep] |
| raw = dict(cond_scalars=scal.copy(), surf_xyz=surf.copy(), s_init=sinit.copy()) |
|
|
| tf.apply_stats(scal, sinit, surf, points, stats) |
| c_emb, f_pred = tf.encode_and_predict_all( |
| model, scal, sinit, torch.from_numpy(surf), v2u, torch.device("cpu")) |
|
|
| gen = torch.Generator().manual_seed(SEED) |
| x0 = torch.randn(len(v2u), n_points, tf.X_DIM, generator=gen) |
| with torch.no_grad(): |
| c_pt = c_emb.repeat_interleave(n_points, dim=0) |
| x1 = tf._ode_integrate(model, x0.reshape(-1, tf.X_DIM), c_pt, n_steps, "euler", amp=False) |
| x1 = x1.reshape(len(v2u), n_points, tf.X_DIM) |
|
|
| np.savez(os.path.join(OUT, "heat_flux_iid.npz"), uuid=uuids(np.arange(n_configs)), |
| config_index=v2u, loss_fraction_pred=f_pred, loss_fraction=lf, n_inner=ni, |
| cond_embedding=c_emb.numpy(), x0=x0.numpy(), x1=x1.numpy(), |
| n_steps=n_steps, sampler="euler", **raw) |
| print(f"heat_flux_iid: {len(v2u)} samples, loss-fraction MAE = {np.abs(f_pred - lf).mean():.4f}") |
|
|
|
|
| def inverse_design(n=8): |
| here = os.path.join(SNAP, "fusion", "shape_generation", "inverse_design") |
| sys.path.insert(0, here) |
| sys.modules.setdefault("wandb", types.ModuleType("wandb")) |
| import dataloading as dl |
| from evaluate_samples import build_model, to_rz |
| from flow_settings import load_flow_settings |
| from sample_loss_sweep import sample_from_noise |
|
|
| ck = torch.load(os.path.join(ORIG, "inverse_design", "boss_diffusion_surf_rz_unet_ep4000_ema.pt"), |
| map_location="cpu", weights_only=False) |
| model, d_design = build_model(ck, "cpu") |
| model.load_state_dict(ck["ema"]) |
| model.eval() |
| sp = load_flow_settings()["sampling"] |
|
|
| ds = dl.StellaratorDataset([DATA], stats=ck["stats"], normalize=True) |
| |
| with h5py.File(DATA, "r") as f: |
| xyz = f["surf_xyz"][:].astype(np.float32) |
| rows = np.array([np.flatnonzero((xyz == ds.surf_xyz[i]).all(axis=1))[0] for i in range(n)]) |
|
|
| props = ds.props[:n] |
| x0 = torch.randn(n, d_design, generator=torch.Generator().manual_seed(SEED)) |
| with torch.no_grad(): |
| x1 = sample_from_noise(model, props, x0, sp["n_steps"], sp["method"], sp["t_min"], sp["t_max"]) |
| R, Z = to_rz(x1, ck["stats"], "cpu") |
|
|
| np.savez(os.path.join(OUT, "inverse_design.npz"), uuid=uuids(rows), rows=rows, |
| props_raw=ds.props_raw[:n], props=props.numpy(), x0=x0.numpy(), x1=x1.numpy(), |
| R=R, Z=Z, R_true=ds.design_raw[:n].reshape(n, 41, 41, 2)[..., 0], |
| Z_true=ds.design_raw[:n].reshape(n, 41, 41, 2)[..., 1], **{f"sampling_{k}": v for k, v in sp.items()}) |
| print(f"inverse_design: {n} designs, sampler {sp}") |
|
|
|
|
| if __name__ == "__main__": |
| torch.manual_seed(SEED) |
| {"loss_fraction": loss_fraction, "heat_flux_iid": heat_flux_iid, |
| "inverse_design": inverse_design}[sys.argv[1]]() |
|
|
