BoSS models

The models from the Benchmark of Stellarator Simulations (the BoSS), a database of stellarator configurations paired with ideal-MHD equilibria and the guiding-center orbits of fusion-born alpha particles. Documentation · Package

folder	model	inputs	size
loss_fraction/	Predicts the loss fraction of the alpha particles born on s = 0.1 within 0.1 s. This can be used for surrogate-based optimization.	the boundary and profiles; no VMEC run	46 MB
heat_flux/	Samples where the lost alpha particles strike the boundary, and predicts the loss fraction, for particles born on s = 0.1, 0.25 or 0.5. This is the i.i.d. flow-matching model of the heat-flux benchmark.	the boundary, profiles and VMEC iota	11 MB
inverse_design/	Generates plasma boundaries with a chosen loss fraction, shape, field and profiles, by conditional flow matching.	target properties	142 MB

Each folder holds config.json (the architecture, inputs and provenance), model.safetensors (the PyTorch weights, with the normalization statistics) and the ONNX export: model.onnx, or encoder.onnx, velocity.onnx and decoder.onnx for the two flow models. All three models were trained on the BoSS train split.

Downloading

models.download in the boss-bench package downloads a model's PyTorch files, and its ONNX files with onnx=True (see below). Without PyTorch installed, use the Hugging Face command line (pip install huggingface_hub):

# one model
hf download mishapadidar/boss --include "heat_flux/*" --local-dir BoSS-models
# the ONNX files only
hf download mishapadidar/boss --include "*/config.json" --include "*.onnx" --local-dir BoSS-models
# everything
hf download mishapadidar/boss --local-dir BoSS-models

or, in Python, huggingface_hub.snapshot_download("mishapadidar/boss", allow_patterns="heat_flux/*", local_dir="BoSS-models"). Each model then sits in its own folder, e.g. BoSS-models/heat_flux. Single files can also be downloaded from the Files and versions tab.

Usage with PyTorch

pip install "boss-bench[models]"

import h5py
import boss_bench as bb
from boss_bench import models
from boss_bench.models import HeatFluxModel, InverseDesignModel, LossFractionModel

with h5py.File("confinement_data_test_split_00000.hdf5", "r") as f:
    record = bb.load_record(f, 0)                            # one BoSS configuration

surrogate = LossFractionModel.load(models.download("loss_fraction", "BoSS-models/loss_fraction"))
surrogate.predict(record)                                    # loss fraction at s0 = 0.1

heat_flux = HeatFluxModel.load(models.download("heat_flux", "BoSS-models/heat_flux"))
zeta, theta, energy, f = heat_flux.sample(record, s0=0.1)    # strike points [rad], energies [MeV]

generator = InverseDesignModel.load(models.download("inverse_design", "BoSS-models/inverse_design"))
conditions = {name: record[name] for name in models.inverse_design.CONDITIONS if name != "loss_fraction"}
boundary_xyz = generator.sample({**conditions, "loss_fraction": 0.01})

• Inputs: every model reads BoSS fields by name, from one record or from a batch of configurations (bb.load).
• Heat-flux model: a sample is a prediction for the package's heat-flux benchmark, bb.heat_flux_metrics.
• Inverse-design model: the generated boundaries are laid out as the field boundary_xyz: 41 × 41 points over one field period.
• Gradients: calling the loss-fraction model on tensors gives the prediction with gradients, for gradient-based optimization of the boundary. Gradients are available in PyTorch only.

Usage with ONNX

The ONNX files run without PyTorch, for example with ONNX Runtime. They give the same predictions as the PyTorch models, but no gradients. Their inputs are BoSS fields under their own names, in float64 (int64 for nfp and vmec_ncurr). The batch dimension is free.

• loss_fraction/model.onnx: the fields go in and the loss fraction comes out.

• The two flow models run in three steps:

  1. encoder.onnx turns the configuration (or the target properties) into the conditioning;
  2. velocity.onnx is integrated from standard normal samples, with numpy for example;
  3. decoder.onnx turns the result into strike points or into boundary_xyz.

  The integration settings are in config.json under onnx and in the sampling metadata of velocity.onnx:

  model	steps	method	time interval
  heat flux	50	Euler	[0, 1]
  inverse design	200	midpoint	[1e-4, 1 - 1e-4]

  examples/onnx_models.py in the package is a complete example.

Performance on the BoSS test split

• loss_fraction: mean squared error 0.0047 and mean absolute error 0.035.
• heat_flux: median energy distance 0.031; median heat-flux error 0.34 (16 × 16 bins) and 0.57 (32 × 32); mean absolute loss-fraction error 0.034 (with 50 Euler steps). The metrics are defined in the package documentation.
• inverse_design: see the paper. Its designs are fitted by a Fourier surface, rescaled to ARIES-CS size and evaluated with VMEC and particle tracing.

Provenance and verification

These are the paper's checkpoints. The weights were copied unchanged from the original training checkpoints, and the normalization statistics moved into the weights files. The code in boss_bench.models reproduces the original training code on published test configurations:

• heat_flux: bit for bit;
• inverse_design: to 6e-7 after 200 integration steps;
• loss_fraction: to 1.3e-6.

The ONNX exports agree with the originals to 2e-6 (loss_fraction), 3e-4 (heat_flux strike points after 50 steps; median 6e-8) and 1.3e-6 (inverse_design). The conversion, export and verification scripts are in verification/.

License

Apache License 2.0. Copyright 2026 The Simons Foundation.