| """A direct predictor with its correcting network as an ONNX file under the benchmark's contract: |
| python export_fb.py NAME OUT.onnx [--check N] |
| |
| Reads ml/runs/NAME.pt (train_fb.py) and its base network, and writes a graph from boundary_r_cos, boundary_z_sin |
| (5, 9), requirement_metrics (4), n_field_periods and n_coils_per_half_period to coil_x_n, coil_y_n, coil_z_n (K, 15) |
| and currents (K) in amperes, with the metadata key supported_tracks set to fixed_shape. The graph runs the base |
| network, computes the feedback of train_fb.py for its coilset with tensor operations (feedback() below, the same |
| quantities as fast.feedback), and applies the corrector, once per round. Prints the largest difference between the |
| graph's outputs in onnxruntime and the networks' in PyTorch, and with --check N the largest difference between |
| feedback() and fast.feedback on N dev instances. |
| """ |
| from __future__ import annotations |
|
|
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| import onnxruntime as ort |
| import torch |
| from torch import nn |
|
|
| from physics import Physics, coil_basis |
| from train import Net |
| from train_fb import Corrector |
|
|
| HERE = Path(__file__).resolve().parent |
|
|
|
|
| def cross(a, b): |
| return torch.stack([a[..., 1] * b[..., 2] - a[..., 2] * b[..., 1], a[..., 2] * b[..., 0] - a[..., 0] * b[..., 2], |
| a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0]], -1) |
|
|
|
|
| class Feedback(nn.Module): |
| """fast.feedback for one coilset, with the pair sums as tensors: constants are the grids' angles, the coil |
| modes at 64 and 128 parameters and the symmetry maps.""" |
|
|
| def __init__(self, nfp, k): |
| super().__init__() |
| ph = Physics(nfp, k, "cpu") |
| self.k, self.J = k, ph.J |
| buf = lambda name, v: self.register_buffer(name, torch.as_tensor(np.asarray(v), dtype=torch.float32)) |
| for name, (th, p) in (("f", ph.grid(32, int(round(3.2 * k)), 0.5, 0.5)), ("s", ph.grid(32, 16 * ph.J, 0.5, 0.5, full=True)), |
| ("a", (ph.grid(64, 48)[0], ph.grid(64, 48)[1] * 2.0))): |
| ang = th[:, None] * ph.m[None] - p[:, None] * ph.n[None] |
| buf(name + "_cos", torch.cos(ang).T) |
| buf(name + "_sin", torch.sin(ang).T) |
| buf(name + "_cp", torch.cos(p)) |
| buf(name + "_sp", torch.sin(p)) |
| buf("m", ph.m) |
| buf("n", ph.n) |
| b64, b128 = coil_basis(64), coil_basis(128) |
| buf("b0", b64[0]) |
| buf("b1", b64[1] * (2 * np.pi / 64)) |
| buf("c1", b128[1]) |
| buf("c2", b128[2]) |
| buf("maps", ph.maps) |
| buf("sign", ph.sign) |
| mine = torch.arange(k).repeat_interleave(64) |
| theirs = torch.cat([mine, torch.full(((ph.J - 1) * k * 64,), -1)]) |
| self.register_buffer("mine", mine.float()) |
| self.register_buffer("theirs", theirs.float()) |
|
|
| def points(self, rc, zs, name): |
| c, s, cp, sp = (getattr(self, f"{name}_{v}") for v in ("cos", "sin", "cp", "sp")) |
| R, Z = rc @ c, zs @ s |
| return R, Z, cp, sp |
|
|
| def forward(self, coef, cur, rc, zs, target): |
| """coef (K, 3, 15), cur (K,) in megaamperes, rc and zs (45,), target (4,): the field feedback (X,) and the |
| slacks (3 K,).""" |
| k = self.k |
| |
| R, Z, _, _ = self.points(rc, zs, "a") |
| R, Z = R.reshape(48, 64), Z.reshape(48, 64) |
| area = 0.5 * (R * torch.roll(Z, -1, 1) - torch.roll(R, -1, 1) * Z).sum(1).abs() |
| a = (area.mean() / np.pi).sqrt() |
| |
| R, Z, cp, sp = self.points(rc, zs, "f") |
| Rt, Zt = -(rc * self.m) @ self.f_sin, (zs * self.m) @ self.f_cos |
| Rp, Zp = (rc * self.n) @ self.f_sin, -(zs * self.n) @ self.f_cos |
| x = torch.stack([R * cp, R * sp, Z], -1) |
| nrm = cross(torch.stack([Rt * cp, Rt * sp, Zt], -1), torch.stack([Rp * cp - R * sp, Rp * sp + R * cp, Zp], -1)) |
| nrm = nrm / (nrm * nrm).sum(-1, keepdim=True).sqrt() |
| r, e = (coef @ self.b0).transpose(1, 2), (coef @ self.b1).transpose(1, 2) |
| rj = torch.einsum("jac,kpc->jkpa", self.maps, r) |
| ej = torch.einsum("jac,kpc->jkpa", self.maps, e) |
| w = (self.sign[:, None] * cur[None, :])[..., None, None] |
| we, rq = (w * ej).reshape(-1, 3), rj.reshape(-1, 3) |
| wc = cross(we, rq) |
| d2 = (x * x).sum(-1)[:, None] + (rq * rq).sum(-1)[None, :] - 2.0 * x @ rq.T |
| inv = d2.clamp(min=1e-10).rsqrt() |
| inv3 = inv * inv * inv |
| field = cross(inv3 @ we, x) - inv3 @ wc |
| signed = (field * nrm).sum(-1) / (field * field).sum(-1).sqrt().clamp(min=1e-30) / target[3] |
| |
| u = r.reshape(-1, 3) |
| d2 = (u * u).sum(-1)[:, None] + (rq * rq).sum(-1)[None, :] - 2.0 * u @ rq.T + 1e6 * (self.mine[:, None] == self.theirs[None, :]).float() |
| cc = d2.reshape(k, -1).amin(1).clamp(min=0).sqrt() / (a * target[0]) |
| R, Z, cp, sp = self.points(rc, zs, "s") |
| surf = torch.stack([R * cp, R * sp, Z], -1) |
| d2 = (u * u).sum(-1)[:, None] + (surf * surf).sum(-1)[None, :] - 2.0 * u @ surf.T |
| cp_ = d2.reshape(k, -1).amin(1).clamp(min=0).sqrt() / (a * target[1]) |
| r1, r2 = (coef @ self.c1).transpose(1, 2), (coef @ self.c2).transpose(1, 2) |
| kap = (cross(r1, r2) ** 2).sum(-1).sqrt() / (r1 * r1).sum(-1).sqrt().clamp(min=1e-12) ** 3 |
| return signed, torch.cat([cc, cp_, kap.amax(1) * (a / target[2])]) - 1.0 |
|
|
|
|
| class Wrapped(nn.Module): |
| def __init__(self, ck, base): |
| super().__init__() |
| self.k, self.rounds = int(base["k"]), int(ck["rounds"]) |
| n_x, n_y = int(base["varies"].sum()), self.k * 46 |
| self.net = Net(n_x, n_y, int(base["width"]), int(base["depth"])) |
| self.net.load_state_dict(base["net"]) |
| self.fix = Corrector(n_x, n_y, 32 * int(round(3.2 * self.k)), 3 * self.k, int(ck["width"]), int(ck["depth"])) |
| self.fix.load_state_dict(ck["fix"]) |
| self.feedback = Feedback(int(base["nfp"]), self.k) |
| self.register_buffer("pick", torch.nonzero(base["varies"]).reshape(-1)) |
| for name in ("x_mean", "x_std", "y_mean", "y_std"): |
| self.register_buffer(name, base[name].float()) |
|
|
| def forward(self, boundary_r_cos, boundary_z_sin, requirement_metrics, n_field_periods, n_coils_per_half_period): |
| rc, zs = boundary_r_cos.reshape(45), boundary_z_sin.reshape(45) |
| x = torch.cat([rc, zs, requirement_metrics.log()]) |
| x = ((x.index_select(0, self.pick) - self.x_mean) / self.x_std).reshape(1, -1) |
| y = self.net(x) |
| for _ in range(self.rounds): |
| full = y.reshape(-1) * self.y_std + self.y_mean |
| field, slack = self.feedback(full[:self.k * 45].reshape(self.k, 3, 15), full[self.k * 45:], rc, zs, requirement_metrics) |
| y = y + self.fix(x, y, field.reshape(1, -1), slack.reshape(1, -1)) |
| full = y.reshape(-1) * self.y_std + self.y_mean |
| full = full + 0.0 * (n_field_periods + n_coils_per_half_period).to(full.dtype) |
| coef = full[:self.k * 45].reshape(self.k, 3, 15) |
| return coef[:, 0], coef[:, 1], coef[:, 2], full[self.k * 45:] * 1e6 |
|
|
|
|
| def main(name, out, *flags): |
| ck = torch.load(HERE / "runs" / f"{name}.pt", map_location="cpu") |
| base = torch.load(HERE / "runs" / f"{ck['base']}.pt", map_location="cpu") |
| model = Wrapped(ck, base).eval() |
| if "--check" in flags: |
| import fast |
| from train import Data |
| n = int(flags[flags.index("--check") + 1]) |
| data, ph = Data(int(base["nfp"]), model.k, "cuda"), Physics(int(base["nfp"]), model.k, "cuda") |
| rows = data.dev_rows[torch.linspace(0, len(data.dev_rows) - 1, n, device="cuda").long()] |
| keys, target = data.keys[data.boundary[rows]], data.target[rows] |
| coef, cur = data.coef[rows].reshape(-1, model.k, 3, 15), data.cur[rows] |
| f, s = fast.feedback(ph, coef, cur, keys[:, :45], keys[:, 45:], target, ph.minor_radius(keys[:, :45], keys[:, 45:])) |
| worst = [0.0, 0.0] |
| with torch.no_grad(): |
| for i in range(n): |
| g = model.feedback(coef[i].cpu(), cur[i].cpu(), keys[i, :45].cpu(), keys[i, 45:].cpu(), target[i].cpu()) |
| worst = [max(worst[0], float((g[0] - f[i].cpu()).abs().max())), max(worst[1], float((g[1] - s[i].cpu()).abs().max()))] |
| print(f"feedback against the kernels on {n} dev instances: largest difference in the field {worst[0]:.2e}, in the slacks {worst[1]:.2e}") |
| feed = (torch.randn(5, 9) * 0.02, torch.randn(5, 9) * 0.02, torch.tensor([0.4, 0.8, 3.0, 0.05]), torch.tensor(3), torch.tensor(5)) |
| feed[0][0, 4], feed[0][1, 4], feed[1][1, 4] = 1.0, 0.15, 0.15 |
| names_in = ["boundary_r_cos", "boundary_z_sin", "requirement_metrics", "n_field_periods", "n_coils_per_half_period"] |
| names_out = ["coil_x_n", "coil_y_n", "coil_z_n", "currents"] |
| torch.onnx.export(model, feed, out, input_names=names_in, output_names=names_out, opset_version=17, dynamo=False, do_constant_folding=False) |
| proto = onnx.load(out) |
| proto.metadata_props.append(onnx.StringStringEntryProto(key="supported_tracks", value="fixed_shape")) |
| onnx.save(proto, out) |
| onnx.checker.check_model(out) |
| opts = ort.SessionOptions() |
| opts.intra_op_num_threads = 1 |
| sess = ort.InferenceSession(out, opts, providers=["CPUExecutionProvider"]) |
| inputs = {n: (v.numpy().astype(np.int64) if v.dtype == torch.int64 else v.numpy()) for n, v in zip(names_in, feed)} |
| got = sess.run(None, inputs) |
| t0 = time.perf_counter() |
| for _ in range(20): |
| sess.run(None, inputs) |
| took = (time.perf_counter() - t0) / 20 |
| with torch.no_grad(): |
| want = model(*feed) |
| worst = max(float(np.abs(g - w.numpy()).max() / max(1.0, float(w.abs().max()))) for g, w in zip(got, want)) |
| print(f"wrote {out}: {Path(out).stat().st_size / 2 ** 20:.1f} MiB, {model.rounds} round(s), dev feasible in training {ck.get('dev_feasible')}, " |
| f"largest relative difference from PyTorch {worst:.2e}, {1e3 * took:.1f} ms a run on one thread") |
|
|
|
|
| if __name__ == "__main__": |
| main(*sys.argv[1:]) |
|
|
