| --- |
| 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/<form>/model_config.json the input and output contract, machine-readable |
| era/<form>/norms.json the statistics baked into the graphs |
| era/<form>/input_ranges.json training ranges per channel, for a units check |
| era/<form>/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. |
|
|
