| """The benchmark's four metrics of a batch of coilsets, in PyTorch on the GPU, differentiable. |
| |
| A coilset is K unique coils of Fourier order 7, (B, K, 3, 15) in DESC's layout [sin 7t .. sin t, 1, cos t .. cos 7t], |
| with currents (B, K); the full coilset is its 2 NFP images (a turn by a field period, and the partner |
| (x, -y, -z) with the current reversed). A boundary is R = sum r_cos[m, n] cos(m theta - n NFP phi), |
| Z = sum z_sin[m, n] sin(m theta - n NFP phi), its coefficients flattened (B, 45) in (m, n) order. |
| |
| Physics(nfp, k) gives: surface() the boundary's points and unit normals, minor_radius(), field_error() the normal |
| field over the field strength at surface points, pair distances for the coil-coil and coil-plasma metrics, and |
| curvature(). metrics() returns the four achieved values normalised as the benchmark's targets are. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
|
|
| ORDER = 7 |
| NM = 2 * ORDER + 1 |
|
|
|
|
| def coil_basis(points): |
| """Values, first and second derivatives of the coil modes at `points` equal parameters, (15, points) each.""" |
| t = 2 * np.pi * np.arange(points) / points |
| n = np.arange(1, ORDER + 1)[:, None] |
| s, c, z = np.sin(n * t), np.cos(n * t), np.zeros((1, points)) |
| return (np.concatenate([s[::-1], z + 1, c]), np.concatenate([(n * c)[::-1], z, -n * s]), |
| np.concatenate([-(n * n * s)[::-1], z, -n * n * c])) |
|
|
|
|
| class Physics: |
| def __init__(self, nfp, k, device="cuda", dtype=torch.float32): |
| self.nfp, self.k, self.device, self.dtype = nfp, k, device, dtype |
| t = lambda a: torch.as_tensor(np.asarray(a), dtype=dtype, device=device) |
| self.t = t |
| self.m = t(np.repeat(np.arange(5), 9)) |
| self.n = t(np.tile(np.arange(-4, 5) * nfp, 5)) |
| ang = 2 * np.pi * np.arange(nfp) / nfp |
| maps = np.stack([np.array([[np.cos(a), -np.sin(a), 0], [np.sin(a), np.cos(a), 0], [0, 0, 1]]) @ np.diag(f) |
| for a in ang for f in ([1.0, 1.0, 1.0], [1.0, -1.0, -1.0])]) |
| self.maps = t(maps) |
| self.sign = t(np.tile([1.0, -1.0], nfp)) |
| self.J = 2 * nfp |
| self._basis = {} |
| self._radius = None |
|
|
| def basis(self, points): |
| if points not in self._basis: |
| self._basis[points] = tuple(self.t(b) for b in coil_basis(points)) |
| return self._basis[points] |
|
|
| |
| def surface(self, rc, zs, theta, phi): |
| """Points and unit normals at angles theta, phi (X,) each: (B, X, 3) each.""" |
| ang = theta[:, None] * self.m[None] - phi[:, None] * self.n[None] |
| c, s = torch.cos(ang).T, torch.sin(ang).T |
| R, Z = rc @ c, zs @ s |
| Rt, Zt = -(rc * self.m) @ s, (zs * self.m) @ c |
| Rp, Zp = (rc * self.n) @ s, -(zs * self.n) @ c |
| cp, sp = torch.cos(phi)[None], torch.sin(phi)[None] |
| x = torch.stack([R * cp, R * sp, Z], -1) |
| et = torch.stack([Rt * cp, Rt * sp, Zt], -1) |
| ep = torch.stack([Rp * cp - R * sp, Rp * sp + R * cp, Zp], -1) |
| nrm = torch.linalg.cross(et, ep) |
| return x, nrm / nrm.norm(dim=-1, keepdim=True) |
|
|
| def surface_points(self, rc, zs, theta, phi): |
| """The boundary's points at angles theta, phi (X,) each, (B, X, 3).""" |
| ang = theta[:, None] * self.m[None] - phi[:, None] * self.n[None] |
| R, Z = rc @ torch.cos(ang).T, zs @ torch.sin(ang).T |
| return torch.stack([R * torch.cos(phi)[None], R * torch.sin(phi)[None], Z], -1) |
|
|
| def grid(self, n_theta, n_phi, du=0.0, dv=0.0, full=False): |
| """Angles of an n_theta by n_phi grid over the half period (or with full the whole torus), moved by du and |
| dv of its spacings: theta, phi (X,) each.""" |
| theta = 2 * np.pi * (torch.arange(n_theta, device=self.device, dtype=torch.float64) + du) / n_theta |
| span = 2 * np.pi if full else np.pi / self.nfp |
| phi = span * (torch.arange(n_phi, device=self.device, dtype=torch.float64) + dv) / n_phi |
| return theta.repeat(n_phi).to(self.dtype), phi.repeat_interleave(n_theta).to(self.dtype) |
|
|
| def minor_radius(self, rc, zs): |
| """The radius of the circle whose area is the mean cross-section's, (B,).""" |
| if self._radius is None: |
| theta, phi = self.grid(64, 48, full=False) |
| phi = phi * 2.0 |
| ang = theta[:, None] * self.m[None] - phi[:, None] * self.n[None] |
| self._radius = (torch.cos(ang).T.contiguous(), torch.sin(ang).T.contiguous()) |
| R, Z = (rc @ self._radius[0]).reshape(-1, 48, 64), (zs @ self._radius[1]).reshape(-1, 48, 64) |
| Rn, Zn = torch.roll(R, -1, 2), torch.roll(Z, -1, 2) |
| area = 0.5 * (R * Zn - Rn * Z).sum(2).abs() |
| return (area.mean(1) / np.pi).sqrt() |
|
|
| |
| def coil_points(self, coef, points, order=0): |
| """The unique coils at `points` equal parameters, (B, K, points, 3): the points, or a derivative.""" |
| return (coef @ self.basis(points)[order]).transpose(2, 3) |
|
|
| def images(self, v): |
| """(B, K, P, 3) under every map: (B, J, K, P, 3).""" |
| return torch.einsum("jac,bkpc->bjkpa", self.maps, v) |
|
|
| def field_error(self, coef, cur, x, nrm, points=64): |
| """|B . n| / |B| at the surface points x with unit normals nrm, (B, X).""" |
| B = coef.shape[0] |
| r, e = self.coil_points(coef, points), self.coil_points(coef, points, 1) * (2 * np.pi / points) |
| rj, ej = self.images(r), self.images(e) |
| w = (self.sign[None, :, None] * cur[:, None, :])[..., None, None] |
| we = (w * ej).reshape(B, -1, 3) |
| wc = (w * torch.linalg.cross(ej, rj)).reshape(B, -1, 3) |
| rq = rj.reshape(B, -1, 3) |
| d2 = (x * x).sum(-1)[:, :, None] + (rq * rq).sum(-1)[:, None, :] - 2.0 * x @ rq.transpose(1, 2) |
| inv3 = d2.clamp(min=1e-10).pow(-1.5) |
| field = torch.linalg.cross(inv3 @ we, x) - inv3 @ wc |
| return (field * nrm).sum(-1).abs() / field.norm(dim=-1).clamp(min=1e-30) |
|
|
| def coil_coil(self, coef, points=64): |
| """Squared distances from each unique coil's points to the points of every other coil of the full coilset, |
| (B, K P, J K P), a coil's distances to itself set large.""" |
| B = coef.shape[0] |
| r = self.coil_points(coef, points) |
| u, v = r.reshape(B, -1, 3), self.images(r).reshape(B, -1, 3) |
| d2 = (u * u).sum(-1)[:, :, None] + (v * v).sum(-1)[:, None, :] - 2.0 * u @ v.transpose(1, 2) |
| own = torch.zeros(self.k, points, self.J, self.k, points, dtype=torch.bool, device=coef.device) |
| for k in range(self.k): |
| own[k, :, 0, k, :] = True |
| return d2.masked_fill(own.reshape(1, self.k * points, -1), 1e6) |
|
|
| def coil_plasma(self, coef, surface_points, points=64): |
| """Squared distances from the unique coils' points to points of the whole boundary, (B, K P, S).""" |
| B = coef.shape[0] |
| u = self.coil_points(coef, points).reshape(B, -1, 3) |
| return (u * u).sum(-1)[:, :, None] + (surface_points * surface_points).sum(-1)[:, None, :] - 2.0 * u @ surface_points.transpose(1, 2) |
|
|
| def curvature(self, coef, points=128): |
| """The unique coils' curvature at `points` equal parameters, (B, K, points).""" |
| r1, r2 = self.coil_points(coef, points, 1), self.coil_points(coef, points, 2) |
| return torch.linalg.cross(r1, r2).norm(dim=-1) / r1.norm(dim=-1).clamp(min=1e-12).pow(3) |
|
|
| def metrics(self, coef, cur, rc, zs, dense=True): |
| """The four achieved metrics as the benchmark normalises them, (B, 4): least coil-coil and coil-plasma |
| distance over the minor radius, largest curvature times it, largest field error.""" |
| a = self.minor_radius(rc, zs) |
| nt, npk, pf, pg, pk, st, sp = (64, 6.4, 160, 128, 256, 64, 32) if dense else (32, 3.2, 64, 64, 128, 32, 16) |
| theta, phi = self.grid(nt, int(round(npk * self.k))) |
| x, nrm = self.surface(rc, zs, theta, phi) |
| ferr = self.field_error(coef, cur, x, nrm, pf).max(1).values |
| cc = self.coil_coil(coef, pg).amin((1, 2)).clamp(min=0).sqrt() / a |
| theta, phi = self.grid(st, sp * self.J, full=True) |
| cp = self.coil_plasma(coef, self.surface(rc, zs, theta, phi)[0], pg).amin((1, 2)).clamp(min=0).sqrt() / a |
| curv = self.curvature(coef, pk).amax((1, 2)) * a |
| return torch.stack([cc, cp, curv, ferr], 1) |
|
|
