chaleong/fusion-equilibrium

Fusion Equilibrium Challenge (NeurIPS 2026): trained weights

These are the trained weights behind our DIII-D entry. The code, the training entrypoints and the docs are in charleneleong-ai/fusion-equilibrium.

What is here

results/ holds every model in the blend configs/blends/v21_16_grids3_fno2.yaml:

Component	Directories
15 base MLPs (10 x 2048 residual MLPs on a 192-coefficient PCA of psi)	v5_loss_mix_i03, v6_ensemble_region_i00-03, v8_ensemble_seeds_i00-03, v14_ramp_i00-05
Shape MLP (shape-tree host)	v21_normalize_seeds_i00
Two XGBoost forests on the PCA coefficients	xgb_top32_full.ubj, xgb_shape_v21s40_top32_full.ubj
Grid models: one U-Net and two FNO seeds	grid_v3, grid_fno, grid_fno_s1
The blend spec as built	blend_v21_16_grids3_fno2/blend.yaml

Each MLP directory has its training config.json, and each grid directory has its grid_config.json.

Use

hf download chaleong/fusion-equilibrium --local-dir weights
export FUSIONEQ_RESULTS=$PWD/weights/results NVIDIA_TF32_OVERRIDE=0
uv run python -m fusioneq.cli build --model $FUSIONEQ_RESULTS/blend_v21_16_grids3_fno2 \
  --configs diii_d_public_test --psi-dtype float32

Build in float32: bf16 autocast and TF32 both cost score on the magnetic axis. Store psi as float32 as well; float16 storage cost 0.00039 S on the Development fold.

Results (Development fold, Codabench)

Entry	DIII-D S
This blend, float16 psi (d5a34274)	0.994192
This blend, float32 psi (0be92e18)	0.994565

Data provenance

• DIII-D models: trained on the challenge's DIII-D training split only.
• MAST: none of these weights produce MAST predictions. MAST comes from a separate physics pipeline that uses the shipped MAST inputs and the machine geometry. No MAST equilibrium reconstructions were used anywhere.