FusionLab ⚛️

An open-data tokamak digital twin. Replay real shots of the UK's MAST tokamak beside a physics model, see where the model is wrong, learn a correction with NVIDIA PhysicsNeMo and report its error honestly, trace the magnetic field lines on the GPU with NVIDIA Warp, and take the shot away as a time-sampled OpenUSD stage. Runs on one workstation, offline, from a clone.

FusionLab is an educational and exploratory twin built on a reduced (0D) model. It compares with the experiment; it does not predict it.

Quick start

Linux, macOS or Windows. An NVIDIA GPU is optional — everything runs CPU-only without one. The repo ships the data cache and the trained models, so nothing needs the network or an API key after the clone.

Install uv if you don't have it. uv also fetches the Python 3.11–3.12 interpreter this project needs, so there is nothing else to install first:

curl -LsSf https://astral.sh/uv/install.sh | sh

Then, from the repo root:

make setup        # uv sync: numpy, torch, nvidia-physicsnemo, usd-core, zarr, s3fs, fastapi
make dev          # http://localhost:8000

No make? (Windows without build tools, say.) The same two steps by hand:

uv sync
cp .env.example .env                     # if you don't have a .env yet
uv run uvicorn fusionlab.api:app --reload --port 8000    # http://localhost:8000

(make setup also wires the git hooks: git config core.hooksPath .githooks.)

On a machine without an NVIDIA GPU, Warp prints Could not find or load the NVIDIA CUDA driver at startup. It is benign — field-line tracing just runs on the CPU.

Other commands:

make test	unit + API tests (uv run pytest -q)
make train	retrain the IPB98 correction, rewrites its metrics
make bench	timings
make usd	OpenUSD export to out/
make fieldlines	Warp field-line q check + benchmark
make validate-virtual	the virtual shot's gates: fetches ~120 held-out shots
make data	re-download the FAIR-MAST cache (the repo already ships it)
uv run python scripts/mode_flag_report.py	derived L/H flag counts + anchors (writes data/mode_flag_report.json)

from fusionlab import mast
from fusionlab.physics import replay

s = mast.load_shot(30166)     # dict of arrays on the EFIT time base, in MA, T, 1e20 m^-3, keV, MW, MJ
r = replay(s)                 # IPB98 / ITER89-P on measured P_loss, H98(t), Greenwald / Troyon / kink fractions
db = mast.clean_db(mast.load_db())   # 6,353 shots at peak current; mode_derived = derived H/L flag (PHYSICS.md)

uv run python scripts/fetch_mast.py 29182 caches any other level-2 shot (~27 s, ~1 MB).

The dev server reads the data cache and the model weights once per process: restart it after re-fetching shots (make data, fetch_mast.py) or retraining (make train) to pick up the new artifacts.

User-provided compute (SSH)

An NVIDIA GPU is optional, and it does not have to be this machine's. FusionLab can run the heavy parts — Warp field-line tracing, the PhysicsNeMo surrogates, benchmarks — on any SSH-reachable host you control: a lab box, a university cluster login, a gaming PC. A small worker process from this same repo runs there (git clone + uv sync --extra remote), the backend talks to it through the SSH tunnel using your existing keys and ~/.ssh/config aliases — no new credentials, nothing stored in FusionLab — and results carry a computed on <host> label. If the remote side is unreachable, work falls back to this machine with a visible banner; the UI never blocks.

# fusionlab.toml (copy fusionlab.toml.example) — or FUSIONLAB_COMPUTE_* env vars, see .env.example
[compute]
provider = "ssh"            # "local" (default): everything in this process, zero config

[compute.ssh]
host = "gpu-lab"            # any ~/.ssh/config alias
remote_dir = "~/FusionLab"  # the clone on that host
autostart = true            # start the worker over SSH when none is answering

More than one GPU box? Define a [compute.hosts.<name>] profile per host (same keys) and switch with FUSIONLAB_COMPUTE_HOST — see docs/COMPUTE.md.

