Accelerating discovery of reduced models for RF heating using graph compilation

Optimizing the heating and current drive of a fusion pilot plant requires fast models to explore the design space. System codes rely on reduced models which sacrifice physics fidelity for computational speed. Higher fidelity modeling is necessary to validate the design before component fabrication. Often higher fidelity modeling reveals problems requiring redesign. This cycle slows progress and moves away from optimal solutions. Recently machine learning methods have shown successful capture of higher fidelity physics in fast to evaluate models. The challenge of machine learning is the large data sets necessary to train an effective model. Using existing ray tracing codes, data sets for ~10⁴ rays can be generated. A new ray-tracing code was developed where physics equations are encoded into a graph of operations. Algebraic simplifications and auto-difference transformations set up the ray update equations. The final optimized graph is Just-In-Time (JIT) compiled to a kernel which can be run on GPU or CPU resources. Using this framework, it is now possible to model ~10⁶ rays using a single Perlmutter GPU node. We will present, physics verification, benchmarking against TORAY, ray tracing in stellarators, and optimization for GPU architectures.