--- license: mit library_name: onnx tags: - onnx - fusion - plasma-physics - tokamak - diii-d - mamba --- # FUSE-EFIT (in era), ONNX A 20 Hz, strictly causal model of the DIII-D plasma edge, driven by the heating, fuelling and field actuators, the measured plasma current, the loop voltage and the EFIT01 boundary shape. Physical units in, physical units out. Code, documentation and the input contract: https://github.com/NoahH72/FUSE-Pedestal-Predictor-Mamba-ONNX (`docs/EXPORT.md`, `docs/INPUT_CONTRACT.md`, `MODEL_CARD.md`). ## Two forms, one calling convention | form | what it is | networks per tick | ONNX files | step call, 1 thread | step call, 8 threads | |---|---|---:|---:|---|---| | `single` | replicate 0 of the in-era finals (seed 100), the frozen delivered network | 1 | 184 MB | 4.87 ms (99th percentile 5.08 ms) | 1.95 ms (99th percentile 2.00 ms) | | `averaged` | the eight in-era finals (seeds 100-107); the mean of their raw outputs at every tick | 8 | 1468 MB | 56.33 ms (99th percentile 59.45 ms) | 18.25 ms (99th percentile 18.60 ms) | Timing: Intel(R) Xeon(R) CPU Max 9470C, onnxruntime 1.30.0, CPU, batch 1. The control cycle is 50 ms per tick. The two forms take and return the same tensors with the same names, shapes and units; only the recurrent state is larger for `averaged` (19.9 MB per shot against 2.5 MB). A host switches form by choosing a folder. ## Files ``` era/single/fuse_efit_single_step.onnx one 50 ms tick per call, state in and out era/single/fuse_efit_single_sequence.onnx one window of 256 ticks era/averaged/fuse_efit_averaged_step.onnx era/averaged/fuse_efit_averaged_sequence.onnx era/
/model_config.json the input and output contract, machine-readable era//norms.json the statistics baked into the graphs era//input_ranges.json training ranges per channel, for a units check era//provenance.json source checkpoints with sha256, code identity, versions SHA256SUMS sha256sum -c SHA256SUMS ``` ## Download and check ```bash git clone https://github.com/NoahH72/FUSE-Pedestal-Predictor-Mamba-ONNX cd FUSE-Pedestal-Predictor-Mamba-ONNX pip install numpy onnxruntime huggingface_hub hf download SCS-Lab/FUSE-PED --local-dir export # this page's files, into the code's export/ folder cd export && sha256sum --check SHA256SUMS && cd .. python examples/verify_integration.py --bundle export/era ``` Every line of the checksum check must say OK, and the integration check must pass, before the model is used. ## Run it ```python from fuse_efit.export.runtime import FuseEfitPredictor # numpy + onnxruntime only model = FuseEfitPredictor("export/era", form="single") # or form="averaged" state = model.init_state() # zeros at the start of every shot for each 50 ms tick: y, state = model.step(inputs_25, indicators_5, state) # exactly once per tick ``` Step graph tensors: `inputs` (batch, 25), `indicators` (batch, 5), `available` (batch, 25), `conv_state`, `ssm_state` in; `predictions` (batch, 9), `conv_state_out`, `ssm_state_out` out. ## Inputs (physical units, this order) | # | channel | unit | |--:|---|---| | 0 | `pinj` | MW | | 1 | `tinj` | N m | | 2 | `ech_mw` | MW | | 3 | `gasa_cal` | Torr L/s | | 4 | `gasb_cal` | Torr L/s | | 5 | `gasc_cal` | Torr L/s | | 6 | `gasd_cal` | Torr L/s | | 7 | `gase_cal` | Torr L/s | | 8 | `bt` | T | | 9 | `ip` | MA | | 10 | `aminor` | m | | 11 | `rsurf` | m | | 12 | `zsurf` | m | | 13 | `kappa` | dimensionless | | 14 | `tritop` | dimensionless | | 15 | `tribot` | dimensionless | | 16 | `drsep` | m | | 17 | `rvsout` | m | | 18 | `gapin` | m | | 19 | `rvsou` | m | | 20 | `rvsiu` | m | | 21 | `shafranov` | m | | 22 | `zeta_ou` | dimensionless | | 23 | `zeta_ol` | dimensionless | | 24 | `vloop` | V | Indicators, 0/1 per tick: `efit_valid`, `vloop_valid`, `rvsout_present`, `rvsou_present`, `rvsiu_present`. Availability, 0/1 per shot and channel (default 1). Where an indicator or the availability flag is 0 the value is ignored; elsewhere supply a real value (not-a-number is read as the training mean). ## Outputs | # | head | unit | |--:|---|---| | 0 | `ne` | 1e19 m^-3 | | 1 | `te_ped` | keV | | 2 | `ti_ped` | keV | | 3 | `t_rot_ped` | krad/s | | 4 | `ne_knee` | 1e19 m^-3 | | 5 | `te_knee` | keV | | 6 | `rho_sym` | rho (normalised radius) | | 7 | `ne_top_loc` | rho (normalised radius) | | 8 | `hmode` | probability in [0, 1] | `ne_knee`, `te_knee`, `rho_sym` and `ne_top_loc` are defined only in H-mode: report them only where `hmode` is at or above a threshold. Each form has its own, set on validation: `single` 0.54 (balanced) and 0.5 (conservative); `averaged` 0.53 and 0.47. On the test shots each keeps about 95-96 % of H-mode ticks with about 5 % L-mode among those kept (`MODEL_CARD.md` §6). ## Accuracy Typical error per test shot (the median over in-era test shots of each shot's RMSE, physical units; `hmode`: its loss), from the single opening of the test set: | output | unit | `single` | `averaged` | |---|---|---|---| | `ne` | 1e19 m^-3 | 0.2928 | 0.2594 | | `te_ped` | keV | 0.05405 | 0.04808 | | `ti_ped` | keV | 0.09957 | 0.08967 | | `t_rot_ped` | krad/s | 4.043 | 3.678 | | `ne_knee` | 1e19 m^-3 | 0.3279 | 0.2973 | | `te_knee` | keV | 0.06058 | 0.05636 | | `rho_sym` | rho (normalised radius) | 0.01047 | 0.009892 | | `ne_top_loc` | rho (normalised radius) | 0.0125 | 0.01188 | | `hmode` | loss | 0.1373 | 0.1251 | | combined loss | | 0.9889 | 0.8986 | ## Export accuracy Every graph agrees with the PyTorch model to 7.4e-06 of each head's training standard deviation or better, on 1,988 validation ticks per form and 2,000-tick runs. ## Limits Needs EFIT01 geometry, the between-shot reconstruction, as trained; it is not for use during the shot. Fed real-time EFIT geometry instead it is about 30 % worse, and far worse on shots where the real-time reconstruction lacks the strike points. No uncertainty estimate. A unit error gives confident nonsense, not an error: run `examples/verify_integration.py` and `check_units` once. Full list: `MODEL_CARD.md` §8. ## Licence and contact Licence: MIT (the `LICENSE` file). Contact: Noah Harris, harrnoah@oregonstate.edu. Code: https://github.com/NoahH72/FUSE-Pedestal-Predictor-Mamba-ONNX.