Gyrokinetic adiabatic-electron turbulence (256 trajectories)

Adiabatic-electron gyrokinetic turbulence simulations (GKW): the full 5D distribution function and electrostatic potentials at every timestep, in bfloat16. This is the dataset used to train the GyroSwin neural surrogates.

• GyroSwin (surrogate model, source of this dataset): https://arxiv.org/abs/2510.07314
• PINC (physics-informed neural compression): https://arxiv.org/abs/2602.04758v2
• Code: https://github.com/ml-jku/neural-gyrokinetics

Parameter scan

Cyclone Base Case (CBC) ion-temperature-gradient turbulence, scanned across the trajectories:

Grid resolution (nvpar, nmu, ns, nkx, nky) = (32, 8, 16, 85, 32), fixed across trajectories.

Storage precision (bf16)

Data is stored in bfloat16. Reconstruction loss vs float32, over ~4900 random snapshots across all trajectories:

Structure

iteration_<n>_ifft_realpotens/
  data/timestep_<t>.bf16.bin   # 5D distribution function f
  data/poten_<t>.bf16.bin      # electrostatic potentials
  metadata_light.npz           # geometry, grid, spectra metadata
  input.dat                    # GKW input deck

Usage

Download the dataset (the full set is large; grab one trajectory first to get started):

from huggingface_hub import snapshot_download
snapshot_download(
    "gerkone/cbc-gyroswin-256traj",
    repo_type="dataset",
    allow_patterns=["iteration_0_ifft_realpotens/*"],   # just ONE trajectory (~12 GB); drop this line for the full set
    local_dir="data/preprocessed_kvikio",
)

Train GyroSwin, from the code repo:

python main.py \
  workflow=gyroswin dataset=cyclone_gyroswin model=multi \
  dataset.path=data/preprocessed_kvikio \
  +dataset.prefer_dtype=bf16 \
  training.batch_size=1

Notes:

• +dataset.prefer_dtype=bf16 is required — this dataset ships bf16 shards (*.bf16.bin); without it the loader looks for fp32 shards and errors.
• The training command uses the dataset config's full trajectory set — download all trajectories first (drop the allow_patterns filter above). To just smoke-test on the single pulled trajectory, add dataset.training_trajectories=[iteration_0] dataset.validation_trajectories=[iteration_0] (trajectory names omit the _ifft_realpotens suffix).
• df normalization is recomputed from the field data on first load (a stats cache is written next to the data); the precomputed stats are intentionally not shipped.
• Both training and evaluation run directly on the bf16 shards (prefer_dtype=bf16 applies to the validation loader too).