File size: 6,397 Bytes
2b87804
 
39d7e43
 
 
 
 
 
 
 
2b87804
39d7e43
3dc632e
39d7e43
3dc632e
 
 
 
39d7e43
3dc632e
39d7e43
3dc632e
 
 
 
39d7e43
3dc632e
 
 
 
39d7e43
3dc632e
39d7e43
3dc632e
 
 
 
 
 
 
 
 
 
39d7e43
 
3dc632e
39d7e43
3dc632e
 
 
 
 
 
 
39d7e43
 
3dc632e
39d7e43
3dc632e
39d7e43
3dc632e
 
39d7e43
3dc632e
 
 
 
 
39d7e43
3dc632e
 
39d7e43
3dc632e
39d7e43
3dc632e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39d7e43
 
 
3dc632e
39d7e43
3dc632e
39d7e43
 
 
3dc632e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
---
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.