phanerozoic/coilstellaration-coil-optimizer

Coil optimizer for Coilstellaration

model.onnx is an entry for the fixed_shape track of the Coilstellaration benchmark (three field periods, five coils per half period). The benchmark asks for a stellarator coilset given a plasma boundary and four engineering targets: the smallest coil-coil distance, the smallest coil-plasma distance, the largest coil curvature and the largest normalized field error on the boundary.

The graph runs in three stages. The first gathers three starts. One is the coilset of the per-coil ensemble per_coil_ensemble_DAr7gipsQu9Mupbnsm2rAaK.onnx of proxima-fusion/coilstellaration-models (MIT licence), embedded with its weights unchanged; it is asked for the instance's own targets, or, where the field-error target is above 0.1, for 1.2 times the two distance targets (at least 0.40 and 0.60 minor radii), 0.85 times the curvature target (at most 3.5) and a field error of at most 0.01. The other two come from the training split of proxima-fusion/coilstellaration (MIT licence), whose 53,206 coilsets of this shape the file holds in half precision: of each of the two training boundaries nearest to the instance's in their Fourier coefficients, the coilset whose targets are closest to the instance's.

The second stage is a Levenberg-Marquardt optimizer on the Fourier coefficients and currents of the five coils, written as ONNX Loops. Its residuals are the part of the normal field over the field strength above 0.85 of the target, at 800 points of the boundary's half period, and penalties on every coil point that is closer to another coil or to the boundary than the distance targets allow, or more curved than the curvature target allows. The 800 points lie on 40 lines of constant poloidal angle, 20 along the half period on each, and each line's points sit 0.618 of their spacing further along than those of the line before, less whole spacings. They are split into four groups spread over the boundary, and each group reads the coils at quadrature points a quarter of their spacing beyond the group before, so that an error of the 40-point quadrature shows as a difference between neighbouring points. The points go through four positions from step to step, half a spacing apart in each angle, and the coils' quadrature points, the coil points and the boundary points of the distance penalties move from step to step as well, so that a residual between the points of one step is met by a later one. The coil-coil distances are taken against the thirty coils of all six images of the five, and the coil-plasma distance asked for is at least 0.3 minor radii. The Jacobians are analytic. The Gauss-Newton matrix also counts each distance and curvature within a tenth of its limit, which holds a step back from a limit it has not yet reached, and each step solves its linear system by twelve conjugate-gradient iterations. A step is kept when the objective on the same points is lower at its end; that check takes the field of the coils' three images nearest the half period anew and keeps the field of the other three from before the step, scaled by each coil's change of current. Each start takes ten steps, and the start whose objective is then lowest takes up to 149 more. Those thirty first steps use fewer points: 240 on the boundary, on 20 lines, 24 quadrature points a coil, and 32 points a coil for the distances and 160 for the curvature, where the later steps have 64 and 400. A start's run ends once its last four steps, one per position of the field grid, each found the field error at most 0.9 of its target and no distance or curvature beyond its target, or at its last step, and it keeps the coilset that step evaluated; at its step 24 and every 45th after it, its field threshold is lowered by as much as its field error over the last four steps exceeds 0.9 of the target. While no start has ended within its targets, the other two, in the order of their objectives after their ten steps, each take up to 250 steps more, and after them the planar circles of the third stage take up to 420 steps as a fourth start. The coilset kept is the first that ended within its targets, else the one whose last field error is lowest among those within their distance and curvature targets.

The third stage returns planar circles sized to the two distance targets when they meet all four targets on the optimizer's grids, and the optimized coilset otherwise; the loops are skipped when the circles are returned.

Every operation of the graph is one that onnxruntime 1.31 runs in the calling thread: each matrix product is two-dimensional and below 65,536 in rows times columns times inner size, and the graph holds no Transpose and no Clip. A run therefore uses one processor whatever the size of the session's thread pool.

Interface

The graph follows the benchmark's ONNX contract. Its inputs are boundary_r_cos and boundary_z_sin (float32, [5, 9]), requirement_metrics (float32, [4]), and n_field_periods and n_coils_per_half_period (int64 scalars, 3 and 5). Its outputs are coil_x_n, coil_y_n and coil_z_n (float32, [5, 15], Fourier order 7 in DESC's [sin(7) .. sin(1), const, cos(1) .. cos(7)] layout) and currents (float32, [5], amperes). The metadata key supported_tracks is fixed_shape.

import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download