By default the worker binds loopback on the remote host and trusts the tunnel. If it might be reachable by more than your own tunnel, set FUSIONLAB_WORKER_TOKEN to the same value on both ends (.env here, the worker's environment there) and every /v1 route on the worker demands a Bearer token — same file, no new services; details in docs/COMPUTE.md.

Full guide — prerequisites, what the chip and banner mean, and a first-run checklist against a real host: docs/COMPUTE.md.

Contents

Quick start · User-provided compute · Why · What's in it · A five-minute demo · Guided and Lab modes · Virtual shot · What it found on real data · Architecture · Performance · Data & licence · Limitations · Security posture · Optional AI agent · Credit

Why

NVIDIA's fusion twins with General Atomics (DIII-D), CFS/Siemens (SPARC) and UKAEA follow one pattern: measured data, AI surrogates, an Omniverse/OpenUSD scene. They are closed and were built on national-lab machines. FusionLab is that pattern on open data (UKAEA's FAIR-MAST archive, 15,969 shots). A researcher gets load_shot(30166) instead of a week of plumbing; a student gets a real discharge to scrub through and watch hit a real limit.

What's in it

Replay a real shot	One screen: EFIT flux surfaces inside the real vessel, Thomson T_e profile, stored energy against the IPB98(y,2) and ITER89-P scaling laws fed the measured power, live Greenwald / Troyon / kink gauges, the session leader's logbook note, and a three.js vessel view with Warp-traced field lines.
Guided study	For people new to fusion: a 15-minute lesson rail beside the live app. It lights one panel at a time, drives the app (picks the shot, scrubs, flips the 3D view, unlocks one slider) and sets tasks that tick themselves. Lab mode is the full app, with glossary tooltips and a ? on every panel that opens its lesson.
Shot database	6,353 real shots against IPB98(y,2), with the learned correction's holdout table beside it.
What-if sandbox	0D power balance on ITER, a SPARC-class design, JET, DIII-D and MAST: move a slider, watch Q, the limits and the operating map.
Virtual shot	Edit a real shot's programme (beam power and timing, density) and re-fly it through a time-dependent 0D energy balance anchored on the measurement. A slider exists only where matched pairs of real shots can test the edit; the result is worded as a distance to a limit, never a forecast.
Equilibrium surrogate	A PhysicsNeMo network reconstructs the flux map ψ(R,Z) from 93 magnetic signals; shown dashed over EFIT's surfaces with the per-slice error.
OpenUSD export	Each shot as a time-sampled stage: real vessel, PF coils, plasma boundary, field lines. OpenUSD export (Omniverse-compatible); opens in usdview, USD Composer or Blender.
CSV export	Every number a replay or compare shows — measured traces, model traces, summary — as one CSV, one row per shot, slice and quantity, each labelled measured / model / learned / computed.
Limit-distance panel	The replayed shot's distance to the Greenwald, Troyon βN and kink limits over time (1.0 = at the limit) — computed distances, never forecasts — with the archive's abort flag overlaid where the catalog carries one.
Python loader	mast.load_shot() and a cleaned shot table, in consistent units on one time base.

Every panel in the UI is tagged measured / model / learned / computed, so data is never mistaken for model.

A five-minute demo

0. First visit: the app asks. New to fusion starts the guided study (about 15 minutes; /?lesson=1.5 jumps into a chapter); I know tokamaks opens the lab. The header's Guided | Lab switch changes it any time.
1. Replay tab. Shot #30166 loads. Read the logbook note ("H-mode from 223 ms"), press ▶ Play, and watch the flux surfaces, Thomson profile, gauges and 3D vessel move together. The measured stored energy leaves the L-mode law and joins the H-mode law.
2. Pick #27257 (held out of all training): the PhysicsNeMo flux surfaces (green, dashed) lie on EFIT's (blue).
3. Vessel view: field lines traced by Warp inside the real vessel, traced q95 next to EFIT's. Drag to rotate; Inside the vessel puts the camera in the tank.
4. Pick #30192: the Troyon gauge turns red before the shot ends, and the logbook says why. What-if in the toolbar: give its beams 30 % more power and read the band on the energy strip and the limit crossing (/?shot=30192&whatif=1&P_nbi=1.3). Current and field have no slider, and the panel says why.
5. Shot database tab: this shot among 6,353, the holdout table, and the exponent caveat.
6. OpenUSD stage in the toolbar downloads the shot; open it in usdview, Omniverse USD Composer or Blender.
7. Sandbox tab: push ITER's density past the Greenwald limit and watch the operating map.

