JAX-in-Cell: A Differentiable Particle-in-Cell Code for Plasma Physics Applications

TL;DR

JAX-in-Cell is a modern, differentiable Particle-in-Cell framework implemented in JAX, enabling efficient plasma physics simulations with gradient-based optimization on various hardware.

Contribution

It introduces a fully differentiable, Python-based PIC framework leveraging JAX for high performance and AI integration in plasma physics simulations.

Findings

Achieves performance comparable to traditional codes on CPUs, GPUs, and TPUs.

Validates accuracy against standard plasma physics benchmarks.

Demonstrates end-to-end differentiability for physics optimization.

Abstract

JAX-in-Cell is a fully electromagnetic, multispecies, and relativistic 1D3V Particle-in-Cell (PIC) framework implemented entirely in JAX. It provides a modern, Python-based alternative to traditional PIC frameworks. It leverages Just-In-Time compilation and automatic vectorization to achieve the performance of traditional compiled codes on CPUs, GPUs, and TPUs. The resulting framework bridges the gap between educational scripts and production codes, providing a testbed for differentiable physics and AI integration that enables end-to-end gradient-based optimization. The code solves the Vlasov-Maxwell system on a staggered Yee lattice with either periodic, reflective, or absorbing boundary conditions, allowing both an explicit Boris solver and an implicit Crank-Nicolson method via Picard iteration to ensure energy conservation. Here, we detail the numerical methods employed, validate against standard benchmarks, and showcase the use of its auto-differentiation capabilities.