gyaradax: Local Gyrokinetics JAX Code

The paper introduces gyaradax, a high-performance JAX/CUDA solver for local flux-tube gyrokinetics that achieves significant speedups and native GPU acceleration over legacy Fortran codes like GKW, while demonstrating how agentic workflows guided by human expertise can rapidly translate complex physics simulations to enable advanced machine learning applications in plasma physics.

The Gist

Imagine you are trying to predict the weather inside a giant, super-hot, swirling star (a fusion reactor). The "weather" here is plasma turbulence—chaotic movements of charged particles that can ruin our ability to generate clean energy.

For decades, scientists have used massive, complex computer programs (written in an old language called Fortran) to simulate this. Think of these old programs as giant, clunky steam engines. They work, but they are heavy, hard to fix, and they run on slow, old-fashioned CPUs. They are also terrible at talking to modern Artificial Intelligence (AI), which needs to be fast and flexible to learn from data.

Enter gyaradax.

The "Tesla" of Plasma Physics

The authors of this paper built gyaradax, a brand-new simulation tool. If the old programs are steam engines, gyaradax is a sleek, electric sports car.

• It's Built for Speed: It runs on modern graphics cards (GPUs), the same chips that power video games and AI. This makes it 10 times faster than the old code.
• It's "AI-Ready": It's written in a modern language called JAX. This is like building a car with a self-driving computer built right into the engine. It allows scientists to use AI to instantly figure out how to tweak the reactor for better performance, something that was nearly impossible with the old "steam engine" code.
• It's Tiny: The old code was like a library of 30,000 books (lines of code). gyaradax does the exact same physics job in just 3,000 lines. It's a "minimalist" masterpiece.

The "Vibe Coding" Revolution

Here is the most fascinating part of the story: How did they build it so fast?

Usually, rewriting complex scientific code takes years of human labor. The authors didn't just write it; they used AI Coding Agents (smart computer programs that write code) to do the heavy lifting.

They call this "Vibe Coding." Imagine you are a conductor (the human expert) and you have a band of incredibly talented but slightly chaotic musicians (the AI agents).

1. The Conductor's Job: You give the musicians a clear sheet of music (the prompt) and a strict rule: "If you play a wrong note, the audience (the computer tests) will boo immediately."
2. The Musicians' Job: The AI agents write the code.
3. The Safety Net: Every time the AI writes a piece of code, the computer runs a "test" (like a math quiz). If the code fails the quiz, the AI has to rewrite it immediately.

Because the scientists set up these strict "quizzes" (unit tests), the AI agents could translate the old, complex Fortran code into the new JAX code in a matter of days, not years. The humans acted as the supervisors, checking the strategy, while the AI did the actual typing and debugging.

Why Does This Matter?

Think of the old code as a locked black box. You could put data in and get results out, but you couldn't easily ask "What happens if I change this one tiny thing?" or "Can we use AI to find the perfect settings?"

gyaradax opens the box.

• Inverse Problems: Instead of just simulating what happens, scientists can now work backward. They can say, "I want the plasma to behave this way," and the code can instantly calculate exactly what settings are needed to make that happen.
• Sensitivity Analysis: It can instantly tell you which knobs on the fusion reactor are the most critical to turn, helping engineers design better machines.

The Bottom Line

This paper is a proof-of-concept that AI can help us build better science tools faster than ever before.

They took a 30-year-old, difficult-to-maintain physics code, handed it to a team of AI agents supervised by humans, and in a very short time, produced a faster, smarter, and more flexible version that bridges the gap between old-school physics and the future of AI. It's not just a new code; it's a new way of doing science.

Technical Summary

1. Problem Statement

Plasma Turbulence Simulation Challenges:

• Legacy Codebases: The field of computational plasma physics relies heavily on established, highly optimized codes like GKW (Gyrokinetic Well) and GENE. However, these are typically written in Fortran, designed for CPU clusters, and rely on rigid MPI parallelization.
• Incompatibility with Modern Workflows: These legacy codes are difficult to maintain, lack native GPU acceleration, and are incompatible with Machine Learning (ML) workflows, specifically automatic differentiation (AD). This hinders research in inverse problems, sensitivity analysis, and physics-informed ML.
• Complexity: GKW consists of over 30,000 lines of Fortran code, making it a barrier to entry for researchers wanting to modify or extend the physics.

2. Methodology

The authors present gyaradax, a minimal, modern solver for local flux-tube gyrokinetics built entirely in JAX (a high-performance numerical computing library for Python).

Core Technical Approach:

• Physics Model: The solver implements the collisionless electrostatic gyrokinetic equations in the local flux-tube limit. It supports both adiabatic and kinetic electron models. It uses the standard $\delta f$ formulation, decomposing the distribution function into a stationary Maxwellian background and a fluctuating component.
• Numerical Discretization:
  • Time Integration: Explicit 4th-order Runge-Kutta (RK4).
  • Spatial Derivatives: 4th-order central and upwinded finite differences for parallel ($s$) and velocity ($v_\parallel$) coordinates.
  • Perpendicular Dimensions: Pseudospectral resolution using 2D FFTs with 3/2-rule dealiasing to handle nonlinear $E \times B$ advection.