Guided and Lab modes

The guided study runs on the real screen, so finishing it leaves you in the UI you now know how to read: nine chapters from why fusion, through the knobs and the three limits on ITER's sandbox, to reading shot 30166 instrument by instrument and watching 30192 go wrong. Tasks check themselves from the app's own events (predict what happens to τ_E when you raise the heating, then do it; scrub to the first slice over the Troyon limit). The lessons are data (web/lessons.json, web/glossary.json); tests/test_guide.py proves every task can be done on the cached shots and is not already done when its step opens. Objectives and sources per chapter: docs/CURRICULUM.md.

Virtual shot

The level of confinement comes from the shot (a blind run with the learned correction and no knowledge of the measured energy, and an anchored run that reproduces it); the response to an edit comes from published scaling-law exponents, never from the network's own derivatives. Blind re-fly of 117 held-out shots puts the flat-top error at 0.175 (median |ln W_refly/W_measured|) with the correction against 0.270 for IPB98 alone. Matched pairs of real shots showed that on MAST the confinement time does not depend on density (958 pairs), so IPB98's n^0.41 is never used, and that routine operation cannot test a current or field edit (14 and 3 pairs), so those sliders do not exist. Tables and the forward-equilibrium spike that did not ship: docs/RESULTS.md; equations: docs/PHYSICS.md.

What it found on real data

Every number comes from code in this repo. Full list, tables and caveats: docs/RESULTS.md.

• The scaling laws land close on real discharges: H89 ≈ 1.15 on ohmic shot 30420, H98 ≈ 1.16 on beam-heated 30166 after its logged H-mode transition.
• Three of the five cached shots run above the Troyon β_N limit before they end. The 0D limits flag the neighbourhood, not the cause.
• The learned correction helps where it has seen data and does not extrapolate. On unseen sessions it has 28% less error than a refit power law; on an unseen campaign, 10% more. A random train/test split would have shown ~40% better.
• Magnetics → flux map: 1.65% median error on 176 shots from unseen sessions, against 2.64% for ridge regression. Feeding it EFIT's own fitted I_p cuts the error 5×, which is a leak, so we don't.
• Warp field lines reproduce EFIT's safety factor to a median 0.11–0.22% on all 209 slices, with no fitted factor.

Architecture

Layer	Tech
Data	UKAEA FAIR-MAST: level-2 Zarr on S3 (zarr, s3fs) + ndjson shot table; cached as .npz
Physics	Python + NumPy, vectorized: 0D power balance (sandbox), scaling laws on measured P_loss (replay). Sources in docs/PHYSICS.md
Learning	NVIDIA PhysicsNeMo FullyConnected on PyTorch/CUDA: τ_E correction, and the equilibrium surrogate
Field lines	NVIDIA Warp kernel: RK4 through the EFIT flux map, one GPU thread per line
Twin export	OpenUSD (usd-core): real vessel + PF coils, time-sampled plasma and field lines
API / UI	FastAPI + Uvicorn; static HTML/JS with Plotly and three.js vendored. No build step, works offline

flowchart LR
  subgraph FAIR-MAST [UKAEA FAIR-MAST · CC BY-SA 4.0]
    Z[level-2 Zarr on S3<br/>EFIT · Thomson · wall · PF coils]
    N[ndjson shot table<br/>15,969 shots + logbook]
  end
  Z --> M[mast.py<br/>SI → project units · EFIT time base · npz cache]
  N --> M
  M -->|shot bundle| R[physics.replay<br/>IPB98 · ITER89-P on measured P_loss<br/>Greenwald · Troyon · kink]
  M -->|6,353 usable shots| S[surrogate.py<br/>PhysicsNeMo residual<br/>session + campaign holdout]
  S -->|correction × IPB98| R
  Z -->|141k EFIT slices| Q[eq_surrogate.py<br/>PhysicsNeMo: 93 magnetic signals → ψ 65×65]
  Q --> A
  M --> F[fieldlines.py<br/>NVIDIA Warp GPU tracing<br/>q check vs EFIT q95]
  F --> U
  F --> A
  M --> U[usd_export.py<br/>time-sampled OpenUSD stage]
  R --> A[FastAPI<br/>/shots /replay /export /db /db/search /surrogate /simulate /map]
  S --> A
  U --> A
  E[physics.simulate<br/>0D power balance<br/>ITER · SPARC · JET · DIII-D · MAST] --> A
  A --> W[Web UI<br/>Replay · Shot database · Sandbox]
  U -.-> O[usdview · Omniverse USD Composer · Blender]

