"""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) # noqa: E731 self.t = t self.m = t(np.repeat(np.arange(5), 9)) # (45,) 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) # (J, 3, 3): T_j v = maps[j] @ v self.sign = t(np.tile([1.0, -1.0], nfp)) # (J,) 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] # the boundary 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] # (X, 45) c, s = torch.cos(ang).T, torch.sin(ang).T # (45, X) 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 # 48 sections over a full period 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() # the coils 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) # (B, J, K, P, 3) w = (self.sign[None, :, None] * cur[:, None, :])[..., None, None] # (B, J, K, 1, 1) 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) # (B, X, Q) field = torch.linalg.cross(inv3 @ we, x) - inv3 @ wc # (B, X, 3) 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)