gyaradax: Local Gyrokinetics JAX Code

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Galletti, Gianluca, Volkmann, Eric, Brandstetter, Johannes
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866914454144286720
author Galletti, Gianluca
Volkmann, Eric
Brandstetter, Johannes
author_facet Galletti, Gianluca
Volkmann, Eric
Brandstetter, Johannes
contents Gyrokinetic simulations are essential for understanding and controlling turbulence in fusion plasmas, yet they are oftentimes implemented in legacy codebases, in many cases CPU-bound. These are both hard to maintain and especially incompatible with optimization and ML workflows. gyaradax is a minimal JAX/CUDA solver for local flux-tube gyrokinetics. We base our implementation on GKW (Peeters et al., 2009), but with added native GPU acceleration and automatic differentiation. We validate gyaradax against analytical cases and empirical benchmarks, achieving formal agreement and statistical parity with GKW alongside a substantial speedup. We deliberately and extensively utilized agentic workflows in this project. A key contribution is showing that coding agents, guided by human expertise, structured prompting, and measurable progress through unit testing enabled extremely fast translation of complex Fortran code, and further optimizations. Gyaradax facilitates research at the intersection of ML and plasma physics. We showcase this through practical examples in inverse problems and sensitivity analysis.
format Preprint
id arxiv_https___arxiv_org_abs_2604_06085
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle gyaradax: Local Gyrokinetics JAX Code
Galletti, Gianluca
Volkmann, Eric
Brandstetter, Johannes
Plasma Physics
Computational Physics
Gyrokinetic simulations are essential for understanding and controlling turbulence in fusion plasmas, yet they are oftentimes implemented in legacy codebases, in many cases CPU-bound. These are both hard to maintain and especially incompatible with optimization and ML workflows. gyaradax is a minimal JAX/CUDA solver for local flux-tube gyrokinetics. We base our implementation on GKW (Peeters et al., 2009), but with added native GPU acceleration and automatic differentiation. We validate gyaradax against analytical cases and empirical benchmarks, achieving formal agreement and statistical parity with GKW alongside a substantial speedup. We deliberately and extensively utilized agentic workflows in this project. A key contribution is showing that coding agents, guided by human expertise, structured prompting, and measurable progress through unit testing enabled extremely fast translation of complex Fortran code, and further optimizations. Gyaradax facilitates research at the intersection of ML and plasma physics. We showcase this through practical examples in inverse problems and sensitivity analysis.
title gyaradax: Local Gyrokinetics JAX Code
topic Plasma Physics
Computational Physics
url https://arxiv.org/abs/2604.06085