Loading

Performance

Measured with make bench and make fieldlines on an AMD Threadripper 3990X + NVIDIA RTX 3090; full tables in docs/RESULTS.md.

Task	Time	Throughput / note
Replay of a real shot (71 EFIT slices), vectorized	0.11 ms	627,247 slices/s, 56x faster than one call per slice
PhysicsNeMo correction, RTX 3090 (1,016,480 rows)	5.86 ms	173,528,297 rows/s
Warp: 4,096 field lines × 20 turns × 256 RK4 steps, RTX 3090	15.4 ms	1.36 × 10⁹ steps/s
Load a cached shot (npz, no network)	7.0 ms	first fetch from S3: ~27 s

Data & licence

MAST spherical tokamak (UKAEA, campaigns M5–M9, 2004–2013) from the open FAIR-MAST archive, https://mastapp.site, licence CC BY-SA 4.0. The files under data/ are derived from it and carry the same licence; data/README.md lists exactly what we changed. No synthetic training data is used. Please cite Jackson et al., SoftwareX 27 (2024) 101869, doi:10.1016/j.softx.2024.101869, and Jackson et al., IEEE Trans. Plasma Sci. (2025), doi:10.1109/TPS.2025.3583419.

Code: MIT.

Limitations

• Reduced model: 0D, steady-state scaling laws. Compared with MAST data, not validated against it, and not predictive.
• The learned correction is trained on one point per shot; it improves 3 of the 5 showcase shots and worsens 2, and the app prints the per-shot error both ways.
• The equilibrium surrogate learns EFIT's answer; it does not check it. Field lines are those of EFIT's axisymmetric reconstruction: no islands, no 3D fields.
• Omniverse: we ship an OpenUSD export (Omniverse-compatible). We have not opened it in an Omniverse Kit app yet.
• The virtual shot models no L–H transition, no fast-ion slowing down and nothing within ~30 ms of a beam edge (the beam is a logged box); slices outside the correction's training range are shaded, not hidden.

The full list and next steps are in docs/RESULTS.md.

Security posture

FusionLab is a single-user tool with no authentication, by design: everything it serves is open FAIR-MAST data (CC BY-SA 4.0), and it needs no API keys (.env.example). make dev starts uvicorn with no --host flag, so the server binds to 127.0.0.1 and is reachable only from your own machine. If you deliberately expose the app on a network, add an authentication gate of your own first — none exists in this repository.

Optional AI agent

There is also a chat panel — off unless you turn it on. Enabled (FUSIONLAB_AGENT_ENABLED=1 and an ANTHROPIC_API_KEY in .env), it answers from the same data the app serves: search the 15,969-shot catalog in plain language (beam-heated shots above 4 MA in campaign M8 — value sorting the UI has no way to do), ask why a what-if slider is admissible or refused, or upload a paper and get a step-by-step protocol draft to review and accept into data/protocols/. It proposes, you confirm: every state-changing action arrives as a card that shows its arithmetic (3.0 MW = 0.60 × the shot's maximum beam power) and nothing applies until you accept it, with the matched-pairs evidence gate enforced server-side exactly as in the virtual-shot UI. Without those two settings the panel is never rendered, /static/agent.js answers 404, and no key is needed — the offline demo is unchanged. Agent answers come from the app's own tools, compared with MAST data, never validated against it and never predictive; limit crossings are distances, not forecasts.

Credit

fusionsimulator.io (Daniel Burgess, Columbia Fusion Research Center) was our reference for what a clear fusion UI feels like. We copied no code, shaders, geometry or assets from it; how the two projects differ is in docs/PRIOR_ART.md.