session = ort.InferenceSession(hf_hub_download("phanerozoic/coilstellaration-coil-optimizer", "model.onnx"))
r_cos = np.zeros((5, 9), np.float32)
z_sin = np.zeros((5, 9), np.float32)
r_cos[0, 4], r_cos[1, 4], z_sin[1, 4] = 1.0, 0.15, 0.15
coil_x_n, coil_y_n, coil_z_n, currents = session.run(
    ["coil_x_n", "coil_y_n", "coil_z_n", "currents"],
    {
        "boundary_r_cos": r_cos,
        "boundary_z_sin": z_sin,
        "requirement_metrics": np.array([0.4, 0.8, 3.0, 0.05], np.float32),
        "n_field_periods": np.asarray(3, np.int64),
        "n_coils_per_half_period": np.asarray(5, np.int64),
    },
)

variable_shape

variable.onnx is an entry for the variable_shape track (two to five field periods, four to eight coils per half period, seventeen shapes) and for fixed_shape. It holds the optimizer once per shape under If nodes on the two shape inputs. A run has up to two starts, taken in turn through the same steps: the first-guess network's coilset, and, if that start does not end within its targets, the training coilset held for the training boundary nearest to the instance's in its Fourier coefficients. The file holds one training coilset per training boundary of each shape, the one whose field-error target is lowest, 31,754 in all, each coefficient as a byte q that stands for its column's median plus q |q| times the column's scale; coilsets with a coefficient beyond 3 are left out. Its copy of the first-guess network holds the weights in half precision. The coilset returned is that of the first start that ended within its targets, else the one whose last field error is lowest among those within their distance and curvature targets. A start is shorter and coarser than one of model.onnx, so that the 10,829 runs of the two tracks fit the evaluator's four hours. Its steps are of two kinds. Up to 30 coarse steps come first, on a field grid of 20 points around the boundary by 2.4 per coil along the half period, with 24 quadrature points a coil, 32 coil points for the distances and 160 for the curvature; they end early once every position of their grid reads the coilset within its targets, and a start whose coarse steps last read the field error above 2.5 of its target takes no further step. Up to 30 fine steps follow, on a grid of 32 points by 3.2 per coil with 40 quadrature points, 64 coil points for the distances and 400 for the curvature, and a start that reaches the last of them with the field error at most 1.25 of its target takes up to 12 more. Both kinds have every line's points at the same toroidal angles and one set of quadrature points, take the coil-coil distances against the three nearby images, count each distance and curvature in the Gauss-Newton matrix from 0.12 below its limit, and judge a step's trial point with the field of the coils' images beyond the three nearest the half period kept from the point before the step, scaled by each coil's change of current. Each step is solved together with the three steps before it of its kind: their Jacobians, taken on the other three positions of the grid, and their field errors, moved to first order by the steps taken since, enter the linear system beside the step's own, and their part of the objective, to first order, counts in judging the trial point. The system takes 40 conjugate-gradient iterations, where that of model.onnx takes twelve. The coil-plasma distance asked for is 3.5% beyond its target. The damping of the first step is 1, where that of model.onnx is 3, and the field threshold follows the field error at step 16 and every twelfth after it. A shape outside the seventeen gets the first-guess network's coilset. Its inputs are those of model.onnx with any of the track's shapes, its outputs have one row per coil, and its metadata key supported_tracks is fixed_shape,variable_shape. wrap.py writes it with these arguments:

python wrap.py variable.onnx all --short-names --half-network --with-dev --set NEIGHBOURS=1 --set FALLBACKS=1 --set STAGE=0 --set TURNS=1 --set DB_BYTES=1 --set DB_PER=1 --set SLOTS=1 --set RING=0 --set GRID_T=32 --set GRID_PK=3.2 --set GRID_SHEAR=-1 --set PHASES=1 --set CC_ALL=0 --set ADAPT_FROM=16 --set ADAPT_EVERY=12 --set EPS_H=0.12 --set LAMBDA0=1.0 --set FAR_HOLD=1 --set COARSE_EXIT=1 --set COARSE_P_F=24 --set COARSE_P_K=160 --set COARSE=30 --set ITERS=60 --set EXTRA=12 --set EXTRA_AT=1.25 --set ABANDON=2.5 --set MEMORY=3 --set CG_ITERS=40 --set MARGIN_CP=0.035

Files

lm.py is the three stages with the settings model.onnx was written with: its functions run on PyTorch tensors and, unchanged, on the tensors of sym.py, which write the ONNX nodes. wrap.py writes the graph and joins it to the first-guess network, whose matrix products it cuts into blocks of columns. Both read the training coilsets and example instances from local tables of the dataset's training split; every setting was chosen on a twentieth of the training boundaries, whose coilsets were left out of the lookup while it was.