• Performance Optimizations:
  • Mixed Precision: Utilizes Float32 for nonlinear FFTs and spatial derivatives (where physical fidelity is maintained) while keeping linear terms and field solvers in Float64.
  • Custom CUDA Kernels: Replaces standard XLA (JAX compiler) operations with custom CUDA kernels.
    • Fused Linear RHS: Combines 9-point parallel stencils, 5-point velocity stencils, and 8 physics terms into a single pass, eliminating gather/predicate overhead.
    • Fused FFTs (LTO): Uses Link-Time Optimization (LTO) callbacks within cuFFT to fuse spectral derivative multiplications, gyro-averaging, and bracket computations directly into the FFT butterfly passes, avoiding intermediate High-Bandwidth Memory (HBM) round-trips.
    • Z2Z Packing: Reduces the number of inverse FFTs by packing two real-valued signals into a single complex-to-complex transform.
• Development Workflow (Agentic/Vibecoding):
  • The project was developed using coding agents (LLMs like GPT-5.3 CODEX, Claude 4.6 Opus) guided by human expertise.
  • Strategy: A "Test-Driven Development" (TDD) loop where agents were tasked with translating Fortran logic, writing unit tests, and validating against analytical benchmarks.
  • Oversight: Human experts provided the initial context (Fortran source, manuals) and verified the "reward signal" (unit tests) to ensure correctness as complexity grew.

3. Key Contributions

1. gyaradax Solver: A fully differentiable, JAX-based implementation of local gyrokinetics (~3,000 lines of code) that achieves numerical parity with the 30,000+ line Fortran GKW code.
2. Hardware Acceleration & Differentiability: Native support for GPU acceleration and end-to-end automatic differentiation, enabling gradient-based optimization and sensitivity analysis directly within the simulation loop.
3. Agentic Software Engineering: A demonstration that coding agents, when guided by rigorous unit testing and human oversight, can successfully translate complex, legacy scientific codebases (Fortran to JAX) and implement advanced optimizations (custom CUDA kernels) at unprecedented speed.
4. Performance Benchmarks: Significant speedups over GKW and standard JAX implementations through mixed precision and custom kernel fusion.

4. Results

Validation:

• Analytical Benchmarks: Successfully passed the Rosenbluth-Hinton zonal flow test (verifying field solver and linear dynamics) and the Cyclone Base Case (linear ITG instability growth rates), matching analytical predictions and GKW reference data.
• Empirical Validation: Compared against GKW across 46 unstable adiabatic configurations and kinetic electron cases.
  • Statistical Agreement: While chaotic systems diverge point-wise due to floating-point differences, gyaradax showed statistical parity with GKW in time-averaged transport fluxes, wavenumber spectra, and growth rates.
  • Metrics: Achieved a relative Mean Absolute Error (rMAE) of 0.14 for heat flux and near-perfect correlation (Pearson $r > 0.95$) for spectral profiles.

Performance Comparison:

• Speedup:
  • JAX (Mixed Precision): ~5.4$\times$ faster than GKW.
  • CUDA (Mixed Precision): 10.5$\times$ faster than GKW.
  • Even the unoptimized JAX backend was ~2.2$\times$ faster than GKW.
• Memory: gyaradax runs on a single Nvidia Blackwell B300 GPU (275GB VRAM), whereas GKW requires a 128-core CPU cluster with 1TB RAM.

Applications (Differentiable Programming):

• Inverse Problems: Successfully recovered the temperature gradient ($R/L_T$) from a target electrostatic potential using gradient descent, converging to the ground truth ($6.908$ vs $6.9$).
• Sensitivity Analysis: Computed exact Jacobians of growth rates with respect to parameters in a single backward pass, revealing high sensitivity of high-$k_y$ modes to drive changes.

5. Significance

• Bridging Physics and ML: gyaradax removes the barrier between traditional plasma physics simulations and modern ML pipelines. It enables physics-informed machine learning, automated parameter optimization, and surrogate model training that were previously impossible with legacy codes.
• Accessibility: By reducing the codebase from 30,000 lines of Fortran to 3,000 lines of JAX, the tool lowers the entry barrier for researchers, making the code easier to extend, distribute, and maintain.
• Future of Scientific Software: The paper serves as a proof-of-concept for agentic-assisted development of high-performance scientific software. It demonstrates that with proper verification loops, AI agents can handle the translation and optimization of complex numerical frameworks, freeing human researchers to focus on fundamental physics rather than software engineering.

Limitations & Future Work:
Currently limited to electrostatic, collisionless plasmas. Future work aims to include electromagnetic perturbations, collisionality, global geometry, and grid parallelism (sharding) for larger problem sizes.