gyaradax: AI Agent-Rewritten Plasma Physics Simulator Integrating JAX Automatic Differentiation and GPU Acceleration

Controlled nuclear fusion is one of the ultimate solutions for human energy, and plasma turbulence is a core challenge in achieving controllable fusion. The gyaradax project uses AI programming agents to rewrite traditional Fortran gyrokinetic simulation code into a JAX/CUDA implementation, integrating native GPU acceleration and automatic differentiation features, and providing a new tool for interdisciplinary research between plasma physics and machine learning. This article will analyze it from dimensions such as background, methodology, verification, and applications.

Nuclear fusion releases energy through the fusion of light nuclei, requiring plasma to be heated to hundreds of millions of degrees and confined. However, heat loss caused by turbulence is a key obstacle. Gyrokinetic simulation is a core tool for studying turbulence, but traditional Fortran code has problems such as poor readability, difficulty in maintenance, and disconnection from modern ML frameworks, which seriously hinder interdisciplinary research.

gyaradax is developed based on the JAX framework, leveraging its features such as automatic differentiation, GPU acceleration, and JIT compilation; its implementation refers to the mature GKW code to ensure physical correctness. The project innovatively adopts an AI agent workflow: human expert guidance + structured prompts + unit test-driven + iterative optimization, efficiently completing the translation and optimization from Fortran to JAX.

gyaradax undergoes three-fold verification: 1. Analytical solution comparison: Simple scenarios are completely consistent with theoretical predictions; 2. Empirical benchmark comparison: Statistically equivalent to GKW; 3. Performance benchmark: Achieves significant acceleration while maintaining accuracy.

The JAX foundation of gyaradax supports ML integration in multiple scenarios: 1. Inverse problem solving: Using automatic differentiation to back-calculate plasma parameters; 2. Sensitivity analysis: Quickly calculating the gradient of outputs with respect to inputs; 3. Neural network integration: Replacing expensive subroutines or optimizing control strategies.

gyaradax brings three insights: 1. Face technical debt: Legacy code hinders innovation and needs proactive modernization; 2. AI agents are accelerators: Improve migration efficiency under guidance and verification; 3. Value of modern frameworks: JAX and others promote compatibility with the ML ecosystem.

gyaradax is not only a new simulation tool but also a paradigm for the modernization of scientific computing. It demonstrates the path of migrating traditional code to modern frameworks, the accelerating role of AI agents, and the potential of ML integration, which will help humans accelerate the mastery of the 'artificial sun' technology.